AI for AI Systems and Chips
In the past decade, computer systems and chips have played a key role in the success of AI. Our vision in Google Brain’s ML for Systems team is to use AI to transform the way systems and chips are designed. Many core problems in systems and hardware design are combinatorial optimization or decision making tasks with state and actions sizes that are orders of magnitude larger than common AI benchmarks in robotics and games. In this talk, I will go over some of our research on tackling such optimization problems. First, I talk about our work on deep reinforcement learning models that learn to do resource allocation, a combinatorial optimization problem that repeatedly appears in systems. Our method is end-to-end and abstracts away the complexity of the underlying optimization space; the RL agent learns the implicit tradeoffs between computation and communication of the underlying resources and optimizes the allocation using only the true reward function (e.g., the runtime of the generated allocation). I will then discuss some of our recent work on deep reinforcement learning methods for sequential decision-making tasks with long horizons and large action spaces, built upon imitation learning and tree search in continuous action spaces. Finally, I discuss our work on deep models that learn to find solutions for the classic problem of balanced graph partitioning with minimum edge cuts. We define an unsupervised loss function and use neural graph representations to adaptively learn partitions based on the graph topology. Our method enables the first generalized partitioner, meaning we can train models that produce performant partitions at inference time on new unseen graphs.