mlsanityData to 13 Oct 2026

RLFlow: Optimising Neural Network Subgraph Transformation with World Models

arXiv 2205.01435 ↗First posted 3 May 2022cs.LGcited by 0

On the map: World models → Latent dynamics models and model-based RL

Abstract (arXiv)

Training deep learning models takes an extremely long execution time and consumes large amounts of computing resources. At the same time, recent research proposed systems and compilers that are expected to decrease deep learning models runtime. An effective optimisation methodology in data processing is desirable, and the reduction of compute requirements of deep learning models is the focus of extensive research. In this paper, we address the neural network sub-graph transformation by exploring reinforcement learning (RL) agents to achieve performance improvement. Our proposed approach RLFlow can learn to perform neural network subgraph transformations, without the need for expertly designed heuristics to achieve a high level of performance. Recent work has aimed at applying RL to computer systems with some success, especially using model-free RL techniques. Model-based reinforcement learning methods have seen an increased focus in research as they can be used to learn the transition dynamics of the environment; this can be leveraged to train an agent using a hallucinogenic environment such as World Model (WM), thereby increasing sample efficiency compared to model-free approaches. WM uses variational auto-encoders and it builds a model of the system and allows exploring the model in an inexpensive way. In RLFlow, we propose a design for a model-based agent with WM which learns to optimise the architecture of neural networks by performing a sequence of sub-graph transformations to reduce model runtime. We show that our approach can match the state-of-the-art performance on common convolutional networks and outperforms by up to 5% those based on transformer-style architectures

Key facts. No citing papers in the map yet (data through 13 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 31 of its 52 references that are in the map. First posted on arXiv 3 May 2022.

Builds on

31 of 52 references.

Closest ancestors · 4

Same direction, newest first

  1. 2020Mastering Atari with Discrete World Models5 Oct
  2. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
  3. 2018Recurrent World Models Facilitate Policy Evolution4 Sep
  4. 2017Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning8 Aug

Other references · 27

  1. 2017Attention Is All You Needcited by 60,106
  2. 2015Deep Residual Learning for Image Recognitioncited by 51,950
  3. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,304
  4. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,043
  5. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,223
  6. 2017Proximal Policy Optimization Algorithmscited by 14,592
  7. 2015Distilling the Knowledge in a Neural Networkcited by 10,180
  8. 2015Rethinking the Inception Architecture for Computer Visioncited by 5,457
  9. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,024
  10. 2013Playing Atari with Deep Reinforcement Learningcited by 3,221
  11. 2014Sequence to Sequence Learning with Neural Networkscited by 2,981
  12. 2016Asynchronous Methods for Deep Reinforcement Learningcited by 2,223
15 more
  1. 2016OpenAI Gymcited by 1,694
  2. 2018Relational inductive biases, deep learning, and graph networkscited by 1,353
  3. 2016SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model sizecited by 1,058
  4. 2017Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithmcited by 728
  5. 2019Solving Rubik's Cube with a Robot Handcited by 702
  6. 2013Generating Sequences With Recurrent Neural Networkscited by 553
  7. 2019Model-Based Reinforcement Learning for Ataricited by 499
  8. 2016The CMA Evolution Strategy: A Tutorialcited by 368
  9. 2015MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systemscited by 195
  10. 2019OpenSpiel: A Framework for Reinforcement Learning in Gamescited by 183
  11. 2018Model-Based Value Estimation for Efficient Model-Free Reinforcement Learningcited by 152
  12. 2017Imagination-Augmented Agents for Deep Reinforcement Learningcited by 138
  13. 2017Device Placement Optimization with Reinforcement Learningcited by 85
  14. 2019Placeto: Learning Generalizable Device Placement Algorithms for Distributed Machine Learningcited by 26
  15. 2019Reinforced Genetic Algorithm Learning for Optimizing Computation Graphscited by 17

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025RLinf: Flexible and Efficient Large-scale Reinforcement Learning via Macro-to-Micro Flow Transformation
  2. 2022MoDem: Accelerating Visual Model-Based Reinforcement Learning with Demonstrations
  3. 2025Accelerating Model-Based Reinforcement Learning with State-Space World Models
  4. 2023An Invitation to Deep Reinforcement Learning
  5. 2022Value Memory Graph: A Graph-Structured World Model for Offline Reinforcement Learning
  6. 2024Scaling Offline Model-Based RL via Jointly-Optimized World-Action Model Pretraining
  7. 2025Scaling Offline RL via Efficient and Expressive Shortcut Models
  8. 2026STO-RL: Offline RL under Sparse Rewards via LLM-Guided Subgoal Temporal Order