Two-Tailed Averaging: Anytime Adaptive Once-in-a-while Optimal Iterate Averaging for Stochastic Optimization
Authors: Gábor Melis
Venue: Submitted to ICLR 2023 (2023)
Keywords: optimization
polyak
iterate averaging
anytime
adaptive
online

Abstract: Tail averaging improves on Polyak averaging's non-asymptotic behaviour by excluding a number of leading iterates of stochastic optimization from its calculations.
In practice, with a finite number of optimization steps and a learning rate that cannot be annealed to zero, tail averaging can get much closer to a local minimum point of the training loss than either the individual iterates or the Polyak average.
However, the number of leading iterates to ignore is an important hyperparameter, and starting averaging too early or too late leads to inefficient use of resources or suboptimal solutions.
Setting this hyperparameter to improve generalization is even more difficult, especially in the presence of other hyperparameters and overfitting.
Furthermore, before averaging starts, the loss is only weakly informative of the final performance, which makes early stopping unreliable.
To alleviate these problems, we propose an anytime variant of tail averaging, that has no hyperparameters and approximates the optimal tail at all optimization steps.
Our algorithm is based on two running averages with adaptive lengths bounded in terms of the optimal tail length, one of which achieves approximate optimality with some regularity.
Requiring only the additional storage for two sets of weights and periodic evaluation of the loss, the proposed two-tailed averaging algorithm is a practical and widely applicable method for improving stochastic optimization.


Under review as a conference paper at ICLR 2023
TWO-TAILED AVERAGING: ANYTIME ADAPTIVE ONCE-IN-A - A-WHILE OPTIMAL WEIGHT AVERAGING FOR BETTER GENERALIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Tail averaging improves on Polyak averaging's non-asymptotic behaviour by ex- cluding a number of leading iterates of stochastic optimization from its calculations. In practice, with a finite number of optimization steps and learning rate that cannot be annealed to zero, tail averaging can get much closer to a local minimum point of the training loss than either the individual iterates or the Polyak average. However, the number of leading iterates to ignore is an important hyperparameter, and starting averaging too early or too late leads to inefficient use of resources or suboptimal solutions. Our work focusses makes setting this hyperparameter even more difficult, especially the presence of other hyperparameters and overfitting. Furthermore, before averaging starts, the loss is only weakly informative of the final performance, which makes early stopping unreliable. To alleviate these problems, propose anytime variant of tail averaging intended for improving generalization not pure optimization, that has no hyperparameters and approximates the optimal tail all optimization steps. Our algorithm is based on two running averages with adaptive lengths bounded in terms of the optimal tail length, one of which achieves approximate optimality with some regularity. Requiring only the additional storage for two sets of weights and periodic evaluation of the loss, the proposed two-tailed averaging algorithm is a practical and widely applicable method for improving generalization.
improving generalization, which
INTRODUCTION
For the series of iterates produced by Stochastic Gradient Descent (SGD) (Robbins and Monro, 1985) to converge to a local minimum point of the training loss, the learning rate must be annealed to zero. Polyak averaging (Polyak and Juditsky, 1992; Ruppert, 1988) improves SGD and achieves a statistically optimal convergence rate by averaging all iterates to produce the final solution. Tail suffix averaging (Jain et al., 2018; Rakhlin al., 2011 1) takes this further and improves the non- asymptotic behaviour by dropping a number of leading iterates from the average, speeding up the decay of the effect of the initial state while allowing the learning rate to stay constant. Both of these properties are advantageous in practice, where a finite number of optimization steps are taken, and because large learning rates may bias optimization towards flatter and wider minima, which improves generalization (Hochreiter and Schmidhuber, 1997; Keskar et al., 2016). Focussing on large learning rates, flat minima, and generalization, Izmailov et al. (2018) propose Stochastic Weight Averaging (SWA), which takes the same form tail averaging but is motivated from ensembling point of view.
Tail averaging starts after given number of optimization steps. Setting this hyperparameter to minimize the training loss already poses some difficulties, which only become more pronounced and numerous in the context of generalization, our primary focus in this work.
Triggering averaging too early is inefficient as the average must grow long for the early weights to matter less.
Triggering averaging too late is inefficient as it does not use valuable information.
Due to interdependencies, tuning of other hyperparameters may become harder.

---

Under review as a conference paper at ICLR 2023
S
L
L
S
L
S
Figure I: Example evolution of the two running averages of weights over optimization steps. Continuous lines indicate optimization iterates being added to averages over course optimization. There always one short (S) and one long (L) average, with the long one having more iterates averaged and a better loss. When the loss with the short one would become better, the short average is renamed long (marked by arrowheads), and the short one is restarted from empty state (marked by discontinuities between double vertical lines). The long average thus always starts from the final state of the preceeding short average and incorporates more iterates until the new short average becomes at least as good in terms of the loss. In any interval labeled with L (from an arrowhead to a vertical bar), there is at least one point where the length of the long average is near optimal.
The learning curves have sudden drop the onset of averaging and not necessarily informa- tive of their eventual performance before that. Early stopping of seemingly unpromising training runs can cull solutions that would benefit more from averaging.
Motivated by these problems, we propose the two-tailed averaging algorithm with the following features.
Anytime: An estimate of the optimal tail is maintained at all optimization steps. Adaptive: It has hyperparameters. The number of weights averaged (the length of the tail) is determined adaptively based on the evolution of generalization performance. Optimal once in a while: The tail length achieves near optimality regularly.
The algorithm is very easy to implement. Its principal cost is the storage required for second running average, and also performs more evaluations generalization performance (e.g. of the validation loss). The main idea, sketched in Figure 1, is to maintain two running averages of optimization iterates: short and long one, with the long average being our estimate of the optimal weights.
RELATED WORKS
2.1
AVERAGING IN PURE OPTIMIZATION
Polyak averaging originally proposed (Ruppert, 1988; Polyak and Juditsky, 1992) computes the equally weighted average
0 t 8
of all iterates 0, from the optimizer up to the current time step t. The convergence rate Or was analyzed in the convex case with appropriately decaying learning rate. Beyond this strictest interpretation, Polyak (or Dolyak-Ruppert averaging may refer to using 6 without the convexity assumption, without decaying learning rate, with another optimizer such as Adam (Kingma and Ba, 2014).
In practice, where finite budget considerations override the asymptotic optimality guarantees offered by theory, Polyak averaging may refer to an exponential moving average (EMA) of the form
0u Ou.
0, (1 - + 3,8,-1
(t 1),
(2)
where 3, < 1 may be constant near 1 or it may be scheduled as by Martens (2020). The idea here is to improve the rate of decay of the effect of the initial error by downweighting early iterates.
Tail averaging (TA) (Jain et al., 2018), also known as suffix averaging (Rakhlin et al., 2011), 1), considers a finite optimization budget of steps with a constant learning rate. At the cost introducing a hyperparameter control the start of averaging, it improves the rate of decay of the effect of the initial error while obtaining near mearmminimax rates on the variance. Tail averaging is defined as
0, 0
8);
2

---

Under review as a conference paper at ICLR 2023
0. = 1 - +
02
(t> s).
We have discussed a few representative averaging methods intended for pure optimization but often repurposed practice for improving generalization by tuning their hyperparameters. Although there are many interesting developments in this area (Shamir and Zhang, 2013; Lacoste-Julics et al., 2012), we now move on to discussing averaging for improving generalization.
2.2 AVERAGING FOR GENERALIZATION
In work parallel tail averaging, Izmailov (2018) propose Stochastic Weight Averaging (SWA), an additional stage of optimization with constant or cyclical learning rate, that computes an equally the following way: With the high
weighted average iterates. SWA can motivated heuristically learning rate, it seeks out wider and flatter basins in the training loss surface to improve generalization, but the high learning rate also prevents from reaching the bottom of the basin, so the weights bounce around it, thus taking their average should land closer to the minimum point. The SWA algorithm is almost identical to tail averaging (except for possibly cyclical learning rate and a periodic subsampling of iterates), but it is motivated from the angle of ensembling and generalization.
Our goal is to improve generalization, so the decision of when to start averaging the weights should depend on generalization performance. Indeed, Merity al. (2017) propose an algorithm mueh like SWA, where averaging is triggered when the validation loss does not improve for a fixed number of optimization steps, which trades one hyperparameter for another and is sensitive to noise in the raw validation loss.
In other related work, Guo et (2022) investigate the repeated application of SWA. Their method is not informed by the validation loss and requires the schedule of multiple SWA stages to be prespecified. Finally, taking the exponential moving average of iterates is also sensitive to its hyperparameter, the decay rate.
In summary, existing averaging methods for generalization which behave well in practice all have one or more hyperparameters to govern the weighting of early iterates. Tuning these hyperparameters can be rather costly, particularly the presence of other hyperparameters and when training runs take a long time. Furthermore, estimate the optimal average at multiple optimization steps in general. We address these issues in the present work.
with their hyperparameters, these methods not flexible enough to
The rest of this paper is structured as follows. In Section 3, we formally define the problem to solve. In Section provide description of the algorithm, whose properties analyzed in Section 5. We verify our analysis experimentally in Section and discuss the validity of our assumptions in Section 5.1.
3 PROBLEM STATEMENT
Let θ be the parameter (or weight) space, 02 E θ (f E No) a series of iterates produced by stochastic optimization with Bg being the initial value, and : θ -> R the generalization loss function. The generalization loss can simply be the validation loss, or it may measure the classification accuracy while training is minimizing the cross-entropy loss, or it might evaluate on question answering language model trained with a masked language modelling loss.
We assume the generalization loss evaluated periodically, every at these points E [0, E, 2E, 3.....] where would like to know how many of the most recent iterates to average to minimize it Here and in the following, t and n (with or without subscripts) are assumed to be from No and [0, E, 2E, J..... respectively, and "loss" always refers f.
N optimization steps, and is
Defining the average of most recent A iterates up to time step f avg(t, A) 1 Ei=41-- O. define the optimal averaging length as O(t) arg min AE[I......] favgg(t A)), and our task is to approximate (n)) all evaluation steps n.
3

---

Under review as a conference paper at ICLR 2023
Two-tailed averaging with 2 running averages, a short one 0 and long one of with & number iterates. the loss with 8 becomes lower or equal than the loss with 8 then we switch the long average over to the short one and start a new short average from the empty state. The grayed out parts reset an average if has not improved for few evaluations. This heuristic makes the algorithm quicker to adapt when Assumption is violated (see Section 5.1).
Algorithm
Require: generalization loss function f. optimization iterates 0₁
Ensure: The sequence len; contains how many of the most recent iterates were averaged to compute loss; at the ith evaluation (i .1......J).
1: S, 05 L,04 0- <- 0, 0,0
b The first evaluation after this init will cause a switch.
2: for 1, 1,2,... do
3: S,0 t S- 58 4 S+1
D Add 0g to the short average
4:
L + 1, L+11
40"
> Add 8, to the long average
5:
ift mod E= then
> Evaluate the loss every E iterates
6:
F1, F$, ,F + f(,),, ((0) if ps is deteriorating then
7:
8:
S + 0,+0
Reset deteriorating $
9:
end
if Fs < or F is deteriorating then S. 0 F + 0, S, 08, FS
Is the short average better? b Switch: set L to S and reset S
11:
end if i 1--18 1/F
13:
pl
pb then
the current iterate better? >> Reset both averages Report the loss and averaging length
15:
S. 0,0
16:
loss, len, f Fr
17:
else
18:
loss, len; t FL, L
> The long average is better, report
19:
end if
20:
end if
21:
end for
There is a trivial algorithm to find Onn), that saves all 0, and performs a search over E [l..... n) to minimize f (avv(n, A)). This has storage and evaluation cost proportional to n, which is prohibitive. Even assuming that f improves monotonically in up to its optimum, the cost is still proportional to Onn). Our proposed algorithm approximates (n) with constant cost.
4 THE ALGORITHM
Algorithm specifies two-tailed averaging (TTA) in pseudocode. The main algorithm (without the grayed out bits) works follows. The outer loop iterates over weights 6g produced by a stochastic optimizer, incorporating them into the short and long running averages Then, every steps, the loss is evaluated with the current weights 9r, with the short average g and with the long average 8², giving Fl, FS and FL. If F$ is better than FL, then we switch: the long average continues from the current short average and the short average restarted (see Figure I). After this, we know that the long average is strictly better than the short one, but Fl might be better still if the raw loss is improving fast, which is the case very early in training. In any case, we report the best of Fl and Fh. for that evaluation.
H² with lengths S and L.
The grayed out logic constitutes the so-called reset heuristic. By resetting the short and long averages if they haven't t improved for a number of evaluations, it is intended to handle cases where the averages become too long, perhaps due to optimization escaping from one basin of attraction to a better one or to the loss surface changing in non-stationary environment.
This algorithm is online as it only accesses the current weights. We analyze its other properties in the next section.
4

---

Under review as a conference paper at ICLR 2023
5 ANALYSIS OF THE ALGORITHM
Our analysis hinges on the following assumptions, which follow from, for example, monotonically decreasing loss and averaging producing diminishing returns as the length increases, and whose validity is discussed in Section 5.1.
Assumption 1. For all n. a function of A E [0,
favvg(, A)) is monotoni-
cally decreasing. where (g(n)= = [O(n)/E] E. That is. for any given evaluation step n, averaging over more iterates from the past monotonically improves f until about the optimum length.
Assumption 2. For all and all n++ (n),,, (avg(n, O &
f(avg(n. (+)), where (m) = [ E E. That is, averaging slighly more iterates than optimal is better than averaging lot more.
Assumption 3. Vn: 3m,: O(n + m.) O(n) nss that is, the optimal average drops some iterates over a sufficiently long interval.
Assumption& 4. (()) & O(n E)). that is, the optimal number of weights to average is monotoni- cally increasing from one evaluation the next.
Let S(1), 0 (t), L(t), and (t) stand for the values of variables S, 05, L, and 0c in Algorithm
respectively, after times through the outer loop. Similarly, let (t) 0 (t), L'(),, and g (t) stand for the values of the same variables at the same iteration but after Line 4 (i.e. before the possible switch of the short and long averages). Furthermore, we introduce the shorthands fx((t) = f(avg(t, (())) for X E SS L, L'L'O) O} with f* (t) too if X(t) O.
Definition 1 (Switch point). We say that is switch point af Line 11 is executed, that is, when the short average becomes at least as good as the long average with respect to the loss, and consequently the short average is made the long average while short average is started. We denote the most recent switch point before iteration t with (t) where SP(t) < t. If there is no such switch point, then SP(t) -1.
Assumption 4 states that the optimal averaging length monotonically increases, without loss of generality, to simplify the analysis, we assume throughout that the raw loss F'l has already been eclipsed by FL the first evaluation. For the analysis, also assume that the reset heuristic (the grayed out parts in Algorithm 1) cannot trigger.
Proposition 2 (Basic properties). Vt > F: and Vn > E:
(i) E (m), E [(n)
(averaging lengths are multiples of the evaluation period) (because they only increase except at switches where S is reset to 0)
(ii) S(m) < L(n) <
(ii) (n) (n)
(due to the switching logic)
(iv) L(t) = ((())+)((SPt))) SP(()) * -1 else L(t) = S(t)
(due to the switching logic)
Proposition 3 (Bounds for the averaging lengths). The lengths of the short and long averages are bounded as ((n) < (nn) and L(n) < 20(n) + E.
Proof. We prove (n) < (n) by contradiction. Suppose ( & S(no) for some ng. As (n)) is monotonically increasing and ((n) increases by E, there exists n & no such that O (n) = 5(n)
(Proposition 2). Since (m)) < L(n) (by (ii) of Proposition 2), from Assumption 2 and 0 5
S(n) < L(n). we have that f (n) < [m), which contradicts (iii) of Proposition
Next we prove L(n) < 20(n) + E. From (iv) of Proposition 2, we have that at the beginning, when there has not yet been a switch, L(t) = (t); else L(t) = S(t + S' (SP(t)) for all t. In the first case, L(n) = S(t) O(n)= 20(n) + E, and we done.
In the second, the usual case, L(n) = S(n) + (SPP())) That is, the length of the current long average is the sum of the lengths of the current and the previously finished short average. Since S(n) < (n) and ('(SP(n)) = P(S)(() -E) + E, so L(n) < 0(n) + (SP(n)) E) + E, from which L(n) < ((n) + (SP((n) - E) + E. Finally, from the monotonicity of O in Assumption 4, O(S((n) E) & O()), we get L(n) < 20(n)= + E.
Proposition 4 (Infinite number of switch points). Switch points keep coming, that is, Vn: En, > n: (n.) -1.
5

---

Under review as a conference paper at ICLR 2023
(()))
[avg(n, A)]
f(avg(n 2))
f(avg(n 2),
L(n)
*
Optimal averaging length (()
Length of the short average S(.)
(
Length of the long average T(.)
loss
5 + E)
Lnn E)
O + E)
2E
5
25
O + 2F
0
2
3
4
6
A/E
Figure 2: Idealized illustration of switching. The three curves show the loss as a function of averaging length at three subsequent ovaluations (at optimization steps n, n E, and n + 2E). The raw loss keeps improving, hence later evaluations have lower loss curves. The optimal averaging length increases, ((n) & O(n + E) 2E, where the loss of the short average dips below the loss of the
O(n 2)), as per Assumption 4, A long average, the long average discarded, the short average hecomes the long average, and new short average is started with zero length.
Proof. This follows from S(n) ((n) and that S(n) increases by between switch points, eventually catching On)), that grows more slowly by Assumption
Proposition 5 (Once-ine in-a-ni--n while optimality). While between any two subsequent switch points n1 and n22 the long average is nearly optimal at least once. Formally, 3n E [2]: L(n) = 0 0*(m), (n).
Proof. Since S(n) < (n) and E S(n) for all n, so ('mm1) < 0 ( Then it is either that S'(n)) = 0 (n,) or S' (m(n)) O(n)). Since at switch points L(n) (n)). in the former case, we conclude the proof with
Considering the latter case, L(n1) S'(m) < O(m). Also, switches happen when f (n) & ft' (n), but as per Assumption 1 this happen only if (m) < L(n). Thus for n2 to be a switch point, it must be that O(nz) & L(n2), and we get that L(n;) < O(n))= & ((ng)= & [(n2). Therefore, considering that L(n) increases linearly between switch points, there must be a point n where L becomes greater than or equal to O, and at that point 0 because E L(n) and O increases monotonically according to Assumption 4.
In short, we have shown that there is an infinite number of switch points, the long average cannot be more than about twice as long as optimal, and between any two switch points the long average gets as close the optimal length as possible given the periodic evaluation scheme. Our results are terms of lengths of averages, and relating the actual loss with the long average f to the loss with the optimal length f would be desirable. Here, informally point out that, all things being equal, the worse 1 gets relative to f the quicker fs is to catch up. making long periods of highly suboptimal solution less likely. Formalizing this notion requires making further assumptions about the www.wvvaaiing-eeeeeeeeeeeee length function (of the kind plotted in Figure 2) and would make analysis considerably more cumbersome.
5.1 ON FAILED ASSUMPTIONS
To augment the theoretical analysis, that is based on idcalized assumptions, make the following observations. The strongest assumption by far is Assumption 1. assumes that increasing the averaging length monotonically improves f until the optimum. In terms of the loss, the algorithm is fairly robust to when the assumption holds only approximately because small deviations of favgg(n A)) from monotonicity change switch times only when fs and ft. are close.
Assumption 2 says that averaging slighly more iterates than optimal (i.e. rounded up to the evaluation period) is better than averaging lot more. This is weak assumption. If it fails, the algorithm can

---

Under review as a conference paper at ICLR 2023
4.25
4.2
4.15
4.1
4.05
4-
3.9
Figure 3: Validation loss with two-tailed averaging (TTA, green), tail averaging (TA, orange), and an exponential moving average (EMA, blue) of weights on language modelling on Penn Treebank. For both TA and EMA, their hyperparameters (the start time and the decay rate) tuned to minimize the final loss, that all three have similar optima. TTA two factors make tuning easier early stopping much more reliable. Additionally, the noise in the raw effectively smoothed out. Note that while the TTA loss decreases monotonically, gentler and steeper slopes are manifest before and after switch points, respectively.
is hyperparameters and produces much better early solutions. These
surprise
fail detect when the short averages becomes longer than optimal, which delays switching until the next evaluation.
Assumption failing means that the optimal average incorporates all new iterates without ever dropping old ones. In this case, the short average, which is always shorter than optimal, will be constant number iterates behind and its loss will converge to the loss of the optimal average. the long average is shorter than optimal, then the same argument applies to it. If the long average is longer than optimal, then eventually a switch will happen. either case, loss the long average converges to that of the optimal average.
Regarding Assumption 4, (n)) & O(n E' can fail if, for example, the improvement of the raw loss accelerates, but that is rather uncommon and temporary occurrence. It may also fail if the raw loss has started to worsen due to overfitting or optimization has escaped from one basin to the next and the average is slowly climbing the ridge separating them. Yet another way for this assumption to fail is when the loss landscape changes during learning in a non-stationary environment. With the exception of accelerating improvement, these cases are likely to be caught by the reset heuristic, wherein the short and long averages are reset if their loss does not improve for a few evaluations (see the grayed out parts in Algorithm 1).
The reset heuristic can trigger when it shouldn't, i.e. when Assumption 4 holds. Such spurious reset makes the estimate of the long average worse either directly indirectly by delaying the next switch. Either way, without further violations of this assumption, the algorithm recovers by the next switch point.
All in all, we can expect the algorithm to display some degree of robustness to minor violations of the assumptions. In practice, we recommend choosing a reasonably large F. to reduce the noise originating from the stochasticity of optimization and the likelihood of the resulting violations.
6 EXPERIMENTS
Tail averaging or Stochastic Weight Averaging have been shown previously to beneficial not only in theory and on simulated data (Jain et al., 2018) but also in language modelling (Merity et al., 2017; Melis al., 2019) and image classification (Izmailov et al., 2018) experiments. Hence, in this work restrict our attention to experiments in single domain to corroborate the analysis in Section 5. Our goals are to i) verify that TTA is on-par with well-tuned TA and EMA. ii) explore the effect of basing the switching logic the training loss instead of the validation loss, fii) demonstrate robustness to the choice of evaluation period, iv) and check how well the assumptions made in Section 5 hold in practice.
7

---

Under review as a conference paper at ICLR 2023
Lenghh of the long average (L) number of evaluations of TTA in Figure 3. Note how the heights of both peaks and valleys (corresponding L(n) and (n(n) increase almost monotonically. Also, the correspondence between this figure and schematic Figure illustrated on one the and L intervals.
Figure 4:
In particular, trained recurrent language model with several hyperparameters Penn Treebank (Mikolov et al., 2010) using the Rectified Adam optimizer (Liu et al., 2019), with periodic evaluation every 1000 optimization steps. The hyperparameters EMA eq. (1). Figure 3 shows that the final validation losses with the three methods are very close, but early losses with TTA are much better. This is expected because TA and EMA not flexible enough to produce optimal averaging lengths at multiple points along the learning curve despite having an extra hyperparameter. Conversely, TTA has at least one nearly optimal solution between any two subsequent peaks (i.e. switch points) in Figure 4 despite having no hyperparameters.
tuned separately for TTA, TA eq. (3), and
We also tried the version of Algorithm where the switching logic was based on comparison of the training losses of the short and long averages instead of the validation losses, but the true validation losses were reported. On this particular language modelling task, the best validation loss with the modified algorithm worsened moderately (3.93 3.92) and was well below the raw validation loss (4.02). Results on test set exhibited the same gap. Since the modified TTA was minimizing the training loss, the smoothness of the reported validation losses observed in Figure 3 were lost in the process. Similar results were obtained by scheduling a learning rate drop without iterate averaging.
To explore the effect of the evaluation period E, tuned models with four times larger and four times smaller than in our previously discussed experiments. As expected, the best final results very close to each other, with shorter periods having advantage early in training the raw loss FL was more quickly eclipsed by FL.
In addition, found that the assumptions made in Section 5 held rather well in practice: the resulting loss FL in Figure 3 and the averaging lengths tended to change monotonically (see heights of peaks and valleys in Figure 4), making our length-based theoretical results more closely linked to the actual loss. When that was not the case, we found that the raw loss F had started to worsen due to overfitting or, much more rarely, optimization had entered new basin, violating Assumption Figure and Figure demonstrate the reset heuristic being triggered in these cases. Finally, TA and TTA having almost identical final validation losses weakly supports our assumptions, although conclusive demonstration would need to plot the results obtained with TA tuned separately for each evaluation, which is beyond our computational budget.
7 CONCLUSION
In summary, we presented variant of tail averaging and Stochastic Weight Averaging based on two running averages. Compared to them, two-tailed averaging requires additional storage for the second running average and relies on periodic evaluation of generalization performance. In return, TTA both removes hyperparameter and provides estimate of the optimal tail at all optimization steps. In effect, it makes hyperparameter tuning easier and early evaluation much more representative final performance, allowing it support early and anytime stopping better. Owing to its simplicity, low implementation cost and adaptivity, TTA is practical and widely applicable method for improving generalization.
8

---

Under review as a conference paper at ICLR 2023
4.2
4.15
4.1
4.05
40
Figure 5: Raw (green) and TTA (orange) validation loss with overfitting. The averaged loss hottoms out due to overfitting. Thus when the averages reset (twice), the validation loss does recover. Although the losses reported after reset are suboptimal, it does not really matter as a better loss was reported already.
(a) Validation loss with the long average (F")..
(b) Length of the long average (L).
Figure 6: Example of optimization entering a new basin. At 40E the validation loss bottoms out starts to slightly worsen. This is detected at 43E, and both averages are reset. With the averages now way too short, the loss spikes but then recovers.
Looking beyond the scope of this work, exploring the relationship between iterate averaging and learning rate schedules is promising direction as existing (Merity et al., 2017) and our own limited experimental results indicate that dropping the learning rate and tail averaging perform comparably.
The properties of our algorithm are particularly compelling for continual learning: by allowing the learning rate remain high and being able adapt the averaging length to changing cireumstances, TTA lets the model maintain high plasticity while reaping the benefits of averaging.
addition, averaging weights can be viewed as a cheap approximation to averaging predictions when the averaged weights reside in a region with suitable geometry. Averaging weights within such regions and averaging predictions over regions (each with its own weight average) could potentially achieve better loss than weight averaging alone at much lower storage and evaluation cost than pure prediction averaging. We leave these avenues for future work to explore.
8 REPRODUCIBILITY
Our main contributions, Algorithm 1 and its analysis are fully contained in the paper.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Hao Guo, Jiyong Jin, and Bin Liu. Stochastic weight averaging revisited. arXiv preprint arXiv:2201.0057,
2022.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1): 12)) 1997.
Pavel Izmailov, Dmitri Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson, Averaging weights leads to wider optima and better generalization. arXiv preprint arXiv:/803.003:30 8888000077, 2018.
Prateek Jain, Sham Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Parallelizing stochastic gradient descent for least squares regression: mini- bani-batching, averaging, and model misspecification. Journal of Machine Learning Research, 18, 2018.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint 1609.04866 2016.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980, 2014.
Simon Lacoste-Julien, Mark Schmidt, and Francis Bach. A simpler approach to obtaining an o (1/t) convergence rate for the projected stochastic subgradient method. arXiv preprint arXiv: 2. 2002, 2012.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint 1.Xiv:990.003260 2019.
James Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research, 21(1):5776-885 2020.
Gábor Melis, Tomáš Kočiský, and Phil Blunsom. Mogrifier LSTM. arXiv preprint arXiv:1909.0773 2019.
Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing LSTM language models. arXiv preprint 1rXiv:1708.02182. 2017.
Tomas Mikolov, Martin Karafiát, Lukas Burget, Jan Cernocky, and Sanjeev Khudanpur. Recurrent neural network based language model. In Interspeech, volume 2. page 3, 2010.
Boris Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal on control and optimization, 30(4):838-855, 1992.
Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. arXiv preprini 1rXiv:1109.566 100.5647, 2011.
Herbert Robbins Sutton Monro. stochastic approximation method. Herbert Robbins Selected Papers, pages 102-109. Springer, 1985.
David Ruppert. Efficient estimations from slowly convergent Robbins- Monno process. Technical report, Cornell University Operations Research and Industrial Engineering. 1988.
Ohad Shamir and Tong Zhang. Stochastic gradient descent non-smooth optimization: Convergence results and optimal averaging schemes. In International conference on machine learning, pages 71-79. PMLR, 2013.
10


Reviews:
[{"review_id": "OlkBlD-qgGF", "reviewer": "Reviewer_bYZd", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Tail-averaging is a very useful trick, in both theory and practice, to improve the performance of iterative learning/optimization method. Yet, it introduces an important hyperparameter, the starting point for the tail-averaging. This work proposes a two-tailed averaging method as an adaptive and hyperparameter-free alternative to the tail-averaging method. Under suitable conditions, it is proved that two-tailed averaging has many good properties, inclduing being (1) anytime, (2) adaptive, and (3) optimal once in a while. Besides to the method and theory, some experiments are conducted to verify the advantage of two-tailed averaging. ", "summary_of_the_review": "Please see above. Due to the mentioned issues I would like to keep my score low. I would like to hear the authors rebuttal to see whether or not I have misunderstood anything. ", "strength_and_weaknesses": "# Strength \n1. Overall I think the paper is written well. I find the problem well motivated, that treating with the hyperparameter in tail-averaging might be painful, so we need a fix to that. \n2.  The proposed two-tailed averaging is novel to my knowledge. \n3. Both theory and experiments are presented for better demonstrating the advantages of  two-tailed averaging. \n\n# Weakness\n1. It is claimed that two-tailed averaging is \"adaptive\" and \"hyperparameter-free\". However I find this a bit misleading, because the procedure needs to be able to evaluate the \"loss\" periodically. Note that the \"loss\" refers to the targeted objective to be minimized if I understand correctly. For example, in the context of stochastic optimization [Jain et al, 2018], the \"loss\" refers to the population loss. But being able to evaluate the true loss is a very strong requirement, and makes the point of \"hyperparameter-free\" less interesting. For example, one can easily come up with a variant of tail-averaging that is \"hyperparameter-free\" provided the ability to evaluate the loss. Let us consider the setting of [Jain et al, 2018] again, one can simply save all the SGD iterates, and evaluate the loss over all possibilities of tail-averaging (or any general weighted averaging of the iterates), and output the optimal one --- this is \"hyperparameter-free\" and \"optimal\". I would like to see some discussions on the requirement of accessing the true loss. I just think this is too strong in many settings. \n\n\n2. I think the current presentation is somewhat \"abstract\". Perhaps the authors can discuss a potential application of the proposed two-tailed averaging in some theory/application problem? As the authors mentioned [Jain et al, 2018] many times, could you formally state a bound for two-tailed averaging in the setting of learning least squares with SGD, and compare it to the one obtained by tail-averaging? Being \"anytime\", \"adaptive\" and \"optimal once in a while\" sounds very cool but I would like to see some more specific benefits of using two-tailed averaging. \n\n\n3. I also think the assumptions could have been checked with a specific exemplar problem. The current assumptions are just very abstract and I do not immediately see how applicable/realistic they are. \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "3uHIR655yc", "reviewer": "Reviewer_bawy", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a tail averaging method for stochastic optimization that maintains two averages of different lengths in order to adaptively choose tail averging length.\nThe average produced by the longer of the two averages is provably close to optimal in length every once in a while.\nThe advantage of the proposed method over existing tail averaging methods is that it does not have tail averaging length as a hyperparameter.\nThe only hyperparameter is the number of iterations that pass between evaluating and comparing the two averages and the proposed method is robust to larger changes of this parameter.\nThe paper provides a theoretical analysis of the algorithm under a set of assumptions and argues how applicable the assumptions are in practice.\nThe paper also provides a demonstration that the method works well in practice also compared to other iterate averaging approaches for a toy example using an LM on Penn Treebank.", "summary_of_the_review": "I recommend to accept the paper due to the clear and elegant theoretical contribution that adds to the tail averaging literature, an important area of stochastic optimization, and the method is also directly applicable and practically feasible for many different machine learning applications.", "strength_and_weaknesses": "Strength:\n- Paper does a great job explaining the algorithm itself and also the motivation for it\n- Assumptions are clearly stated and even though they seem to be too strict for a nonconvex setting they should still be motivated in practice close to optimum and otherwise treated by the heuristics\n- Great figures, they take some time to think about but then make sense and illustrate useful properties once you understand them\n- Figure 3 nicely demonstrates that the approach competes with hyperparameter tuned alternatives but in a purely adaptive fashion\n- Figure 5 nicely demonstrates the TTA has the potential to improve over not using any tail averaging\n- Figure 6 nicely demonstrates the effect of a required reset heuristic from nonconvexity and again how much too-short-averaging-length hurts the validation loss\n\nWeakness:\n- Authors fairly point out that the analysis would become more cumbersome when trying to make a statement about the worst case difference between f^O and f^L but explain why this should not be a problem in practice. Under simplified assumptions where f is convex extending the analysis could potentially not be as cumbersome? Perhaps the loss is also monotonous in the averaging length before the averaging length becomes suboptimal? Having this analysis would present a nice addition for the paper making the theoretical analysis even more high quality.\n- Figure quality (e.g. of Figure 3, Figure 5) could be improved (the numbers on the axes look blury)\n- Would be nice to see the experimental evaluation also on modern architectures and in general on a broader set of evaluations\n- The structure of the experiments section could be clearer by separating the different research questions into separate sections making the presentation cleaner\n\nQuestions:\n- How much wall-clock time cost is added in practice by maintaining the additional sets of weights?\n- In this paper about anytime tail averaging: https://arxiv.org/pdf/1902.05083.pdf the averaging length k seems to depend on t (k=ct) which makes maintaining the average more expensive and still requires tuning c. But such approach does not seem to be used for comparison in the author's paper / they seem to compare to a constant averaging length \"only\". Does the k=ct approach behave qualitatively differen\ntly / more similar to the author's proposed adaptive method before the last iterate for which it was tuned?\n- The authors mention that tail averaging performs similarly to learning rate decay. Are there comparable adaptive methods for learning rate decay and how does the presented method compare in terms of com\nputational cost (memory and wall-clock time)?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "0Y5M1FgxGpj", "reviewer": "Reviewer_NMLB", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This submission proposes a new iterative averaging scheme for stochastic\ngradient methods. \nThe scheme, which works by maintaining two iterate averages\nat any given time, manages to be parameter free by setting the average length\nparameters dynamically based on the validation loss. \nAs such, it is parameter-free unlike suffix/tail averaging, where the tail \nlength parameter must be specified a-priori.\nThe authors prove that their \"two-tailed\" averaging scheme is optimal at\nspecific iterates infinitely often (\"once in a while\") under assumptions\non the behavior of the validation loss with respect to the length of the \naverage.\nThe submission concludes with several experiments on the Penn Treebank dataset.", "summary_of_the_review": "=== Update ===\n\nI appreciate the discussion with the authors, who have clarified several issues in the manuscript. I think the paper is improved, but I remain skeptical of the main assumptions required for the theoretical analysis. In particular, I think that Assumption 1 is not justified and I am reluctant to accept this submission without (i) an in-depth experimental justification for the assumption, or (ii) a revised theoretical analysis. Although I would normally increase my score to four reflect our discussion, I cannot due to the review form. As such, I will maintain my score of 3.\n\n=====\n\nThis submission is an interesting attempt to develop anytime methods for iterate \naveraging. I like two-tailed averaging scheme; perhaps it can be shown to admit\na near optimal convergence rate for the training loss.\nInstead of pursing this direction, the authors rely on difficult-to-trust\nassumptions on the validation for their theoretical analysis, which shows that the averaging\nscheme is once-in-a-while optimal.\nAs the experiments are limited to a single dataset and do not verify the assumptions, \nI feel the proposed method is not well-justified by theory or experiment. \nI recommend the submission be rejected.", "strength_and_weaknesses": "The main strength of the submission is the relative simplicity and elegance of\nthe averaging method.\nBy maintaining two averages --- one over a short time scale and one over a long time scale ---\nthe averaging scheme can capture both local improvement in the objective\nand long-term behavior of the optimizer. \nI think it is possible that the length of the short/long intervals could be\nset to give near-optimal anytime convergence rates for the training objective,\ne.g. something like Theorem 5 of Rakhlin et al. [1].\n\nThe major weaknesses of the submission are the following: (i) the theoretical\njustification for two-tailed averaging is unrealistic and \n(ii) the experimental evaluation is to small in scope to justify the \nmethod without strong theory. \n\n**Theoretical Justification**: The theory in this submission does not justify\nthe proposed method.\nMy major complaint here is that the authors assume exactly the properties of \nthe validation loss that they need for their averaging scheme to work.\nIn particular, Assumptions 1 --- monotonic behavior of the validation loss \nwith respect to the length of the average --- is deeply unrealistic for several\nreasons:\n\n- It is unlikely to be satisfied for stochastic optimization of non-convex models,\n    where monotone improvement of any train/validation/test metric is unlikely.\n- It makes a very strong assumption on the connection between the training objective\n    and the validation objective. That is, progress on the training objective\n    should roughly equate to progress on the validation metric. \n\nBecause of the weakness of the theoretical analysis, I can only evaluate\ntwo-tailed averaging as a heuristic method for setting the suffix parameter.\n\n**Empirical Evaluation**: The empirical evaluation is small scale and does\nnot investigate the assumptions on which the theory is based.\nOnly two experiments are presented, both on Penn treebank.\nNo comparison is made against other anytime averaging methods \n(see \"Missing References\"), making it difficult to establish a baseline for\nthe performance of two-tailed averaging.\n\nWhile two-tailed averaging performs well on Penn Treebank, one experiment does \nnot justify it's wide-spread use.\nIn particular, I disagree with the authors statement that two-tailed averaging \"is approximately equivalent \nto a series of TA which were tuned for each evaluation step separately\" and thus requires no other experiments;\nthis only holds under their assumptions.\nWhile the assumptions were tested via downstream metrics like behavior\nof the average lengths, this is only (approximately) addresses Assumption 3. \nChecking the assumptions properly requires carefully plotting the behavior of\nthe validation loss and optimal average length on multiple heterogeneous datasets.\n\nFinally, some experiments appear to be omitted from the manuscript.  \nSpecifically, the authors say: \"To explore the effect of the evaluation period E, we tuned models with four times larger and four times smaller E than in our previously discussed experiments.\"\nI do not see these results presented anywhere, so we can only take the authors\non their word that $E$ doesn't strongly affect performance. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject