Timezone: »
Most compilers for machine learning (ML) frameworks need to solve many correlated optimization problems to generate efficient machine code. Current ML compilers rely on heuristics based algorithms to solve these optimization problems one at a time. However, this approach is not only hard to maintain but often leads to sub-optimal solutions especially for newer model architectures. Existing learning based approaches in the literature are sample inefficient, tackle a single optimization problem, and do not generalize to unseen graphs making them infeasible to be deployed in practice. To address these limitations, we propose an end-to-end, transferable deep reinforcement learning method for computational graph optimization (GO), based on a scalable sequential attention mechanism over an inductive graph neural network. GO generates decisions on the entire graph rather than on each individual node autoregressively, drastically speeding up the search compared to prior methods. Moreover, we propose recurrent attention layers to jointly optimize dependent graph optimization tasks and demonstrate 33%-60% speedup on three graph optimization tasks compared to TensorFlow default optimization. On a diverse set of representative graphs consisting of up to 80,000 nodes, including Inception-v3, Transformer-XL, and WaveNet, GO achieves on average 21% improvement over human experts and 18% improvement over the prior state of the art with 15x faster convergence, on a device placement task evaluated in real systems.
Author Information
Yanqi Zhou (Google Brain)
Sudip Roy (Google)
Amirali Abdolrashidi (UC Riverside)
Daniel Wong (Carnegie Mellon University)
Peter Ma (Google)
Qiumin Xu (Google)
Hanxiao Liu (Google Brain)
Phitchaya Phothilimtha (Google Brain)
Shen Wang (Google Inc)
Anna Goldie (Google Brain / Stanford)
Azalia Mirhoseini (Google Brain)
James Laudon (Google)
Related Events (a corresponding poster, oral, or spotlight)
-
2020 Oral: Transferable Graph Optimizers for ML Compilers »
Fri. Dec 11th 02:30 -- 02:45 AM Room Orals & Spotlights: Health/AutoML/(Soft|Hard)ware
More from the Same Authors
-
2023 Poster: DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining »
Sang Michael Xie · Hieu Pham · Xuanyi Dong · Nan Du · Hanxiao Liu · Yifeng Lu · Percy Liang · Quoc V Le · Tengyu Ma · Adams Wei Yu -
2023 Poster: Conditional Adapters: Parameter-efficient Transfer Learning with Fast Inference »
Tao Lei · Junwen Bai · Siddhartha Brahma · Joshua Ainslie · Kenton Lee · Yanqi Zhou · Nan Du · Vincent Zhao · Yuexin Wu · Bo Li · Yu Zhang · Ming-Wei Chang -
2023 Poster: Learning Large Graph Property Prediction via Graph Segment Training »
Kaidi Cao · Phitchaya Phothilimtha · Sami Abu-El-Haija · Dustin Zelle · Yanqi Zhou · Charith Mendis · Jure Leskovec · Bryan Perozzi -
2023 Poster: TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs »
Phitchaya Phothilimtha · Sami Abu-El-Haija · Kaidi Cao · Bahare Fatemi · Charith Mendis · Bryan Perozzi -
2023 Workshop: Machine Learning for Systems »
Xinlei XU · Dan Zhang · Phitchaya Phothilimtha · Beidi Chen · Yawen Wang · Divya Mahajan -
2023 Workshop: New Frontiers in Graph Learning (GLFrontiers) »
Jiaxuan You · Rex Ying · Hanjun Dai · Ge Liu · Azalia Mirhoseini · Smita Krishnaswamy -
2022 : Invited talk: Azalia Mirhoseini »
Azalia Mirhoseini -
2022 Poster: Mixture-of-Experts with Expert Choice Routing »
Yanqi Zhou · Tao Lei · Hanxiao Liu · Nan Du · Yanping Huang · Vincent Zhao · Andrew Dai · zhifeng Chen · Quoc V Le · James Laudon -
2022 Poster: TabNAS: Rejection Sampling for Neural Architecture Search on Tabular Datasets »
Chengrun Yang · Gabriel Bender · Hanxiao Liu · Pieter-Jan Kindermans · Madeleine Udell · Yifeng Lu · Quoc V Le · Da Huang -
2021 : Closing Remarks »
Jonathan Raiman · Mimee Xu · Martin Maas · Anna Goldie · Azade Nova · Benoit Steiner -
2021 : ML for Autotuning Production ML Compilers »
Phitchaya Phothilimtha -
2021 : Opening Remarks »
Jonathan Raiman · Anna Goldie · Benoit Steiner · Azade Nova · Martin Maas · Mimee Xu -
2021 Workshop: ML For Systems »
Benoit Steiner · Jonathan Raiman · Martin Maas · Azade Nova · Mimee Xu · Anna Goldie -
2021 Poster: Representing Long-Range Context for Graph Neural Networks with Global Attention »
Zhanghao Wu · Paras Jain · Matthew Wright · Azalia Mirhoseini · Joseph Gonzalez · Ion Stoica -
2020 Workshop: Machine Learning for Systems »
Anna Goldie · Azalia Mirhoseini · Jonathan Raiman · Martin Maas · Xinlei XU -
2020 Poster: Evolving Normalization-Activation Layers »
Hanxiao Liu · Andy Brock · Karen Simonyan · Quoc V Le -
2020 Spotlight: Evolving Normalization-Activation Layers »
Hanxiao Liu · Andy Brock · Karen Simonyan · Quoc V Le -
2020 Poster: PyGlove: Symbolic Programming for Automated Machine Learning »
Daiyi Peng · Xuanyi Dong · Esteban Real · Mingxing Tan · Yifeng Lu · Gabriel Bender · Hanxiao Liu · Adam Kraft · Chen Liang · Quoc V Le -
2020 Oral: PyGlove: Symbolic Programming for Automated Machine Learning »
Daiyi Peng · Xuanyi Dong · Esteban Real · Mingxing Tan · Yifeng Lu · Gabriel Bender · Hanxiao Liu · Adam Kraft · Chen Liang · Quoc V Le -
2020 Poster: Rethinking Pre-training and Self-training »
Barret Zoph · Golnaz Ghiasi · Tsung-Yi Lin · Yin Cui · Hanxiao Liu · Ekin Dogus Cubuk · Quoc V Le -
2020 Oral: Rethinking Pre-training and Self-training »
Barret Zoph · Golnaz Ghiasi · Tsung-Yi Lin · Yin Cui · Hanxiao Liu · Ekin Dogus Cubuk · Quoc V Le -
2019 : Poster Session 2 »
Hanson Wang · Yujun Lin · Yixiao Duan · Aditya Paliwal · Ameer Haj-Ali · Ryan Marcus · Tom Hope · Qiumin Xu · Nham Le · Yuxiang Sun · Ross Cutler · Vikram Nathan · Min Sun -
2019 : Coffee Break & Poster Session 1 »
Yan Zhang · Jonathon Hare · Adam Prugel-Bennett · Po Leung · Patrick Flaherty · Pitchaya Wiratchotisatian · Alessandro Epasto · Silvio Lattanzi · Sergei Vassilvitskii · Morteza Zadimoghaddam · Theja Tulabandhula · Fabian Fuchs · Adam Kosiorek · Ingmar Posner · William Hang · Anna Goldie · Sujith Ravi · Azalia Mirhoseini · Yuwen Xiong · Mengye Ren · Renjie Liao · Raquel Urtasun · Haici Zhang · Michele Borassi · Shengda Luo · Andrew Trapp · Geoffroy Dubourg-Felonneau · Yasmeen Kussad · Christopher Bender · Manzil Zaheer · Junier Oliva · Michał Stypułkowski · Maciej Zieba · Austin Dill · Chun-Liang Li · Songwei Ge · Eunsu Kang · Oiwi Parker Jones · Kelvin Ka Wing Wong · Joshua Payne · Yang Li · Azade Nazi · Erkut Erdem · Aykut Erdem · Kevin O'Connor · Juan J Garcia · Maciej Zamorski · Jan Chorowski · Deeksha Sinha · Harry Clifford · John W Cassidy -
2019 Workshop: ML For Systems »
Milad Hashemi · Azalia Mirhoseini · Anna Goldie · Kevin Swersky · Xinlei XU · Jonathan Raiman · Jonathan Raiman