Batch size selection by stochastic optimal control
Abstract
Stochastic gradient descent (SGD) and its variants are widespread in the ﬁeld of machine learning.
Although there is extensive research on the choice of step-size schedules to guarantee convergence of these methods, there is substantially less work examining the inﬂuence of the batch size on optimization.
The latter is typically kept constant and chosen via experimental validation.
In this work we take a stochastic optimal control perspective to understand the effect of the batch size when optimizing non-convex functions with SGD.
Speciﬁcally, we deﬁne an optimal control problem, which considers the entire trajectory of SGD to choose the optimal batch size for a noisy quadratic model.
We show that the batch size is inherently coupled with the step size and that for saddles there is a transition-time t ⇤ , after which it is beneﬁcial to increase the batch size to reduce the covariance of the stochastic gradients.
We verify our results empirically on various convex and non-convex problems.
Introduction
Stochastic gradient methods are extremely popular in the ﬁeld of machine learning [4] [3] [10].
Despite their simplicity they require carefully tuning some hyper-parameters, such as the step size and the batch size.
While the choice of step size has been an extensive area of research, including cyclic step sizes [1][9] or adaptive step sizes, such as AdaGrad [7], RMSProp [20], Adam [12], etc., there has been noticeably less work regarding the choice of the batch size.
This is somewhat surprising given that batch size tuning has been shown to have signiﬁcant advantages over tuning the step size [19][8].
A notable exception is a work by Balles et al.
[2] that proposes a greedy batch size selection based on maximizing the bound on the expected gain of a single SGD step.
Another work by De et al.
[6] proposes different increasing batch-size schedules, motivated by approximately constant SNR in gradient approximations, and provides theoretical guarantees for convergence.
[11] theoretically analyze the inﬂuence of batch size, step size and gradient covariance on the properties of the achieved minima.
They show that the ratio of step size to batch size determines a trade-off between the width of the minima, measured by the trace of the Hessian, and the expected terminal loss.
Incontrast to prior work that optimizes greedy objectives, we investigate the use of stochastic optimal control in order to study what is the optimal batch size when considering the entire trajectory of the stochastic process.
To the best of our knowledge, our approach to selecting the batch size is novel.
In fact, stochastic optimal control has only rarely been used in the ﬁeld of optimization, with the exception of [14], but for the problem of step size selection in one dimension. Their extension to the Has it Trained Yet?
Workshop at the Conference on Neural Information Processing Systems (NeurIPS 2022).
multi-dimensional case is based on a local diagonal-quadratic assumption.
Similarly, we consider the continuous-time representation of SGD and deﬁne a continuous control problem that selects the optimal batch size by solving the related Hamilton–Jacobi-Bellman equation (HJBE) analytically, without requiring a diagonal-quadratic assumption.
Furthermore, we verify the validity of our theoretical analysis empirically on various convex and non-convex problems.
We show that the derived optimal batch-size schedule only evaluates 19.7 % as many samples compared to switching to the maximal batch size for a 2D saddle point, and only 48.3 % as many samples compared to switching to the maximal batch size for a saddle point indimensions with two descent directions.
Continuous-time models for SGD
A common way to analyze SGD is to model it as a stochastic differential equations (SDE), which is a well-established approach in the ﬁeld of stochastic approximation [13, 14].
In the ﬁeld of machine learning, this approach was taken in [16] to examine the stationary distributions of a stochastic process and in [11] to determine factors inﬂuencing the minima found by SGD.
In [17] the authors use continuous-time models of mini-batch SGD and SVRG to derive convergence bounds.
In the case, where the loss function f : R d ! R can be written as the sum of individual functions f i , each corresponding to some data point i 2 [1 , . . . , n ] , that is min ✓ ⇥ f
( ✓ ) :
= 1 n
P ni [-=
-]=
[Equation]
f i ( ✓ [Equation] ⇤ , an update step in mini-batch SGD is of the form
[Equation]
where ↵⌘ is the step-size, in which ⌘ is the maximal allowed step-size, and ↵ is the adjustment factor as was also done in [14], and
[Equation]
where |B k | = m k for some m k ⌧ n .
Let the empirical covariance of r f i ( ✓ ) be denoted by ⌃ ( ✓ ) :
=n
P ni =1 ( r f i ( ✓ )  r f ( ✓ ))
( r f i ( ✓ )  r f ( ✓ ))
T , then by the assumption above, the covariance of r f B k ( ✓ ) is cov ( r f B k ( ✓ )) =
⌃ ( ✓ ) /m
k .
From this, an SDE of the following form can be derived (see A.1 for details):
[Equation]
where dB t is Brownian motion and m t denotes the time-dependent batch size Now that a continuoustime model of SGD is derived, different tools such as optimal control theory can be applied to analyze the effect of the batch-size when optimizing non-convex functions.
Optimal control
Now consider a dynamical system with state vector x t 2 R d and a control vector m t 2 R l
[Equation]
with a given function f : R d ⇥ R l ⇥ R !
R d .
The control m t is limited to the admissible set M on the ﬁxed time interval [0 , T ] , which is a time-invariant, closed, and convex subset of R l .
The cost functional being considered is
[Equation]
with given functions K : R d !
R and L : R d ⇥ R l ⇥ [0 , T ] !
R .
m ( t !
T, x ) denotes the forward ﬂow map [18] following the system of ODEs in Eq.
(4) with some batch-size schedule m t and the initial condition X t = x and ending at X T =: m ( t !
T, x ) .
The optimization problem
[Equation]
can be solved via the Hamilton-Jacobi-Bellman equation
[Equation]
In the next section we will formulate a continuous control problem that selects the optimal batch-size by optimizing the desired objective function and derive an optimal batch-size schedule.
Method
Consider the objective f ( ✓ ) =
✓ T A ✓ with ✓ 2 R d and A 2 R d ⇥ d .
Moreover, we assume for the sake of simplicity that the f i ’s are such that ⌃ ( ✓ ) =
⌃ = diag (  1 , . . .
d ) is diagonal with constants  i  0 , for i = 1 , . . .
, d .
Then the continuous-time representation of SGD is deﬁned by the update
[Equation]
Further, let A 2 R d ⇥ d be symmetric and thus diagonalizable, i.e. A = V T ⇤ V with orthogonal V and ⇤ =
diag (  1 , . . .
,  d ) .
(Otherwise we can choose A 0 = 12 ( A + A T ) ).
Note that the  i can be both positive and negative.
Theorem 4.1 The average dimension-decoupled dynamics are given by (see A.2 for a derivation):
[Equation]
with
P
[Equation]
)] , where the expectation is over the Brownian motion.
In a convex setting, a large batch size is generally desired because it corresponds to a smaller variance of the stochastic gradient, which in turn speeds up the convergence.
However, larger batch sizes also come at the expense of more expensive computation time.
In the following, we therefore optimize the loss value but also add a term that penalizes the size of the batch.
Additionally, we introduce a cost-weight-factor  2 [0 , 1] to weigh between using small batch sizes and having a low terminal loss.
Thecontinuous control problem with the above dynamics then becomes
[Equation]
where we use the short-hand notation J ( g, t ) :
=
J ( g 1 , . .
. , g d , t ) , m min and m max are the feasible range of batch sizes with 1  m min < m max  n , and mi ( t !
T, g i ) is the forward ﬂow map of the respective ODE in Eq.
(9) starting at g i : g i ( t ) and ending at mi ( t !
T, g i ) :
=
g i ( T ) .
This control problem can be solved via the HJB-equation from which we can derive the following batch-size schedule, depending on the eigenvalues  i of f ( ✓ ) (see A.3 for the derivation)
Theorem 4.
Depending on the sign of the eigenvalues  i , for i = 1 , . . .
, d we have the following batch-size schedule:
[Equation]
With a non-convex objective f ( ✓ ) we can assume w.l.o.g. that the eigenvalues are ordered, such that  1 < . . .
<
p  0 <  p
+1 < . . .
<
d .
Then the batch-size schedule is
[Equation]
with t ⇤ such that
[Equation]
In the case where d = 2 and  1 < 0 <  2 we can express t ⇤ explicitly as
[Equation]
We can see in Eq.
(12) and (13) that the step size is proportionally compensated by the batch size (ignoring ↵ in the exponent), which was already observed in [11] and in [2].
Looking at (15), we see that a change in the schedule only occurs, if |  1  1 | < |  2  2 | .
That is, moving along the negative eigendirection is less likely than along the positive eigendirection, which makes it more challenging to escape saddles of this type.
In the following we will verify our results empirically.
Experimental results
We validated our results empirically for a saddle point in 2D, which can be found in Fig.
We ran another experiment for d = 40 , which can be found in A.6.
The loss function f ( ✓ ) = 12 ✓ T A ✓ was chosen with A = diag (  0 .
001 , 0 . 1) , such that the positive eigenvalue is two orders of magnitude larger than the negative eigenvalue.
The experiment was repeated for 1000 runs of each 2000 iterations.
Given the parameter choice in (74), the transition time was calculated to be at t ⇤ = 1886 .
The batch-sizes were restricted to the interval [1 , 1000] .
The batch-size returned by the optimal schedule was rounded to the closest integer.
A more detailed description of the experiment setup can be found in A.5.
Wecompared the schedule to using just the minimal/maximal batch-size for the entire run and switching to the maximal batch-size at t ⇤ , which corresponds to the case if there is no cost for the size of the batch-size.
(An analysis on this can be found in A.4).
Additionally, we compared our schedule to the CABS rule proposed in [2] on a convex quadratic function for d = 40 dimensions. We can see in Fig. 2b that our proposed schedule only evaluates 10% as many samples as the CABS rule, but reaches approximately the same average terminal loss. Moreover, if we look in Fig.
2a at the average loss value reached when the CABS rule evaluates as many samples as our proposed schedule we can see that it is 77% higher than the terminal loss achieved by our schedule.
This indicates that the learning efﬁciency with respect to per sample evaluation varies over the course of optimization.
Of course in order to fully evaluate learning efﬁciency other aspects such as the number of parameter updates and the cost of the schedule evaluation needs to be taken into consideration as well.
[Figure]
[Equation]
[Figure]
Discussion
In this work we tried to understand what the theoretical optimal batch size is, which takes into account the entire trajectory of the continuous-time model of SGD, for optimizing both convex and non-convex noisy quadratic functions in multiple dimensions.
In practice, this approach could also be used to optimize more general functions in machine learning, by approximating the function with a local quadratic model.
However there are a few limitations to this approach.
In practice, the eigenvalues of the Hessian are unknown and need to be estimated.
The work of [21] presents a framework in which this could be potentially done.
Alternatively, one could optimize quadratic models that are typically used in trast-region methods [5].
Another minor limitation is that the derived schedule is continuous, but a batch-size can only be integer.
It is also worth mentioning, that in practice, there are other considerations for choosing the mini batch to achieve optimal performance, for instance to powers of two [15].
Future work could focus on validating the schedule on more general functions and popular machine learning benchmarks, such as MNIST, Cifar-10 or SVHN, where efﬁcient estimation of eigenvalues becomes relevant.