We apply a reinforcement learning (RL) based approach to learning optimal synchronization policies used for Parameter Server-based distributed training of machine learning models with Stochastic Gradient Descent (SGD). Utilizing a formal synchronization policy description in the PS-setting, we are able to derive a suitable and compact description of states and actions, allowing us to efficiently use the standard off-the-shelf deep Q-learning algorithm. As a result, we are able to learn synchronization policies which generalize to different cluster environments, different training datasets and small model variations and (most importantly) lead to considerable decreases in training time when compared to standard policies such as bulk synchronous parallel (BSP), asynchronous parallel (ASP), or stale synchronous parallel (SSP). To support our claims we present extensive numerical results obtained from experiments performed in both simulated and real cluster environments. In our experiments training time is reduced by 44% on average in simulated and 28% in real cluster environments, respectively. Moreover, the learned policies generalize to multiple unseen circumstances.
This paper studies how to improve the synchronization policy for parameter server based distributed SGD algorithm. Existing work focus on either synchronous or asynchronous policy, which often results in straggler or staleness problems. Recent research proposes different ways to make the policy design fully adaptive and autonomous. This paper proposes a reinforcement learning (RL) approach and shows promising results to reduce the total training time on various scenario. 
A major challenge is to design the state and action spaces in the reinforcement learning setting. It requires the design can be generalized to different scenario, while ensuing efficient policy learning. Compared to existing policies such as BSP, ASP and SSP, RL has an advantage to adapt to non-stationary situations over the training process. Compared to other existing approaches, RL could be fully data-driven. 
The paper formalizes an RL problem by minimizing the total training time to reach a given validation accuracy. To minimize the number of actions, only 3 actions coming from BSP, ASP and SSP are used. The state space is designed to capture similar loss curves, and to be independent of the number of workers (as much as possible). A policy network is used make decisions after trained with exiting methods. 
Numerical results validate that the RL policy improves the training time compared to BSP, ASP and SSP. Different number of works and models, dataset are also tested to show the RL policy is generalizable to unseen scenario. Although all the results are simulated in a controlled environment, Figure 4 gives a very interesting illustration showing the advantage of using the RL policy. I still have detailed comments (see below), but I find the paper well written, and the author(s) has obtained promising results.
Detailed comments:
-	The validation accuracy 88% on MINST seems pretty low to me to stop the algorithm, in particular when training multiple layer neural networks. What would happen if the accuracy is increased, can the RL approach still find a good policy? What about the validation accuracy on CIFAR-10?
-	I still have some concern of the computation time obtain the RL state per step. In particular, the time cost to compute the loss L on different weights w. How do you address this issue? Is L computed on the validation set? What is its size? This parameter seems to me highly sensitive when the policy is used for different dataset, in particular the dataset vary. It would be better to have more discussions in the paper or appendix. 
-	What is the final test accuracy on the trained models using different policies? This allows us to see whether the approach has not over-fitted to the training/validation set. 
Some typo:
Page 2 line 2 AP -> ASP
Page 5, last line 4: converge -> convergence