Optimizing Neural Networks: How to Straighten Tangled Computation Graphs ✛
Context
Deep neural networks for recommender systems are computational behemoths. The more accurate the model, the more tangled its inner landscape: thousands of operations interweave into graphs with whimsical topologies. Each unnecessary twist, each repeated splitting and merging of tensors, is like an uncleared trail slowing down a traveler. At the hardware level, this translates into a cascade of tiny kernel launches on the GPU, eating up time and memory. Engineers try to tidy up manually, but models change weekly and architectures grow ever fancier. An automatic “graph trail straightener” became essential—one that can find and weed out inefficient code fragments without human intervention.
Methods
At the heart of Optimus is a collection of over thirty optimization patterns, each describing a typical “junk” motif in the graph: for example, the pointless chain “slice a tensor and immediately glue it back together.” The search for such motifs is driven by a greedy algorithm that starts from “anchors”—key operations like split, or those common to recommendation architectures. From an anchor, the algorithm visits neighboring nodes, comparing them to a reference pattern, and on a match triggers a replacement: instead of the clunky construction, a compact equivalent is inserted. There are two kinds of fixes: horizontal—where several independent identical operations (say, a dozen linear layers) are fused into one; and vertical—where redundant transitions like the infamous split-then-cat are excised. The whole procedure is baked into the PyTorch 2.x compiler and works on any of its intermediate representations (Torch-IR, Aten-IR), making it universal across a diverse zoo of models.
Results
Across five production recommendation models at Meta, spanning various architectural styles (DHEN-like, Wukong-like, Interformer-like, and others), Optimus consistently boosted throughput—up to 63% more training examples per second. Vertical optimizations contributed the most: eliminating parasitic split-cat cycles axed hundreds of redundant Triton kernels. The gains from horizontal fusions depended on the number of identical operations—the more there are, the bigger the win. Peak memory usage on large models with activation quantization dropped by 2–6%, and cold compilation time simply nosedived: in one test, from 4000+ seconds to around 3600, a saving of over 400 seconds. On an open illustrative example from TorchBench, the number of training kernels plummeted from 12 772 to 24, and inference kernels from 6 320 to 6, yielding an 11x training speedup and a 364x inference speedup. Comparisons with alternatives (TASO, TENSAT) showed they couldn’t optimize this graph at all, whereas Optimus succeeded thanks to its flexible rule system.
Implications
Optimus isn’t a one-off hack; it’s a modular framework embedded in the open-source PyTorch. Any researcher or engineer can add a new optimization pattern by registering it via a standard interface and immediately apply it to their neural network graph. This approach turns the once-tedious hand-tuning of code into systematic housekeeping at the compiler level. For industry, this means pricey GPUs get used more efficiently without rewriting models—and the carbon footprint of training gigantic recommender systems shrinks as a bonus.
Future development
Currently the core rules operate on Torch-IR and Aten-IR, but the authors plan to extend the method to Inductor-IR—a lower-level representation where even more performance can be squeezed out. Another avenue is automatic selection of the order and set of patterns for a specific model: today, the optimization sequence is chosen manually; the goal is to hand this off to an algorithm that analyzes the graph structure.
Impact
The technology will directly speed up training and deployment of recommendation systems—the backbone of modern internet platforms. It can also be adapted for any other deep neural networks, especially those with branchy computation graphs: natural language processing, computer vision, graph neural networks.
Next steps
Porting advances to the Inductor-IR level and implementing an automatic scheduler for pattern sequences. The pattern library will also be expanded to cover new architectural motifs.
Key open problems
The problem of efficiently running neural networks on modern hardware boils down to the gap between architectural complexity and limited GPU memory bandwidth. Optimus is a step toward automatically closing that gap at the compiler level, without requiring developers to rework their code. This fits into the broader quest for scalable numerical optimization methods for computation.
The knowledge machine's view
This section is addressed to the author of the paper and their colleagues. Neighbouring works were found by meaning-based search across our archive, not by word overlap. We do not review: we show the field around the work.
What the work does
This paper introduces Optimus, a modular framework for automatic optimization of computation graphs in PyTorch 2.x, targeting recommendation models. It uses a collection of over thirty templates to find and replace inefficient patterns (e.g., split-then-cat) via a greedy anchor-based algorithm. On five industrial Meta models, it achieved up to 63% training speedup, a 2–6% reduction in peak memory usage, and a significant decrease in compilation time (from 4000+ to ~3600 seconds in one test).
Where its strength is
The strength of this work lies in its practical integration with the PyTorch compiler and its flexible rule-based system, which automatically identifies and eliminates inefficient patterns in real industrial graphs, where counterparts like TASO and TENSAT fall short.
Where it could go next
- It might be interesting to look at SpecFANN's experience in accelerating computations through specialized architectures: it uses a fast emulator to generate spectra, which resonates with the idea of replacing heavy fragments with lightweight equivalents. Perhaps the principles of replacement based on trained surrogates could be adapted to expand Optimus's template set.based on: SpecFANN: Lightning-Fast Sketch of a Star's Spectrum arXiv:2607.17348v1
- Related work on a neural network that senses stress in parts of any shape applies graph analysis methods to identify critical nodes. This could suggest ways to improve the anchor selection heuristic in Optimus to find more global optimizations.based on: A neural network that senses stress in parts of any shape arXiv:2606.08287
Worth trying
- Try using a trained surrogate to predict the benefit of applying a template, in order to prioritize replacements.
- Test the analysis of graph 'stress' (analogous to the work on parts) to select anchors that are more likely to lead to cascading optimizations.
- Investigate whether new templates can be automatically generated by generalizing frequent patterns in real graphs.
Nearby works in our archive
- A neural network that senses stress in parts of any shapesimilarity 0.62
- SpecFANN: Lightning-Fast Sketch of a Star's Spectrumsimilarity 0.62
- How a Digital Twin Saves a Brain in Minutessimilarity 0.62
- Neural Network Learns to Flow Through a Sievesimilarity 0.60
- Graph Labels: When Quantum Methods Yield to Simplicitysimilarity 0.60
- Neurons Reveal Their Secrets in Minutes, Not Hourssimilarity 0.59
The closest work in the archive sits at 0.62 by meaning — the usual level for our corpus.
Original abstract
as written by the authors arXiv:2607.02945
In large-scale industrial applications, deep learning models that power recommendation and ranking have complex and diverse model architectures. These models are continuously developed and refined by large teams of machine learning engineers, rendering manual optimization infeasible. Consequently, graph-based optimization techniques have become an industry standard for boosting performance, with PyTorch FX transformations leading the charge. These transformations typically rely on a set of human-engineered module-level rewrite rules which are not scalable to diverse model architectures. To address this limitation, we introduce Optimus, a general-purpose model transformation framework built in the PyTorch 2.x (PT2) machine learning compiler. With a concise set of predefined patterns, Optimus applies an efficient greedy search algorithm for pattern matching and replacement, while preserving model semantic. It is designed and implemented as a highly customizable and extensible framework integrated into the PT2 stack. Our evaluation shows that the framework can achieve up to 63% speedup, 6% peak memory reduction, and over 400 second compile time decrease for our industry-scale recommendation models compared to baselines. Optimus is open-sourced together with PyTorch 2.x as a customizable model transformation layer.
Shaving 400 seconds off compilation is like skipping the ritual of making coffee, checking email, and waiting for your IDE to open: you just dive straight in. And the inference kernel count for the illustrative model dropped from 6320 to 6—a compaction of the graph comparable to replacing a swarm of couriers with a single truck.
Key numbers
- throughput boost (QPS): up to 63%
- peak memory reduction: up to 6%
- compilation time reduction: over 400 seconds
- kernels in TorchBench training (before/after): 12 772 -> 24
- kernels in TorchBench inference (before/after): 6 320 -> 6