-
Optimizing Memory Mapping Using Deep Reinforcement Learning
Authors:
Pengming Wang,
Mikita Sazanovich,
Berkin Ilbeyi,
Phitchaya Mangpo Phothilimthana,
Manish Purohit,
Han Yang Tay,
Ngân Vũ,
Miaosen Wang,
Cosmin Paduraru,
Edouard Leurent,
Anton Zhernov,
Po-Sen Huang,
Julian Schrittwieser,
Thomas Hubert,
Robert Tung,
Paula Kurylowicz,
Kieran Milan,
Oriol Vinyals,
Daniel J. Mankowitz
Abstract:
Resource scheduling and allocation is a critical component of many high impact systems ranging from congestion control to cloud computing. Finding more optimal solutions to these problems often has significant impact on resource and time savings, reducing device wear-and-tear, and even potentially improving carbon emissions. In this paper, we focus on a specific instance of a scheduling problem, n…
▽ More
Resource scheduling and allocation is a critical component of many high impact systems ranging from congestion control to cloud computing. Finding more optimal solutions to these problems often has significant impact on resource and time savings, reducing device wear-and-tear, and even potentially improving carbon emissions. In this paper, we focus on a specific instance of a scheduling problem, namely the memory mapping problem that occurs during compilation of machine learning programs: That is, mapping tensors to different memory layers to optimize execution time.
We introduce an approach for solving the memory mapping problem using Reinforcement Learning. RL is a solution paradigm well-suited for sequential decision making problems that are amenable to planning, and combinatorial search spaces with high-dimensional data inputs. We formulate the problem as a single-player game, which we call the mallocGame, such that high-reward trajectories of the game correspond to efficient memory mappings on the target hardware. We also introduce a Reinforcement Learning agent, mallocMuZero, and show that it is capable of playing this game to discover new and improved memory mapping solutions that lead to faster execution times on real ML workloads on ML accelerators. We compare the performance of mallocMuZero to the default solver used by the Accelerated Linear Algebra (XLA) compiler on a benchmark of realistic ML workloads. In addition, we show that mallocMuZero is capable of improving the execution time of the recently published AlphaTensor matrix multiplication model.
△ Less
Submitted 17 October, 2023; v1 submitted 11 May, 2023;
originally announced May 2023.
-
Competition-Level Code Generation with AlphaCode
Authors:
Yujia Li,
David Choi,
Junyoung Chung,
Nate Kushman,
Julian Schrittwieser,
Rémi Leblond,
Tom Eccles,
James Keeling,
Felix Gimeno,
Agustin Dal Lago,
Thomas Hubert,
Peter Choy,
Cyprien de Masson d'Autume,
Igor Babuschkin,
Xinyun Chen,
Po-Sen Huang,
Johannes Welbl,
Sven Gowal,
Alexey Cherepanov,
James Molloy,
Daniel J. Mankowitz,
Esme Sutherland Robson,
Pushmeet Kohli,
Nando de Freitas,
Koray Kavukcuoglu
, et al. (1 additional authors not shown)
Abstract:
Programming is a powerful and ubiquitous problem-solving tool. Developing systems that can assist programmers or even generate programs independently could make programming more productive and accessible, yet so far incorporating innovations in AI has proven challenging. Recent large-scale language models have demonstrated an impressive ability to generate code, and are now able to complete simple…
▽ More
Programming is a powerful and ubiquitous problem-solving tool. Developing systems that can assist programmers or even generate programs independently could make programming more productive and accessible, yet so far incorporating innovations in AI has proven challenging. Recent large-scale language models have demonstrated an impressive ability to generate code, and are now able to complete simple programming tasks. However, these models still perform poorly when evaluated on more complex, unseen problems that require problem-solving skills beyond simply translating instructions into code. For example, competitive programming problems which require an understanding of algorithms and complex natural language remain extremely challenging. To address this gap, we introduce AlphaCode, a system for code generation that can create novel solutions to these problems that require deeper reasoning. In simulated evaluations on recent programming competitions on the Codeforces platform, AlphaCode achieved on average a ranking of top 54.3% in competitions with more than 5,000 participants. We found that three key components were critical to achieve good and reliable performance: (1) an extensive and clean competitive programming dataset for training and evaluation, (2) large and efficient-to-sample transformer-based architectures, and (3) large-scale model sampling to explore the search space, followed by filtering based on program behavior to a small set of submissions.
△ Less
Submitted 8 February, 2022;
originally announced March 2022.
-
MuZero with Self-competition for Rate Control in VP9 Video Compression
Authors:
Amol Mandhane,
Anton Zhernov,
Maribeth Rauh,
Chenjie Gu,
Miaosen Wang,
Flora Xue,
Wendy Shang,
Derek Pang,
Rene Claus,
Ching-Han Chiang,
Cheng Chen,
Jingning Han,
Angie Chen,
Daniel J. Mankowitz,
Jackson Broshear,
Julian Schrittwieser,
Thomas Hubert,
Oriol Vinyals,
Timothy Mann
Abstract:
Video streaming usage has seen a significant rise as entertainment, education, and business increasingly rely on online video. Optimizing video compression has the potential to increase access and quality of content to users, and reduce energy use and costs overall. In this paper, we present an application of the MuZero algorithm to the challenge of video compression. Specifically, we target the p…
▽ More
Video streaming usage has seen a significant rise as entertainment, education, and business increasingly rely on online video. Optimizing video compression has the potential to increase access and quality of content to users, and reduce energy use and costs overall. In this paper, we present an application of the MuZero algorithm to the challenge of video compression. Specifically, we target the problem of learning a rate control policy to select the quantization parameters (QP) in the encoding process of libvpx, an open source VP9 video compression library widely used by popular video-on-demand (VOD) services. We treat this as a sequential decision making problem to maximize the video quality with an episodic constraint imposed by the target bitrate. Notably, we introduce a novel self-competition based reward mechanism to solve constrained RL with variable constraint satisfaction difficulty, which is challenging for existing constrained RL methods. We demonstrate that the MuZero-based rate control achieves an average 6.28% reduction in size of the compressed videos for the same delivered video quality level (measured as PSNR BD-rate) compared to libvpx's two-pass VBR rate control policy, while having better constraint satisfaction behavior.
△ Less
Submitted 14 February, 2022;
originally announced February 2022.
-
Learning and Planning in Complex Action Spaces
Authors:
Thomas Hubert,
Julian Schrittwieser,
Ioannis Antonoglou,
Mohammadamin Barekatain,
Simon Schmitt,
David Silver
Abstract:
Many important real-world problems have action spaces that are high-dimensional, continuous or both, making full enumeration of all possible actions infeasible. Instead, only small subsets of actions can be sampled for the purpose of policy evaluation and improvement. In this paper, we propose a general framework to reason in a principled way about policy evaluation and improvement over such sampl…
▽ More
Many important real-world problems have action spaces that are high-dimensional, continuous or both, making full enumeration of all possible actions infeasible. Instead, only small subsets of actions can be sampled for the purpose of policy evaluation and improvement. In this paper, we propose a general framework to reason in a principled way about policy evaluation and improvement over such sampled action subsets. This sample-based policy iteration framework can in principle be applied to any reinforcement learning algorithm based upon policy iteration. Concretely, we propose Sampled MuZero, an extension of the MuZero algorithm that is able to learn in domains with arbitrarily complex action spaces by planning over sampled actions. We demonstrate this approach on the classical board game of Go and on two continuous control benchmark domains: DeepMind Control Suite and Real-World RL Suite.
△ Less
Submitted 13 April, 2021;
originally announced April 2021.
-
Online and Offline Reinforcement Learning by Planning with a Learned Model
Authors:
Julian Schrittwieser,
Thomas Hubert,
Amol Mandhane,
Mohammadamin Barekatain,
Ioannis Antonoglou,
David Silver
Abstract:
Learning efficiently from small amounts of data has long been the focus of model-based reinforcement learning, both for the online case when interacting with the environment and the offline case when learning from a fixed dataset. However, to date no single unified algorithm could demonstrate state-of-the-art results in both settings. In this work, we describe the Reanalyse algorithm which uses mo…
▽ More
Learning efficiently from small amounts of data has long been the focus of model-based reinforcement learning, both for the online case when interacting with the environment and the offline case when learning from a fixed dataset. However, to date no single unified algorithm could demonstrate state-of-the-art results in both settings. In this work, we describe the Reanalyse algorithm which uses model-based policy and value improvement operators to compute new improved training targets on existing data points, allowing efficient learning for data budgets varying by several orders of magnitude. We further show that Reanalyse can also be used to learn entirely from demonstrations without any environment interactions, as in the case of offline Reinforcement Learning (offline RL). Combining Reanalyse with the MuZero algorithm, we introduce MuZero Unplugged, a single unified algorithm for any data budget, including offline RL. In contrast to previous work, our algorithm does not require any special adaptations for the off-policy or offline RL settings. MuZero Unplugged sets new state-of-the-art results in the RL Unplugged offline RL benchmark as well as in the online RL benchmark of Atari in the standard 200 million frame setting.
△ Less
Submitted 13 April, 2021;
originally announced April 2021.
-
Monte-Carlo Tree Search as Regularized Policy Optimization
Authors:
Jean-Bastien Grill,
Florent Altché,
Yunhao Tang,
Thomas Hubert,
Michal Valko,
Ioannis Antonoglou,
Rémi Munos
Abstract:
The combination of Monte-Carlo tree search (MCTS) with deep reinforcement learning has led to significant advances in artificial intelligence. However, AlphaZero, the current state-of-the-art MCTS algorithm, still relies on handcrafted heuristics that are only partially understood. In this paper, we show that AlphaZero's search heuristics, along with other common ones such as UCT, are an approxima…
▽ More
The combination of Monte-Carlo tree search (MCTS) with deep reinforcement learning has led to significant advances in artificial intelligence. However, AlphaZero, the current state-of-the-art MCTS algorithm, still relies on handcrafted heuristics that are only partially understood. In this paper, we show that AlphaZero's search heuristics, along with other common ones such as UCT, are an approximation to the solution of a specific regularized policy optimization problem. With this insight, we propose a variant of AlphaZero which uses the exact solution to this policy optimization problem, and show experimentally that it reliably outperforms the original algorithm in multiple domains.
△ Less
Submitted 24 July, 2020;
originally announced July 2020.
-
Approximate exploitability: Learning a best response in large games
Authors:
Finbarr Timbers,
Nolan Bard,
Edward Lockhart,
Marc Lanctot,
Martin Schmid,
Neil Burch,
Julian Schrittwieser,
Thomas Hubert,
Michael Bowling
Abstract:
Researchers have demonstrated that neural networks are vulnerable to adversarial examples and subtle environment changes, both of which one can view as a form of distribution shift. To humans, the resulting errors can look like blunders, eroding trust in these agents. In prior games research, agent evaluation often focused on the in-practice game outcomes. While valuable, such evaluation typically…
▽ More
Researchers have demonstrated that neural networks are vulnerable to adversarial examples and subtle environment changes, both of which one can view as a form of distribution shift. To humans, the resulting errors can look like blunders, eroding trust in these agents. In prior games research, agent evaluation often focused on the in-practice game outcomes. While valuable, such evaluation typically fails to evaluate robustness to worst-case outcomes. Prior research in computer poker has examined how to assess such worst-case performance, both exactly and approximately. Unfortunately, exact computation is infeasible with larger domains, and existing approximations rely on poker-specific knowledge. We introduce ISMCTS-BR, a scalable search-based deep reinforcement learning algorithm for learning a best response to an agent, thereby approximating worst-case performance. We demonstrate the technique in several two-player zero-sum games against a variety of agents, including several AlphaZero-based agents.
△ Less
Submitted 3 November, 2022; v1 submitted 20 April, 2020;
originally announced April 2020.
-
Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model
Authors:
Julian Schrittwieser,
Ioannis Antonoglou,
Thomas Hubert,
Karen Simonyan,
Laurent Sifre,
Simon Schmitt,
Arthur Guez,
Edward Lockhart,
Demis Hassabis,
Thore Graepel,
Timothy Lillicrap,
David Silver
Abstract:
Constructing agents with planning capabilities has long been one of the main challenges in the pursuit of artificial intelligence. Tree-based planning methods have enjoyed huge success in challenging domains, such as chess and Go, where a perfect simulator is available. However, in real-world problems the dynamics governing the environment are often complex and unknown. In this work we present the…
▽ More
Constructing agents with planning capabilities has long been one of the main challenges in the pursuit of artificial intelligence. Tree-based planning methods have enjoyed huge success in challenging domains, such as chess and Go, where a perfect simulator is available. However, in real-world problems the dynamics governing the environment are often complex and unknown. In this work we present the MuZero algorithm which, by combining a tree-based search with a learned model, achieves superhuman performance in a range of challenging and visually complex domains, without any knowledge of their underlying dynamics. MuZero learns a model that, when applied iteratively, predicts the quantities most directly relevant to planning: the reward, the action-selection policy, and the value function. When evaluated on 57 different Atari games - the canonical video game environment for testing AI techniques, in which model-based planning approaches have historically struggled - our new algorithm achieved a new state of the art. When evaluated on Go, chess and shogi, without any knowledge of the game rules, MuZero matched the superhuman performance of the AlphaZero algorithm that was supplied with the game rules.
△ Less
Submitted 21 February, 2020; v1 submitted 19 November, 2019;
originally announced November 2019.
-
Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm
Authors:
David Silver,
Thomas Hubert,
Julian Schrittwieser,
Ioannis Antonoglou,
Matthew Lai,
Arthur Guez,
Marc Lanctot,
Laurent Sifre,
Dharshan Kumaran,
Thore Graepel,
Timothy Lillicrap,
Karen Simonyan,
Demis Hassabis
Abstract:
The game of chess is the most widely-studied domain in the history of artificial intelligence. The strongest programs are based on a combination of sophisticated search techniques, domain-specific adaptations, and handcrafted evaluation functions that have been refined by human experts over several decades. In contrast, the AlphaGo Zero program recently achieved superhuman performance in the game…
▽ More
The game of chess is the most widely-studied domain in the history of artificial intelligence. The strongest programs are based on a combination of sophisticated search techniques, domain-specific adaptations, and handcrafted evaluation functions that have been refined by human experts over several decades. In contrast, the AlphaGo Zero program recently achieved superhuman performance in the game of Go, by tabula rasa reinforcement learning from games of self-play. In this paper, we generalise this approach into a single AlphaZero algorithm that can achieve, tabula rasa, superhuman performance in many challenging domains. Starting from random play, and given no domain knowledge except the game rules, AlphaZero achieved within 24 hours a superhuman level of play in the games of chess and shogi (Japanese chess) as well as Go, and convincingly defeated a world-champion program in each case.
△ Less
Submitted 5 December, 2017;
originally announced December 2017.