U NDERSTANDING S HORT -H ORIZON B IAS IN S TOCHASTIC M ETA -O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Careful tuning of the learning rate, or even schedules thereof, can be crucial to effective neural net training.
There has been much recent interest in gradient-based meta-optimization, where one tunes hyperparameters, or even learns an optimizer, in order to minimize the expected loss when the training procedure is unrolled.
But because the training procedure must be unrolled thousands of times, the metaobjective must be deﬁned with an orders-of-magnitude shorter time horizon than is typical for neural net training.
We show that such short-horizon meta-objectives cause a serious bias towards small step sizes, an effect we term short-horizon bias.
We introduce a toy problem, a noisy quadratic cost function, on which we analyze short-horizon bias by deriving and comparing the optimal schedules for short and long time horizons.
We then run meta-optimization experiments (both ofﬂine and online) on standard benchmark datasets, showing that meta-optimization chooses too small a learning rate by multiple orders of magnitude, even when run with a moderately long time horizon (100 steps) typical of work in the area.
We believe short-horizon bias is a fundamental problem that needs to be addressed if metaoptimization is to scale to practical neural net training regimes.
I NTRODUCTION
The learning rate is one of the most important and frustrating hyperparameters to tune in deep learning.
Too small a value causes slow progress, while too large a value causes ﬂuctuations or even divergence.
From optimization theory perspective, learning rate decay is required to achieve convergence guarantee for stochastic gradient methods under certain conditions (Bottou, 1998).
While a ﬁxed learning rate often works well for simpler problems, good performance on the ImageNet (Russakovsky et al., 2015) benchmark requires a carefully tuned schedule.
A variety of decay schedules have been proposed for different architectures, including (Szegedy et al., 2014), staircase (VGG,
ResNet) (Simonyan & Zisserman, 2015; He et al., 2016).
Clever learning rate heuristics have resulted in large improvements in training efﬁciency (Goyal et al., 2017; Smith, 2017).
A related hyperparameter is momentum; typically ﬁxed to a reasonable value such as 0.9, careful tuning can also give signiﬁcant performance gains (Sutskever et al., 2013).
While optimizers such as Adam (Kingma & Ba, 2015) are often described as adapting coordinate-speciﬁc learning rates, in fact they also have global learning rate and momentum hyperparameters analogously to SGD, and tuning at least the learning rate can be important to good performance.
Because of this, it is not surprising that there have been many attempts to adapt learning rates, either online during optimization (Schraudolph, 1999; Schaul et al., 2013), by learning a policy ofﬂine corresponding to a learning rate schedule (Maclaurin et al., 2015), or even with an optimizer (Andrychowicz et al., 2016; Li & Malik, 2017; Chelsea Finn, 2017; ?
All of these approaches are forms of meta-optimization, where one deﬁnes a meta-objective, the expected loss after some number of optimization steps, and tunes the hyperparameters to minimize this meta-objective.
But because gradient-based meta-optimization can require thousands of updates, each of which unrolls the entire base-level optimization procedure, the meta-optimization is thousands of times more expensive than the base-level optimization.
Therefore, the meta-objective must be deﬁned with a much smaller time horizon (e.g. hundreds of updates) than we are ordinarily interested in for large-scale optimization.
The hope is that the learned hyperparameters or optimizer will generalize well to much longer time horizons.
However, we show that this is not achieved because of a strong tradeoff between short-term and long-term performance, which we refer to as short-horizon bias .
The basic reason for short-horizon bias can be seen in Figure 1, which shows the meta training losses under different horizon (number of look ahead steps) for various initial learning rate α 0 following a learning rate schedule α t
= α 0 1+ t 5 . One can see that the optimal learning rate varies with the length of the horizon, and shorter horizon prefers a more conservative learning rate.
As a result, short-horizon meta-optimizer would want to aggressively decay the learning rate to achieve immediate gains.
But this would come at the expense of long-term optimization progress.
The goals of this paper are to make this short-horizon bias more precise, provide empirical evidence for it, and to suggest a toy problem that can be used to analyze, understand, and potentially overcome this bias.
We ﬁrst deﬁne a toy problem to illustrate the fundamental tradeoff between short-horizon and longhorizon performance.
In particular, we analyze a noisy quadratic cost function based on (Schaul et al., 2013).
We consider this a good proxy for neural net training because second-order optimization algorithms have been shown to train neural networks in orders-of-magnitude fewer iterations (Martens, 2010), suggesting that much of the difﬁculty of SGD training can be explained by quadratic approximations to the cost.
In our noisy quadratic problem, the dynamics of SGD with momentum can be analyzed exactly, allowing us to derive the 1-step-horizon-optimal learning rate and momentum in closed form, as well as to ﬁt the long-horizon-optimal schedule using gradient descent.
We analyze the differences between the short-horizon and long-horizon schedules.
Because short-horizon performance is dominated by high-curvature directions, the short-horizon optimizer neglects low-curvature directions, leading to poor long-term performance.
Next, we consider the setting of gradient-based meta-optimization.
To measure short-horizon bias in a practical setting, we consider two idealized meta-optimization algorithms: an ofﬂine algorithm which ﬁts a learning rate decay schedule by running optimization many times from scratch, and an online algorithm which adapts the learning rate and momentum during training.
Since our interest is in studying the effect of the meta-objective itself, we allow the meta-optimizers sufﬁcient time to optimize their meta-objectives well.
We show that short-horizon meta-optimizers dramatically underperform a hand-tuned ﬁxed learning rate, and sometimes cause base-level optimization progress to slow to a crawl, even with moderately long time horizons (e.g. 100 or 1000 steps) similar to those used in prior work on gradient-based meta-optimization .
[Figure]
In short, we expect that any meta-objective which doesn’t correct for short-horizon bias will probably fail dramatically when run for a much longer time horizon than it was trained on.
There are applications where shorthorizon meta-optimization is directly useful, such as fewshot learning (Lake et al., 2011).
In those settings, shorthorizon bias is by deﬁnition not an issue.
But much of the appeal of meta-optimization comes from the possibility of using it to speed up or simplify the training of large neural networks.
In these settings, short-horizon bias is a fundamental obstacle that must be addressed for metaoptimization to be practically useful.
N OISY QUADRATIC PROBLEM
In this section, we develop a toy problem which demonstrates the short-horizon bias, but that we can analyze analytically.
In particular, we borrow the noisy quadratic model of Schaul et al.
; the true function being optimized is a quadratic, but in each iteration, we observe a “noisy” version with perturbed center but exact curvature.
We analyze the dynamics of SGD with momentum on this example, determine both the optimal and short-horizon-optimal learning rate schedules.
2. B ACKGROUND
Approximating the cost surface of a neural network with a quadratic function has yielded powerful insights and algorithms.
Second-order optimization methods such as Newton-Raphson and natural gradient (Amari, 1998) iteratively minimize quadratic approximations to a cost function.
Hessian-free (H-F) optimization (Martens, 2010) is an approximate natural gradient method which approximately minimizes the quadratic approximation using conjugate gradients; this was able to ﬁt deep neural networks in orders-of-magnitude fewer updates than SGD, suggesting that much of the difﬁculty of neural net optimization is captured by quadratic models.
In the setting of Bayesian neural networks, quadratic approximations to the log-likelihood motivated the Laplace approximation (MacKay, 1992).
Koh & Liang (2017) used quadratic approximations to analyze the sensitivity of a neural networks’ predictions to particular training labels, thereby yielding insight into adversarial examples.
Our primary interest is to ﬁnd a shared optimal learning rate and momentum across all dimensions.
The reason we focus on this instead of coordinate-speciﬁc learning rates is because we are assuming doing gradient descent in a preconditioned space.
Diagonal methods (Kingma & Ba, 2015), and even K-FAC (Martens & Grosse, 2015), can be viewed as gradient descent in a preconditioned space.
Since current methods do not provide a perfect preconditioning, we do consider various curvature on each dimension in our model.
Quadratic approximations have directly yielded insight into optimization hyperparameters.
In a deterministic setting, under certain conditions, second-order optimization algorithms can be run with a learning rate of 1; for this reason, H-F was able to eliminate the need to tune learning rate or momentum hyperparameters.
Martens & Grosse (2015) observed that for a deterministic quadratic cost function greedily choosing the learning rate and momentum to minimize the error on the next step is equivalent to conjugate gradients.
Therefore, by a general property of conjugate gradient (CG), the greedily chosen learning rates and momenta are optimal, in the sense that the greedy sequence achieves the minimum possible loss value of any sequence of learning rates and momenta.
This property fails to hold in the stochastic setting, however, and as we show in this section, the greedy choice of learning rate and momentum can do considerably worse than optimal.
2. A NALYSIS 2.2. N OTATIONS
We will primarily focus on the SGD with momentum algorithm in this paper.
The update is written as follows,
[Equation]
where L is the loss function, α ( t ) is the learning rate, and v ( t ) called velocity, with decay constant µ ( t ) , called momentum, at time t .
For simplicity, we also denote A ( · ) =
E [ · ] 2 + V [ · ] , where E and V denote expectation and variance respectively.
2.2. P ROBLEM FORMULATION
The problem is to ﬁnd the optimal learning rate and momentum for optimizing a noisy quadratic problem.
For every quadratic, one can choose a coordinate system such that the Hessian is diagonal.
We assume the gradient noise variance to be co-diagonalizable with Hessians.
The loss is written as,
[Equation]
where the expectation is taken over a data distribution P , and θ i ,
c i are the i th coordinate of the parameter θ and the sample c , and E c ∼P
[ c i ]
= θ ∗ i , V [ c i ] =
σ 2 i .
The stochastic gradient is written as ∂ L [Equation] i =
h i [Equation] θ i − c i ) .
WLOG, we assume θ ∗ i = 0 , so the stochastic gradient becomes,
[Equation]
The expected loss at timestep t is,
[Equation]
2.2. C OORDINATE -S HARING HYPERPARAMETERS
We ﬁrst consider the setting where the curvature already accounts for preconditioners, and hence the learning rate and momentum are shared across all dimensions.
We derive the greedy optimal learning rate and momentum by from the dynamics of the SGD with momentum.
Several observations allow us to compactly model the dynamics following SGD with momentum on the noisy quadratic model.
First, the expected loss can be expressed in terms of the variables E
[ θ i ] and V [ θ i ] .
Second, because the Hessians and noise covariances are diagonal, each coordinate evolves independently of the others.
Third, the second-order statistics of the iterate θ ( t ) i and the velocity v ( t ) i can be expressed in terms of the second-order statistics of θ ( t − 1) i and v ( t − 1) i .
these observations, we model the dynamics of SGD with momentum as a deterministic recurrence relation with sufﬁcient statistics V [ θ ( t ) ] , V [ v ( t ) ] , and Σ ( t ) θ,v =
Cov( θ ( t ) , v ( t ) ) .
Note that we drop the dimension subscripts because the dimensions evolve independently.
The dynamics are as follows:
Theorem 1 (Mean and variance dynamics) .
The expectations of the parameter θ and the velocity v are updated as,
[Equation]
The variances of the parameter θ and the velocity v are updated as,
[Equation]
Theorem 1 allows us to simulate a forward dynamics of how loss changes over time given any schedule of { α ( t ) , v ( t ) } t , and express the loss at timestep T as L ( θ ( T ) )
( { α ( t ) , v ( t ) } T − 1 t =1 ) .
It turns the stochastic problem into a deterministic one, and so we are able to optimize learning rate and momentum.
Further, we can get an analytic solution if we only do 1 -step lookahead, i.e., we can solve for the optimal learning rate α ( t ) ∗ and momentum µ ( t ) ∗ that optimizes the loss at timestep t + 1 , L ( θ ( t ) ) .
We call this as the greedy optimal learning rate and momentum.
Theorem 2 (Greedy optimal learning rate and momentum) .
The global optimal learning rate and momentum is given by,
[Equation]
Note that the greedy optimal learning rate derived in Schaul et al.
(2013) is a special case of Theorem 2, in which the momentum is assumed to be 0 .
Analytically solving for the optimal learning rate and momentum for greater than 1 -step look ahead is infeasible.
However, since Theorem 1 describes how the loss changes over time, we can numerically solve the problem of T -steps lookahead with
Autodiff softwares and do gradient descent on the loss at timestep T with respect to the learning rate and momentum.
2.2. C OORDINATE - SPECIFIC HYPARAMETERS
Now we consider solving the problem (Eq.(3)) using coordinate-speciﬁc hyperparameters.
This reduces to a one dimensional problem as each coordinate evolves independently of others.
In the vanilla SGD case, the greedy optimal learning rate is indeed optimal, in the sense that it achieves the minimum possible loss value of any sequence of learning rates.
[Figure]
Theorem 3 (Optimal learning rate) .
For all T ∈ N , the sequence of learning rate { α ( t ) ∗ } T − 1 t =1 that minimizes the loss at time step T , L ( θ T ) is given by,
[Equation]
Moreover, this agrees with the greedy optimal learning rate
[Figure]
Note that when all the coordinates have equal curvature and noise, the optimal learning rate is shared across all dimensions, so that applying the greedy optimal learning rate is optimal.
Secondorder methods try to precondition the problem so that the curvature is spherical.
The noise variance becomes spherical when the Fisher is a good approximation to the Hessian, so that the noise variance matches the curvature.
Therefore, one-step lookahead will work well with a sufﬁciently good natural gradient method.
2. E XPERIMENTS
We now compare the results we obtain if we do one-step lookahead vs. T -steps lookahead, where
T = 250 .
We chose a 1000 dimensional quadratic with the curvature distribution adopted from Li (2005), that which CG achieves the worst convergent rate.
Since the Fisher matrix is an approximation to the Hessian matrix, we assume that h
i =
V [ ∂ L ∂θ i ] , and hence σ i = 1 √ h
i .
We applied the greedy learning rate and momentum to this problem.
We then compared it with the learning rate and momentum obtained by minimizing the T -step losses formed by the dynamics given in Theorem 1.
We used Adam (Kingma & Ba, 2015), with learning rate 0.003 and 500 steps.
We set an upper bound for the learning rate to avoid the loss on any curvature direction becoming larger than its initial value, so as to assure the quadratic assumption.
We also showed the loss curves for a ﬁxed learning rate and momentum which was also optimized by Adam.
The results and the corresponding learning rate and momentum are shown in Figure 2 (a).
The optimal total loss decreases to a much lower value (4.25) than the loss obtained by greedy schedule (63.86) and ﬁxed schedule (42.19).
We observe that the optimal learning rate and momentum are much higher than the greedy learning rate and momentum in the beginning of the training.
We then look at the sum of the losses on thehighest curvature directions and 50 lowest curvature directions separately.
We ﬁnd that by following the optimal schedule, the loss on the high curvature directions do not decrease initially However, by Forward Backward L w g +
Forward Backward L w g + t+ ⍺ - + + - t keeping at a high learning rate and momentum, the loss on the low curvature directions decreases signiﬁcantly. On the other hand, by following a greedy schedule, the learning rate and momentum become small very early, which makes the loss on the high curvature directions decrease instantly, along with the total loss.
However, in the long term, since the learning rate is too small to make the loss on the low curvature direction decrease, it converged to a much higher loss in the end.
This gives a valuable insight into the reason behind the short-horizon bias we see in meta optimization.
As we will also see later in the real datasets, short horizon objective will often encourage the learning rate and momentum to be conservative, so as to achieve largest gain in short term, but performing poorly in the long term.
On the other hand, the short term horizon bias is also the result of stochasticity in the objective.
As observed by Martens & Grosse (2015), the greedy learning rate is the optimal learning rate schedule in the deterministic case.
We also validated it by setting the noise σ i to zero and run greedy schedule compared to the schedule found by gradient descent.
Both curves aligned, shown in Figure 2 (b).
The noise in the problem adds uncertainty to the objective, resulting in failures of greedy schedule.
G RADIENT - BASED META - OPTIMIZATION
We now turn our attention to gradient-based hyperparameter optimization.
A variety of approaches have been proposed which tune hyperparameters by doing gradient descent on a meta-objective (Schraudolph, 1999; Maclaurin et al., 2015; Andrychowicz et al., 2016).
We empirically analyze an idealized meta-optimization algorithm which is simpliﬁed in two ways.
First, we drop the algorithmic tricks used in prior work, and instead allow the meta-optimizer more computation than would be economical in practice.
Second, we limit the representational power of our meta-model: whereas
Andrychowicz et al.
(2016) aimed to learn a full optimization
algorithm, we focus on the much simpler problem of adapting learning rate and momentum hyperparameters, or schedules thereof.
The aim of these two simpliﬁcations is that we would like to do a good enough job of optimizing the meta-objective that any failures arise from deﬁciencies in the meta-objective itself (i.e. short-horizon bias) rather than incomplete meta-optimization.
However, we believe our ﬁndings are applicable to the more general setting without these two simpliﬁcations.
We don’t see any reason why cheaper meta-optimization methods would address the shortcomings of the meta-objective.
The weights of a full meta-optimizer would implicitly encode learning rate schedules.
For instance, Duchi et al.
(2011) implicitly uses a polynomial decay schedule because the denominator of the update uses the sum of the gradient variances rather than the mean.
Therefore, any sufﬁciently powerful meta-optimization method needs to take into account short-horizon bias.
3. B ACKGROUND 3.1. S TOCHASTIC M ETA D ESCENT
The basic idea of stochastic meta descent (SMD) (Schraudolph, 1999) is to perform gradient descent on learning rate, or any other differentiable hyper-parameters.
Any gradient based optimization can be unrolled as a computation graph (see Figure 3), and automatic differentiation is readily available in most deep learning libraries.
Two types of automatic differentiation (auto diff) methods are commonly used: forward mode and reversed mode .
Suppose a node on the computation graph has inputs x and outputs y .
In forward mode auto diff, gradients are computed as a dual to regular computations: y (cid:48) =
J x (cid:48) , where J is the Jacobian.
In contrast, reverse mode auto diff (a.k.a. backpropagation) computes the gradients in reversed order: ∇ x = J (cid:62) ∇ y .
Since J depends on the values of x for any non-linear functions, reverse mode auto diff needs to store the values of x until the gradients are computed.
Meta-optimization using reverse mode can be computationally demanding due to memory constraints.
Maclaurin et al.
(2015) got around this by cleverly exploiting reversibility to minimize the memory cost of activations.
Since we are optimizing only two hyperparameters, forward mode auto diff can be computationally cheap.
We provide the forward differentiation equations for obtaining the gradient of vanilla SGD learning rate.
Let ∂ θ t ∂α be θ (cid:48)
α,t , and ∂ L t ∂α be g α,t , and the Hessian at step t to be H t .
By chain rule, we get,
[Equation]
The Hessian-vector product in Equation 6 can be computed efﬁciently using reverse-on-reverse (Werbos, 1988) or forward-on-reverse automatic differentiation (Pearlmutter, 1994), which is linear in the running time of a forward pass, without explicitly computing the full Hessian.
Gradient based meta-optimization is feasible after the gradients to the hyperparameters are computed.
It is worth noting that, although SMD was originally proposed for optimizing SGD, in practice it can optimize any optimizer with few differentiable hyper-parameters, e.g. SGD with momentum, RMSProp (Tieleman & Hinton, 2012), Adam (Kingma & Ba, 2015), etc.
Moreover, the meta updates can be done by any type of gradient based optimizers as well.
[Table]
We provide details of SMD algorithm in Algorithm 1.
α is a set of hyperparameters (e.g. learning rate), and α 0 are inital hyperparameter values.
θ is a set of optimization
intermediate variables, such as weights and velocity.
η is a set of meta optimizer hyperparameters (e.g. meta learning rate).
Additionally, T is the look ahead window size, and M is the number of meta updates.
Simpliﬁcations from the original SMD algorithm.
In the original algorithm (Schraudolph, 1999), 1) each parameter dimension has an individual adaptive learning rate, 2) intermediate gradients ( θ (cid:48) α ) are accumualted throughout the process of training, 3) only one meta update step is applied.
Many of those designs are based on speed reasons, and as a consequence the original algorithm does not compute the exact gradients for meta- optimization.
Since we can afford extra computation to study the results of meta-optimization, we adapted the original version such that 1) only global hyper-parameters are tuned and exact gradients are computed, 2) M meta update steps are applied.
3. O FFLINE META - OPTIMIZATION
To understand the sensitivity of the best hyper-parameter towards the number of look ahead steps, we ﬁrst design an ofﬂine experiment on a multi-layered perceptron (MLP) on MNIST (LeCun et al., 1998).
Since the dynamics of training are different at the very start compared to later stages, we pretrain the network ﬁrst with manually tuned learning rate and momentum. [Figure] [Figure] [Figure]
Learnable decay schedule.
We include a parameterized learning rate decay schedule, on top of the SGD with momentum updates.
A standard decay schedule is the inverse time decay (Welling & Teh, 2011): α t
= α 0 (1+ tK ) β , where t is the number of training steps, β is the learning rate decay exponent, and K is the time constant.
In the SMD procedure, we optimize β together with α and µ .
Experimental details.
The network layer size is [784-100-100-10], with ReLU activations.
Weights are initialized with zero- mean Gaussian distribution with standard deviation 0.1.
The network is pretrained with 1000 SGD with momentum steps, with α = 1 , µ = 0 .
We train all hyper-parameters on log space using Adam optimizer, with 2k meta steps.
After one meta-step, we restore the weights back to its pretrained initialization (see Algorithm 1).
We re-parameterized the momentum with (1 − µ ) , the learning rate with the effective learning rate α eff =
α 1 − µ .
Figure 4 plots SMD optimization trajectories on the loss surface, which are obtained by using random samples of hyperparameters.
Multiple trajectories were run to ensure the local minima match with the ones plotted on the surface.
The ﬁnal hyperparameter values are not inﬂuenced much by their initializations.
Importantly, looking ahead for longer steps ends up with larger initial learning rate α , and much smaller learning rate decay exponent β .
The loss surface exhibit distinct landscapes, and the ﬁnal chosen β differs by two orders of magnitude between 100 and 20k steps.
If we use the optimal hyperparameters for 100 steps, the resulting training loss at 20k steps is three orders of magnitude larger than 20k steps.
3. O NLINE M ETA -O PTIMIZATION
In this section, we study whether online adaptation also has a short-horizon bias.
Different from the previous section, here we use Algorithm 1) to adapt the learning rate and momentum parameters during each stage of the regular training.
We experimented with both MLP on MNIST and CNN on
CIFAR-10 (Krizhevsky, 2009).
Experimental details.
The MLP network layer size is [784-100-100-10], with ReLU activations.
Weights are initialized with zero- mean Gaussian distribution with standard deviation 0.1.
The CNN network is the CIFAR-10 network adapted from (Jia et al., 2014), with channel size [3, 32, 32, 64], 2 × 2 max pooling after every convolution layer, ReLU activations, and fully connected layers of size [1024, 100, 10].
The models were pretrained to 1000 steps using the best manually tuned learning rate and momentum parameters.
100 Adam meta steps is applied every 10 steps of regular training.
We ﬁxed the learning rate decay exponent to be zero, and applied manual decay in the middle of the training such that all runs end up with effective learning rate 1e-4.
Figure 5 shows online SMD training curves compared to manually chosen SGD with momentum.
Online SMD training converge to similar learning rate values despite different initializations.
SMD experiments obtain similar learning rate schedules to the greedy noisy quadatic schedules in Figure 2.
The results can be explained by our noisy quadratic example: SMD tend to lower the learning rate aggressively in the beginning to obtain large decrease in high curvature directions, but later cannot move very far on small curvature direction.
It does not choose to increase the learning rate afterwards since otherwise it will increase loss on high curvature directions again.
C ONCLUSION
In this paper we look at the short-horizon bias in the meta-optimization objective.
We try to understand the problem by analyzing a noisy quadratic model.
We ﬁrst derive the greedy learning rate and momentum analytically, and compared the performance on with the optimal learning rate found by gradient descent.
We ﬁnd that the greedy schedule starts with a very low learning rate, to achieve instant decrease in high curvature directions, but it can no longer optimize the losses on the low curvature directions.
The optimal schedule, on the other hand, maintains a high learning rate in the beginning, so the losses on the low curvature directions decrease signiﬁcantly.
Due to this progress on the low curvature directions, the optimal schedule eventually reaches a much better solution.
We also showed that if we have a sufﬁcient second order method for preconditioning, the greedy optimal learning rate is indeed optimal.
Moreover, we ﬁnd the bias comes from the noise in the objective.
For this reason, the greedy optimal learning rate and momentum is optimal for deterministic objectives.
We ﬁnd the same effect of short-horizon bias in larger datasets, such as MNIST/CIFAR.
A P ROOFS OF T HEOREMS A. M ODEL D YNAMICS
The stochastic gradient descent with momentum is deﬁned as follows,
[Equation]
A.1. D YNAMICS OF THE EXPECTATION
We calculate the mean of the velocity v ( t +1) ,
[Equation]
We calculate the mean of the parameter θ ( t +1) ,
[Equation]
Let’s assume the following initial conditions:
[Equation]
Then Eq.(7) and Eq.(8) describes how E [ θ ( t ) ] changes over time t .
A.1. D YNAMICS OF THE VARIANCE
We calculate the variance of the velocity v ( t +1) ,
V [Equation] ( t
[Equation]
[Equation] [ µ ( t ) [Equation]
( t )
[Equation]
( t
[Equation]
t )
[Equation]
t
[Equation]
( µ
[Equation]
[Equation]
( t ) ]
+ (
[Equation]
[Equation]
[Equation]
t ) ]
[Equation]
[Equation]
The variance of the parameter θ ( t +1) is given by,
[Equation]
We also need to derive how the covariance of θ and v changes over time:
[Equation]
Let’s assume the following initial conditions:
[Equation]
Combining Eq.(9-11), we obtain the following dynamics (from t = 0 , . .
, T − 1 ):
[Equation]
A. G REEDY
A.2. U NIVARIATE CASE
The loss at time step t is,
L ( t +1)
=2 h ( E [ θ ( t +1) ] 2 + V [ θ ( t +1) ])
=2 h (cid:104)
( E [ θ ( t ) ] + µ ( t ) E [ v ( t ) ] − ( α ( t ) h ) E
[ θ ( t ) ])
2 + V [ θ ( t ) ] + ( µ ( t ) )
2 V [ v ( t ) ]
+ ( [-hσ ) 
 =α-]α ( t ) h ) 2 V [ θ ( t ) ] − 2 µ ( t ) α
( t ) h · Cov ( θ ( t ) , v ( t ) )
+ ( α ( t ) hσ ) 2 + 2( µ
( t ) Cov( θ ( t ) , v ( t ) ) − α ( t )
h V
[ θ ( t ) ]) (cid:105) =
2 h (cid:104) ((1 − α ( t ) h ) E [ θ ( t ) ] + µ ( t ) E [ v ( t ) ]) 2 + (1 − α ( t ) h ) 2 V [ θ ( t ) ]
+ ( µ ( t ) )
2 V [ v ( t ) ] + 2 µ ( t ) (1 − α ( t ) h )
Cov( θ ( t ) , v ( t ) )
+ ( α ( t ) hσ ) 2 (cid:105) =2 h (cid:104) (1 − α ( t ) h ) 2 (cid:16) E
[ θ ( t ) ] 2 + V [ θ ( t ) ] (cid:17) + ( µ ( t ) ) 2 (
cid:16) E [ v ( t ) ] 2 + V [ v ( t ) ] (cid:17) + 2 µ ( t ) (1 − α ( t ) h ) (cid:16)
E
[ θ ( t ) ]
E [ v ( t ) ] + Cov( θ ( t ) , v ( t ) ) (cid:17)
+ ( α ( t ) hσ ) 2 (cid:105) For simplicity, we denote A ( · ) =
E [ · ] 2 + V [ · ] , and notice that E [ θ ( t ) v ( t ) ] =
E [ θ ( t ) ]
E [ v ( t ) ]
+
Cov( θ ( t ) , v ( t ) ) , hence,
[Equation] In order to ﬁnd the optimal learning rate and momentum for minimizing L ( t +1) [Equation]
we take the derivative with respect to α ( t ) and µ ( t ) , and set it to 0 : [Equation]
A.2. HIGH DIMENSION CASE
The loss is the sum of losses on each direction:
L ( t +1)
= (cid:88)
i
[Equation]
Now we obtain optimal learning rate and momentum by setting the derivative to 0 ,
[Equation]

[Equation]
[Table]
A. C OORDINATE - SPECIFIC O PTIMALITY IN SGD
We now consider a dynamic programming approach to solve the problem.
We formalize the optimization problem of { α i } as follows.
We ﬁrst denote f ( t ) as the minimum expected loss at times step T (i.e., under the optimal learning rate) as a function of the mean E [ θ ( t ) ] and variance V [ θ ( t ) ] of θ ( t ) at step t .
[Equation]
Because we can express the E [ θ ( t ) ] and V [ θ ( t ) ] in terms of the E [ θ ( t − 1) ] and V [ θ ( t − 1) ] of θ ( t − 1) and the learning rate α ( t − 1) at the previous time step,
[Equation]
⇒ ( E [ θ ( t ) ] , ( V [ θ ( [Equation] [Equation] = (
[Equation] [Equation] E [ θ
[Equation] [Equation] − [Equation]
−
[Equation]
[Equation] ) V [ [Equation] we can solve the minimization problem recursively by,
[Equation]
Now, let’s ﬁrst derive the form of f for T, T − 1 , T − 2 for illustration.
Let A ( t ) =
( E [ θ ( t ) ]) 2 + ( V [ θ ( t ) ]) 2 .
We can ﬁnd a recurrence relation in terms of A :
[Equation]
Since f depends on E [ θ ] and σ only through A , so we instead write f as a function of A :
[Equation]
Now since A ( T ) is a function of α ( T − 1) , if we take the derivative of f ( t ) w.r.t.
α ( T − 1) and setting it to zero, we get:
[Equation]
Thus we can write A ( T ) in terms of A ( T − 1) and the optimal α ( T − 1) ∗ :
[Equation]
Therefore,
[Equation]
Now since A ( T − 1) is a function of α ( T − 2) , if we take the derivative of f ( T − 1) w.r.t.
α ( T − 2) and setting it to zero, we get:
[Equation]
Theorem 4.
For all T ∈ N , and k ∈ N , 1 ≤ k ≤ T , we have,
[Equation]
Therefore, the optimal learning α ( t ) at timestep t is given as,
[Equation]
Proof.
The form of f T − k can be easily proven by induction on k , and use the identity that,
[Equation]
The learning rate then follows immediately by taking the derivative of f T − k w.r.t.
α ( T − k − 1) and setting it to zero.