[Paper Review] REGAL: Transfer Learning For Fast Optimization of Computation Graphs.

Aditya Paliwal, Felix Gimeno+5 more|arXiv (Cornell University)|May 7, 2019

Parallel Computing and Optimization Techniques20 references12 citations

TL;DR

REGAL proposes a transfer learning approach using deep reinforcement learning to optimize neural network computation graphs for reduced execution time and peak memory usage. It trains an optimizer offline to generalize across unseen graphs, achieving high-quality optimizations in seconds—dramatically faster than hours required by prior methods—while outperforming classical and learning-based baselines.

ABSTRACT

We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.

Motivation & Objective

• To address the inefficiency of existing optimization methods that require retraining for each new graph.
• To enable fast, high-quality optimization of computation graphs without fine-tuning on unseen models.
• To minimize both execution time and peak memory usage in neural network computation graphs.
• To generalize a single trained optimizer across diverse real-world TensorFlow graphs.

Proposed method

• A deep reinforcement learning agent is trained offline on a diverse set of computation graphs to learn optimization policies.
• The agent learns to make decisions such as operator fusion, layout transformation, and memory allocation.
• Transfer learning is applied to enable the same policy network to generalize to previously unseen graphs without further training.
• The method uses a reward function that combines execution time and memory usage to guide optimization decisions.
• The policy network is trained using off-policy reinforcement learning with experience replay to improve sample efficiency.
• The framework is integrated into an optimizing compiler pipeline to apply decisions at inference time.

Experimental results

Research questions

• RQ1 Can a single reinforcement learning policy be trained once and generalized to optimize previously unseen computation graphs?
• RQ2 How does the proposed method compare to classical optimization techniques in terms of execution time and memory usage?
• RQ3 To what extent can transfer learning reduce the need for retraining when optimizing new models?
• RQ4 Can the method achieve significant speedups in optimization time while maintaining or improving optimization quality?

Key findings

• REGAL achieves optimization quality comparable to or better than classical and learning-based baselines on both execution time and peak memory usage.
• The method reduces optimization time from hours to seconds on real-world TensorFlow graphs.
• The trained policy generalizes effectively to unseen graphs without fine-tuning, demonstrating strong transferability.
• The approach outperforms strong baselines in minimizing execution time and memory footprint across diverse models.
• The reward-based RL framework successfully balances trade-offs between speed and memory efficiency.