Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

Authors: Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, Oriol Vinyals

ICLR 2020 | Conference PDF | Archive PDF | Plain Text | LLM Run Details

Reproducibility Variable Result LLM Response
Research Type Experimental 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 Tensor Flow 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.
Researcher Affiliation Industry Aditya Paliwal Google Research adipal@google.com Felix Gimeno Deep Mind fgimeno@google.com Vinod Nair Deep Mind vinair@google.com Yujia Li Deep Mind yujiali@google.com Miles Lubin Google Research mlubin@google.com Pushmeet Kohli Deep Mind pushmeet@google.com Oriol Vinyals Deep Mind vinyals@google.com
Pseudocode No The paper describes algorithms but does not contain structured pseudocode or algorithm blocks (clearly labeled algorithm sections or code-like formatted procedures).
Open Source Code No The paper states: 'The synthetic data in Cost Graph Def format is available at https://github.com/deepmind/deepmind-research/tree/master/regal.' This link is specifically for data, not the open-source code for the methodology described in the paper.
Open Datasets Yes For reproducibility, we have released a synthetic dataset of computation graphs with 10000 training, 1000 validation, and 1000 test cases. The synthetic data in Cost Graph Def format is available at https://github.com/deepmind/deepmind-research/tree/master/regal.
Dataset Splits Yes The dataset is split into {train, valid, test} sets containing {60%, 20%, 20%} graphs, respectively.
Hardware Specification Yes Table 3 shows the average running times of the various algorithms on the Tensor Flow test set and the XLA dataset, as measured on an Intel Xeon E5-1650 3.60GHz machine.
Software Dependencies No The paper mentions frameworks like Tensor Flow, MXNet, Py Torch and a solver (CP-SAT solver of Google OR-tools) but does not provide specific version numbers for these software components. For example, it mentions 'CP-SAT solver of Google OR-tools (Google, 2019)' but not a version of OR-tools.
Experiment Setup Yes The graph neural network had a state size of 32 for each node and edge, 16 propagations, all networks MLPn MLPe MLPnode MLPmsg MLP msg being two layers of size 32, the aggregation used was mean pooling. For faster training, the reward of the training set was made with 1000 fitness evaluations for REGAL and BRKGA (4600 for REGAL and 5000 for BRKGA for the validation and test sets). Training lasted 100000 gradient steps with each step having a mini-batch of size 4 and with gradient clipping by L2 norm with value 10 . The baseline mean squared error term s contribution to the overall loss was weighted by 0.0001 . The optimizer was Adam with beta1 0.9 beta2 0.999 epsilon 1e 8 and learning rate 0.0001 . The number of devices (for the memory model) was 2.