# Transferable Graph Optimizers for ML Compilers

arXiv: https://arxiv.org/abs/2010.12438 · HTML: https://mlsanity.com/paper/2010.12438/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 21 Oct 2020.

## Built on by (most cited first, 29 of 29)

- [Full Stack Optimization of Transformer Inference: a Survey](https://mlsanity.com/paper/2302.14017/) (2023)
- [Robust Scheduling with GFlowNets](https://mlsanity.com/paper/2302.05446/) (2023)
- [SurCo: Learning Linear Surrogates For Combinatorial Nonlinear Optimization Problems](https://mlsanity.com/paper/2210.12547/) (2022)
- [A Full-Stack Search Technique for Domain Optimized Deep Learning Accelerators](https://mlsanity.com/paper/2105.12842/) (2021)
- [Automap: Towards Ergonomic Automated Parallelism for ML Models](https://mlsanity.com/paper/2112.02958/) (2021)
- [Neural Topological Ordering for Computation Graphs](https://mlsanity.com/paper/2207.05899/) (2022)
- [Learning Compiler Pass Orders using Coreset and Normalized Value Prediction](https://mlsanity.com/paper/2301.05104/) (2023)
- [HASCO: Towards Agile HArdware and Software CO-design for Tensor Computation](https://mlsanity.com/paper/2105.01585/) (2021)
- [Enabling Resource-efficient AIoT System with Cross-level Optimization: A survey](https://mlsanity.com/paper/2309.15467/) (2023)
- [A Graph Policy Network Approach for Volt-Var Control in Power Distribution Systems](https://mlsanity.com/paper/2109.12073/) (2021)
- [Neural Architecture Search using Property Guided Synthesis](https://mlsanity.com/paper/2205.03960/) (2022)
- [Optimizing Memory Mapping Using Deep Reinforcement Learning](https://mlsanity.com/paper/2305.07440/) (2023)
- [REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving](https://mlsanity.com/paper/2506.01374/) (2025)
- [A Transferable Approach for Partitioning Machine Learning Models on Multi-Chip-Modules](https://mlsanity.com/paper/2112.04041/) (2021)
- [Understanding Curriculum Learning in Policy Optimization for Online Combinatorial Optimization](https://mlsanity.com/paper/2202.05423/) (2022)
- [X-RLflow: Graph Reinforcement Learning for Neural Network Subgraphs Transformation](https://mlsanity.com/paper/2304.14698/) (2023)
- [Saturn: An Optimized Data System for Large Model Deep Learning Workloads](https://mlsanity.com/paper/2309.01226/) (2023)
- [Optimizing Automatic Differentiation with Deep Reinforcement Learning](https://mlsanity.com/paper/2406.05027/) (2024)
- [A Graph Deep Learning Framework for High-Level Synthesis Design Space Exploration](https://mlsanity.com/paper/2111.14767/) (2021)
- [Terra: Imperative-Symbolic Co-Execution of Imperative Deep Learning Programs](https://mlsanity.com/paper/2201.09210/) (2022)
- [End-to-end Mapping in Heterogeneous Systems Using Graph Representation Learning](https://mlsanity.com/paper/2204.11981/) (2022)
- [Attention mechanisms for physiological signal deep learning: which attention should we take?](https://mlsanity.com/paper/2207.06904/) (2022)
- [Celeritas: Fast Optimizer for Large Dataflow Graphs](https://mlsanity.com/paper/2208.00184/) (2022)
- [Target-independent XLA optimization using Reinforcement Learning](https://mlsanity.com/paper/2308.14364/) (2023)
- [A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling](https://mlsanity.com/paper/2603.23249/) (2026)
- [Reinforcement Learning Approach for Mapping Applications to Dataflow-Based Coarse-Grained Reconfigurable Array](https://mlsanity.com/paper/2205.13675/) (2022)
- [Expediting Distributed DNN Training with Device Topology-Aware Graph Deployment](https://mlsanity.com/paper/2302.06126/) (2023)
- [Network Representation Learning for Biophysical Neural Network Analysis](https://mlsanity.com/paper/2410.11503/) (2024)
- [Learning to Shard: RL for Co-optimizing the Parallelism Degrees and Per-operator Sharding Dimensions in Distributed LLM Inference](https://mlsanity.com/paper/2509.00217/) (2025)

## Abstract (arXiv)

Most compilers for machine learning (ML) frameworks need to solve many correlated optimization problems to generate efficient machine code. Current ML compilers rely on heuristics based algorithms to solve these optimization problems one at a time. However, this approach is not only hard to maintain but often leads to sub-optimal solutions especially for newer model architectures. Existing learning based approaches in the literature are sample inefficient, tackle a single optimization problem, and do not generalize to unseen graphs making them infeasible to be deployed in practice. To address these limitations, we propose an end-to-end, transferable deep reinforcement learning method for computational graph optimization (GO), based on a scalable sequential attention mechanism over an inductive graph neural network. GO generates decisions on the entire graph rather than on each individual node autoregressively, drastically speeding up the search compared to prior methods. Moreover, we propose recurrent attention layers to jointly optimize dependent graph optimization tasks and demonstrate 33%-60% speedup on three graph optimization tasks compared to TensorFlow default optimization. On a diverse set of representative graphs consisting of up to 80,000 nodes, including Inception-v3, Transformer-XL, and WaveNet, GO achieves on average 21% improvement over human experts and 18% improvement over the prior state of the art with 15x faster convergence, on a device placement task evaluated in real systems.