Forecasting Deep Learning Dynamics with Applications to Hyperparameter Tuning

1 INTRODUCTION . Machine learning researchers working with deep neural networks spend time looking at training curves and asking themselves the question : How would the results change if I modified some hyperparameters ? They run many experiments and , with time , develop a better understanding of how modifications in various hyperparameters affect the learning dynamics of our models . In this work , we attack this problem of understanding deep learning dynamics with deep learning itself . To this end , we collect a data-set of training curves of deep learning models with a diverse set of hyperparameter schedules . Next , we train an autoregressive deep model to predict the training curves conditioned on the schedule . More concretely , our model observes the setting of a small number of hyperparameters , such as learning rate , weight decay , or dropout rate . Then , it predicts a small number of values that a researcher would usually look at , such as training and validation loss and accuracy . We predict how these values will change a few hundred training steps later . Then , the model gets new settings for the next few hundred training steps , and predicts the next values . The model is autoregressive , so it uses the history of its own predictions from previous time-steps to predict the next ones . We study how to do time-series prediction with autoregressive models even in the presence of stochastic behaviour . We study different losses and introduce a loss weight scheme that allows to use Transformer models for stochastic time-series predictions . We show that the new method we introduce outperforms baselines and is able to faithfully approximate the learning dynamics of deep learning models , as shown in the examples in Figure 1 . Having a model of learning dynamics is interesting in its own right , as it allows to study the training process much faster , without running full experiments . But it also opens up the possibility to learn hyperparameter schedules by using reinforcement learning in the model . We experiment with this and validate our forecaster : the policies learned using reinforcement learning and the forecaster work well when training real models . Running the forecaster is of course much faster than running real experiments , so learning a policy in the forecaster is thousands of times faster . 2 TIME-SERIES MODELING WITH TRANSFORMER . Transformer ( Vaswani et al. , 2017 ) is an autoregressive neural architecture for discrete sequence prediction and transduction . Its variants achieve state-of-the-art results in machine translation and language modeling ( Shoeybi et al . ( 2019 ) , Radford et al . ( 2019 ) ) . The architecture is notably composed almost entirely of self-attention layers , with no recurrent connections . This allows the model to consider the entire sequence when predicting the next token , while being very efficient to train using GPUs , due to the massive paralleism allowed by attention layers . In this work , we show that Transformer can also be used to predict time series – concretely , training curves of deep learning models . One obvious way to apply Transformer to time-series prediction would be to try to directly predict the next value in the sequence with L2 loss . We show that this is not a good strategy when sequences can be stochastic and introduce a discretization technique that reduces the problem to the standard discrete sequence prediction setting and that yields good results . 2.1 DISCRETIZATION . In order to model time series using Transformer , we represent each number in the sequence in fixed-precision , base-k encoding . We concatenate the consecutive base-k digits corresponding to each number , starting from the most significant one . To predict multiple time series at a time , we arrange the representations of a single element of each time series in a fixed order , and then concatenate the representations of each time step , as shown at Figure 2 . This way we can model the joint probability distribution of multiple time series using a single model . We train the Transformer decoder to predict consecutive symbols in this discrete sequence autoregressively . We use cross-entropy as the loss function , weighted with exponential decay according to the significance of each digit : L ( t , p ) = n−1∑ i=0 βiH ( ti , pi ) where H denotes the cross-entropy , ti is are the true digits , pi is the model prediction and i indexes digits in the fixed precision , base-k encoding of a number , starting from the most significant one . 2.2 RESULTS ON SYNTHETIC DATA . We present an experiment on synthetic data , demonstrating the ability of Transformer to model a distribution over time series . The experiment is simple by design , devised so that it is easy to measure the diversity of data sampled by the model , while retaining high-level similarity with the real-world data . Diversity in our case is crucial because , as described in the next section , the learned model is used to provide experience for training an reinforcement learning agent . We want this experience to be as diverse as in the real environment , so the agent can learn to be robust to various future scenarios . We train Transformer on a dataset of synthetically generated curves calculated using the formula : xi = 1− 1 1 + iα +N ( 0 , σ2 ) where i is an integer from the interval [ 1 , N ] and σ is scale of the noise , set to 0.01 . The dataset is designed to mimic training accuracy curves , starting from 0 and converging to 1 in the limit . Rate of convergence is controlled by α , sampled uniformly from the interval [ 0.5 , 5 ] for each curve . To measure the diversity of the generated curves , we estimate the parameter α from a curve predicted by the model , by averaging over pointwise estimates along the curve : α̂ = 1 N N−1∑ i=0 t ( 1 xi − 1 ) We then visualize the distribution of α̂ in a histogram . As shown in Figure 3 , Transformer is able to generate diverse and convincingly-looking curves . As a baseline , we consider modeling the curves without discretization , using the Transformer decoder without embedding and predicting a sequence of real numbers , training the model with L2 loss . During inference , we add Gaussian noise with variance σ2 to the prediction at each step , to introduce stochasticity . As shown at Figure 4 , the curves generated this way tend to collapse to a single shape . 3 HYPERPARAMETER TUNING USING REINFORCEMENT LEARNING . As is standard in reinforcement learning , we frame the problem of controlling model hyperparameters during training as a partially-observable Markov decision process . Each transition in the MDP corresponds to n steps of training followed by an evaluation on a held-out set of 5 % of the training data . Observations are the current values of training and evaluation accuracy and loss . Note that the agent does not have access to full state of the environment , because that would encompass the current values of all parameters of the model , which is intractable as an input for the agent . Rewards are differences between the values of an optimized metric between two timesteps , so that the cumulative reward is equal to the final value of that metric . In all experiments we optimize for validation accuracy . Actions are relative changes in each controlled hyperparameter . Instead of predicting continuous actions using a parametric policy , we use a simple discretization scheme . Discretization of the action space allows the policy distribution to be multi-modal and has been empirically shown to improve training of reinforcement learning agents in terms of average performance , stability and robustness to hyperparameters ( Tang & Agrawal , 2019 ) . Concretely , we predict the change independently for each hyperparameter , out of a discrete set of values { −50 % , −20 % , −5 % , 0 , +5 % , +25 % , +100 % } . The action space conveys the intuition that a hyperparameter schedule should be approximately continuous in the log space , while allowing for both significant changes in the hyperparameters between timesteps , and for more fine-grained control . Opposite values of relative change cancel each other out , so that a random walk in the action space has a median relative change of approximately 1 for every hyperparameter . 3.1 MODEL-FREE APPROACH . As a basic model-free approach to solving this problem , we use the Proximal Policy Optimization algorithm ( Schulman et al. , 2017 ) . PPO is a widely-used policy gradient method , notable for being stable and easy to implement , while significantly improving upon the sample-efficiency of vanilla policy gradient . It allows to perform multiple gradient updates based on the same amount of collected experience by optimizing a surrogate objective , ensuring that the policy distribution does not drift too far from the original one . We set the discount factor γ to 1 to get an unbiased estimate of the expected return . This does not cause the return to diverge , as the rewards are bounded and the trajectory length is constant . We calculate advantages with Generalized Advantage Estimation ( Schulman et al. , 2015b ) with λ = 0.95 . As a policy network , we use the decoder part of Transformer , without the embedding , so it accepts continuous input . The network has two heads , one predicting the distribution over actions , and the other the value ( an estimate of the future return ) . We add to the PPO objective the L2 loss of the value head with weight 1 and an entropy bonus with weight 0.1 . 3.2 MODEL-BASED APPROACH . As a more sample-efficient , model-based approach , we use Simulated Policy Learning ( SimPLe , Kaiser et al . ( 2017 ) ) . The algorithm consists of three phases , repeated in a loop , as shown in Figure 5 . First , a set of trajectories is collected , either using the current policy network ( which is randomly initialized ) or using an external data-set . Then a predictive model of the environment is trained on the collected data . After that , the policy is trained using PPO in the environment simulated by the model . The improved policy is used to collect new data , and the loop continues . We use a variant of SimPLe where we start the loop from training the model on a set of pre-collected data , and then resume the standard SimPLe algorithm , collecting new data in each iteration . To be able to use the Transformer decoder as an environment model , we discretize the entire history of observations and actions and concatenate their representations into one long sequence o1a1o2a2 ... on . Observations are discretized as described in subsection 2.1 . Actions are already discrete , so we just rewrite them using one discrete symbol per each controlled hyperparameter . We train the model to predict just the observation symbols , masking out the loss terms corresponding to action symbols . During infer- ence , we calculate the reward based on the difference of an appropriate metric in two consecutive predicted observations . For the policy network , we use the same architecture as for the environment model , except for the action and value prediction heads . This lets us pre-initialize the policy parameters with those of the trained environment model , which significantly accelerates the learning process . We conjecture that the attention heads in the policy network needs to attend to similar positions in the sequence as the attention heads in the environment model , so a policy network with parameters initialized in such a way needs to learn mostly the action and value prediction heads , which is much faster .

Summary: The paper investigates the possibility of learning a model to predict the training behaviour of deep learning architectures from hyperparameter information and a history of training observations. The model can then be used by researchers or a reinforcement learning agent to make better hyperparameter choices. The paper first adapts the Transformer model to be suitable to this prediction task by introducing a discretization scheme that prevents the transformer decoder's predictions from collapsing to a single curve. Next, the problem is formalized as a partially-observable MDP with a discrete action set, and PPO and SimPLe are introduced. The proposed model-based method is compared against a human and a model-free baseline training a Wide ResNet on CIFAR-10. The model-based method achieves better validation error than the other baselines that use actual data. Next, the method is compared against a human and a model-free baseline training Transformer models on the Penn Treebank dataset. While the human achieves the best performance at the end of the run, the proposed method appears to learn more quickly than the others and finishes with performance comparable to the model-free baseline.