Post: Device Placement with Cross-Entropy Minimization and Proximal Policy Optimization
Training deep neural networks requires an exorbitant amount of computation resources, including a heterogeneous mix of GPU and CPU devices. It is critical to place operations in a neural network on these devices in an optimal way, so that the training process can complete within the shortest amount of time. The state-of-the-art uses reinforcement learning to learn placement skills by repeatedly performing Monte-Carlo experiments. However, due to its equal treatment of placement samples, we argue that there remains ample room for significant improvements. In this paper, we propose a new joint learning algorithm, called Post, that integrates cross-entropy minimization and proximal policy optimization to achieve theoretically guaranteed optimal efficiency. In order to incorporate the cross-entropy method as a sampling technique, we propose to represent placements using discrete probability distributions, which allows us to estimate an optimal probability mass by maximal likelihood estimation, a powerful tool with the best possible efficiency. We have implemented Post in the Google Cloud platform, and our extensive experiments with several popular neural network training benchmarks have demonstrated clear evidence of superior performance: with the same amount of learning time, it leads to placements that have training times up to 63.7% shorter over the state-of-the-art.
This is a great work as it tackles an important problem: graph partitioning in heterogeneous/multi-device settings. There is an increasing number of problems that could benefit from resource allocation optimization techniques such as the one described in this work.  ML and specifically RL techniques have been recently developed to solve the problem of  device placement. This work addresses one of the main deficiencies of the prior work by making more sample efficient (as demonstrated by empirical results).   The novelty is in the way the placement parameters are trained: As oppose to directly train a placement policy for best runtime,  a softmax is used to model the distribution of op placements on devices (for each device among the pool of available devices.) At each iteration, the current placement distribution is refined by the conditional distribution representing the top 10% of the best placements. A KL-divergence loss is applied to minimize the distance between the old and new distribution. PPO is also used for further data sample reuse efficiency. The strengths: Great results: The results outperform existing approach based on Policy gradient by upto 58%.  The overhead of placement is low given that only 2400 samples were needed to achieve the results. This makes this work pottentially deployable.  Control experiments show improvement as a result of both KL minimization and PPO optimization. The results also show a clear win over prior policy gradient based work.  The paper is clear and easy to follow. The code is provided and looks very modular and easy to use.   Questions/suggestions:   What if the sampled placements from the conditional distribution at time t+1 produces placements that are all worse than the placements produces at iteration t? Has it ever happened in your experiments?  It would be good to show results on the amount of saving the optimized placements bring to the end-to-end training of the target models.  Th proposed model does not seem to generalize across graphs as it uses a one hot representation for nodes. Have the authors considered ways to incorporate graph embeddings into the model?  The reductions shows in Table 1 should compare Post with best baseline results, not with policy gradient.