# Neural Topological Ordering for Computation Graphs

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 44 references that are in the map.
- First posted on arXiv 13 Jul 2022.

## Builds on (most cited first, 24 of 24)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Relational inductive biases, deep learning, and graph networks](https://mlsanity.com/paper/1806.01261/) (2018)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [A Generalization of Transformer Networks to Graphs](https://mlsanity.com/paper/2012.09699/) (2020)
- [Attention, Learn to Solve Routing Problems!](https://mlsanity.com/paper/1803.08475/) (2018)
- [Machine Learning for Combinatorial Optimization: a Methodological Tour d'Horizon](https://mlsanity.com/paper/1811.06128/) (2018)
- [Neural Combinatorial Optimization with Reinforcement Learning](https://mlsanity.com/paper/1611.09940/) (2016)
- [Learning Combinatorial Optimization Algorithms over Graphs](https://mlsanity.com/paper/1704.01665/) (2017)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [A Survey on The Expressive Power of Graph Neural Networks](https://mlsanity.com/paper/2003.04078/) (2020)
- [NeuroLKH: Combining Deep Learning Model with Lin-Kernighan-Helsgaun Heuristic for Solving the Traveling Salesman Problem](https://mlsanity.com/paper/2110.07983/) (2021)
- [D-VAE: A Variational Autoencoder for Directed Acyclic Graphs](https://mlsanity.com/paper/1904.11088/) (2019)
- [Directed Acyclic Graph Neural Networks](https://mlsanity.com/paper/2101.07965/) (2021)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices](https://mlsanity.com/paper/2003.02369/) (2020)
- [A Bi-Level Framework for Learning to Solve Combinatorial Optimization on Graphs](https://mlsanity.com/paper/2106.04927/) (2021)
- [Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs](https://mlsanity.com/paper/1905.02494/) (2019)
- [Low-variance Black-box Gradient Estimates for the Plackett-Luce Distribution](https://mlsanity.com/paper/1911.10036/) (2019)
- [Reparameterizing the Birkhoff Polytope for Variational Permutation Inference](https://mlsanity.com/paper/1710.09508/) (2017)
- [Neural Simulated Annealing](https://mlsanity.com/paper/2203.02201/) (2022)

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

- [RL4CO: an Extensive Reinforcement Learning for Combinatorial Optimization Benchmark](https://mlsanity.com/paper/2306.17100/) (2023)
- [Robust Scheduling with GFlowNets](https://mlsanity.com/paper/2302.05446/) (2023)
- [Learning to Scale Logits for Temperature-Conditional GFlowNets](https://mlsanity.com/paper/2310.02823/) (2023)
- [Transformers over Directed Acyclic Graphs](https://mlsanity.com/paper/2210.13148/) (2022)
- [A Theory of Non-Acyclic Generative Flow Networks](https://mlsanity.com/paper/2312.15246/) (2023)
- [Semi-Supervised Learning for Deep Causal Generative Models](https://mlsanity.com/paper/2403.18717/) (2024)
- [MARCO: A Memory-Augmented Reinforcement Framework for Combinatorial Optimization](https://mlsanity.com/paper/2408.02207/) (2024)
- [NN-Former: Rethinking Graph Structure in Neural Architecture Representation](https://mlsanity.com/paper/2507.00880/) (2025)
- [A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling](https://mlsanity.com/paper/2603.23249/) (2026)
- [Moccasin: Efficient Tensor Rematerialization for Neural Networks](https://mlsanity.com/paper/2304.14463/) (2023)
- [Network Representation Learning for Biophysical Neural Network Analysis](https://mlsanity.com/paper/2410.11503/) (2024)
- [CaTs and DAGs: Integrating Directed Acyclic Graphs with Transformers for Causally Constrained Predictions](https://mlsanity.com/paper/2410.14485/) (2024)

## Abstract (arXiv)

Recent works on machine learning for combinatorial optimization have shown that learning based approaches can outperform heuristic methods in terms of speed and performance. In this paper, we consider the problem of finding an optimal topological order on a directed acyclic graph with focus on the memory minimization problem which arises in compilers. We propose an end-to-end machine learning based approach for topological ordering using an encoder-decoder framework. Our encoder is a novel attention based graph neural network architecture called Topoformer which uses different topological transforms of a DAG for message passing. The node embeddings produced by the encoder are converted into node priorities which are used by the decoder to generate a probability distribution over topological orders. We train our model on a dataset of synthetically generated graphs called layered graphs. We show that our model outperforms, or is on-par, with several topological ordering baselines while being significantly faster on synthetic graphs with up to 2k nodes. We also train and test our model on a set of real-world computation graphs, showing performance improvements.