# Robust Scheduling with GFlowNets

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

## Key facts

- Cited by 69 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 43 references that are in the map.
- First posted on arXiv 17 Jan 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Machine Learning for Combinatorial Optimization: a Methodological Tour d'Horizon](https://mlsanity.com/paper/1811.06128/) (2018)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022)
- [How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks](https://mlsanity.com/paper/2009.11848/) (2020)
- [pymoo: Multi-objective Optimization in Python](https://mlsanity.com/paper/2002.04504/) (2020)
- [Understanding the impact of entropy on policy optimization](https://mlsanity.com/paper/1811.11214/) (2018)
- [Biological Sequence Design with GFlowNets](https://mlsanity.com/paper/2203.04115/) (2022)
- [Bayesian Structure Learning with Generative Flow Networks](https://mlsanity.com/paper/2202.13903/) (2022)
- [Learning to Dispatch for Job Shop Scheduling via Deep Reinforcement Learning](https://mlsanity.com/paper/2010.12367/) (2020)
- [Learning Scheduling Algorithms for Data Processing Clusters](https://mlsanity.com/paper/1810.01963/) (2018)
- [Generative Flow Networks for Discrete Probabilistic Modeling](https://mlsanity.com/paper/2202.01361/) (2022)
- [Observe and Look Further: Achieving Consistent Performance on Atari](https://mlsanity.com/paper/1805.11593/) (2018)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learning](https://mlsanity.com/paper/1906.08879/) (2019)
- [Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs](https://mlsanity.com/paper/1905.02494/) (2019)
- [Neural Topological Ordering for Computation Graphs](https://mlsanity.com/paper/2207.05899/) (2022)
- [Optimizing Memory Placement using Evolutionary Graph Reinforcement Learning](https://mlsanity.com/paper/2007.07298/) (2020)

## Built on by (most cited first, 30 of 69)

- [A theory of continuous generative flow networks](https://mlsanity.com/paper/2301.12594/) (2023)
- [Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization](https://mlsanity.com/paper/2310.02679/) (2023)
- [Learning diverse attacks on large language models for robust red-teaming and safety tuning](https://mlsanity.com/paper/2405.18540/) (2024)
- [Ant Colony Sampling with GFlowNets for Combinatorial Optimization](https://mlsanity.com/paper/2403.07041/) (2024)
- [Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network](https://mlsanity.com/paper/2305.19366/) (2023)
- [Local Search GFlowNets](https://mlsanity.com/paper/2310.02710/) (2023)
- [Improved off-policy training of diffusion samplers](https://mlsanity.com/paper/2402.05098/) (2024)
- [Discrete Probabilistic Inference as Control in Multi-path Environments](https://mlsanity.com/paper/2402.10309/) (2024)
- [GFlowNet-EM for learning compositional latent variable models](https://mlsanity.com/paper/2302.06576/) (2023)
- [Let the Flows Tell: Solving Graph Combinatorial Optimization Problems with GFlowNets](https://mlsanity.com/paper/2305.17010/) (2023)
- [FlowRL: Matching Reward Distributions for LLM Reasoning](https://mlsanity.com/paper/2509.15207/) (2025)
- [Stochastic Generative Flow Networks](https://mlsanity.com/paper/2302.09465/) (2023)
- [Thompson sampling for improved exploration in GFlowNets](https://mlsanity.com/paper/2306.17693/) (2023)
- [Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets](https://mlsanity.com/paper/2412.07775/) (2024)
- [Flow of Reasoning: Training LLMs for Divergent Reasoning with Minimal Examples](https://mlsanity.com/paper/2406.05673/) (2024)
- [Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training](https://mlsanity.com/paper/2503.18929/) (2025)
- [Distributional GFlowNets with Quantile Flows](https://mlsanity.com/paper/2302.05793/) (2023)
- [Learning to Scale Logits for Temperature-Conditional GFlowNets](https://mlsanity.com/paper/2310.02823/) (2023)
- [Improving GFlowNets for Text-to-Image Diffusion Alignment](https://mlsanity.com/paper/2406.00633/) (2024)
- [PhyloGFN: Phylogenetic inference with generative flow networks](https://mlsanity.com/paper/2310.08774/) (2023)
- [Goal-conditioned GFlowNets for Controllable Multi-Objective Molecular Design](https://mlsanity.com/paper/2306.04620/) (2023)
- [Pre-Training and Fine-Tuning Generative Flow Networks](https://mlsanity.com/paper/2310.03419/) (2023)
- [Learning to Sample Effective and Diverse Prompts for Text-to-Image Generation](https://mlsanity.com/paper/2502.11477/) (2025)
- [GFlowVLM: Enhancing Multi-step Reasoning in Vision-Language Models with Generative Flow Networks](https://mlsanity.com/paper/2503.06514/) (2025)
- [Pessimistic Backward Policy for GFlowNets](https://mlsanity.com/paper/2405.16012/) (2024)
- [Order-Preserving GFlowNets](https://mlsanity.com/paper/2310.00386/) (2023)
- [Can a Bayesian Oracle Prevent Harm from an Agent?](https://mlsanity.com/paper/2408.05284/) (2024)
- [torchgfn: A PyTorch GFlowNet library](https://mlsanity.com/paper/2305.14594/) (2023)
- [Revisiting Non-Acyclic GFlowNets in Discrete Environments](https://mlsanity.com/paper/2502.07735/) (2025)
- [Adversarial Generative Flow Network for Solving Vehicle Routing Problems](https://mlsanity.com/paper/2503.01931/) (2025)

## Abstract (arXiv)

Finding the best way to schedule operations in a computation graph is a classical NP-hard problem which is central to compiler optimization. However, evaluating the goodness of a schedule on the target hardware can be very time-consuming. Traditional approaches as well as previous machine learning ones typically optimize proxy metrics, which are fast to evaluate but can lead to bad schedules when tested on the target hardware. In this work, we propose a new approach to scheduling by sampling proportionally to the proxy metric using a novel GFlowNet method. We introduce a technique to control the trade-off between diversity and goodness of the proposed schedules at inference time and demonstrate empirically that the pure optimization baselines can lead to subpar performance with respect to our approach when tested on a target model. Furthermore, we show that conditioning the GFlowNet on the computation graph enables generalization to unseen scheduling problems for both synthetic and real-world compiler datasets.