Butterfly Effects of SGD Noise: Error Amplification in Behavior Cloning and Autoregression
Authors: Adam Block
Dylan J Foster
Akshay Krishnamurthy
Max Simchowitz
Cyril Zhang
Venue: ICLR 2024 poster (2024)
Keywords: Deep Learning
Imitation Learning
Stability
Exponential Moving Average
Optimization
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Small fluctuations from minibatch SGD noise are amplified catastrophically in unstable feedback loops. Variance reduction and (especially) iterate averaging help a lot.

Abstract: This work studies training instabilities of behavior cloning with deep neural networks. We observe that minibatch SGD updates to the policy network during training result in sharp oscillations in long-horizon rewards, despite negligibly affecting the behavior cloning loss. We empirically disentangle the statistical and computational causes of these oscillations, and find them to stem from the chaotic propagation of minibatch SGD noise through unstable closed-loop dynamics.  While SGD noise is benign in the single-step action prediction objective, it results in catastrophic error accumulation over long horizons, an effect we term *gradient variance amplification* (GVA).  We demonstrate that many standard mitigation techniques do not alleviate GVA, but that taking an exponential moving average (EMA) of iterates is surprisingly effective at doing so.  Furthermore, we illustrate the generality of the phenomenon by showing both the existence of GVA and its amelioration by EMA in autoregressive language generation.  Finally, we provide theoretical vignettes both exhibiting the benefits of EMA in alleviating GVA and illustrating the extent to which classical convex models help in understanding the benefits of iterate averaging in deep learning.

Published as a conference paper at ICLR 2024
BUTTERFLY EFFECTS OF SGD NOISE: ERROR AMPLIFICATION IN BEHAVIOR CLONING AND AUTOREGRESSION
Adam Block
Cyril Zhang, Dylan Foster & Akshay Krishnamurthy
Department of Mathematics MIT
Microsoft Research New York, NY, 10012
Cambridge, MA. 02139 ablock@m t. edu
Jcyrilzhang, dylanfoster, akshaykr) cmccrooott.coom
Max Simchowitz MIT
Cambridge, MA, 02139 ms www.hhwwmiitt.mmmmmmmmmmm.mm edu
ABSTRACT
This work studies training instabilities of behavior cloning with deep neural net- works. We observe that minibatch SGD updates to the policy network during training result in sharp oscillations in long-horizon rewards, despite negligibly affecting the behavior cloning loss. We empirically disentangle the statistical and computational causes of these oscillations, and find them to stem from the chaotic propagation of minibatch SGD noise through unstable closed-loop dy- namics. While SGD noise benign in the single-step action prediction objective, it results in catastrophic error accumulation gradient variance amplification (GVA). We show that many standard mitigation techniques do not alleviate GVA, but find an exponential moving average (EMA) of iterates to be surprisingly effective at doing so. We illustrate the generality this phenomenon by showing the existence of GVA and its amelioration by EMA in both continuous control and autoregressive language generation. Finally, we provide theoretical vignettes that highlight the benefits of EMA in alleviating GVA and shed light on the extent to which classical convex models can help in understanding the benefits of iterate averaging deep learning.
long horizons, effect term
INTRODUCTION
Deep neural networks are increasingly used in machine learning tasks that contain feedback loops as a defining characteristic: outputs of language models depend on previously predicted tokens (Vaswani al., 2017), recommendation systems influence the users to whom they give suggestions (Krauth et al., 2020; Dean & Morgenstern, 2022), and robotic policies take actions in reactive control environments (Ross & Bagnell, 2010: Laskey et al., 2017). Because these tasks are so complex, is standard practice to optimize surrogate objectives, such as next-token prediction, that typically ignore feedback loops altogether (Pomerleau, 1988; Vaswani et al., 2017; Florence et al., 2022).
When training deep models by gradient updates on the surrogate objective, surrogate performance often improves more or less monotonically as training progresses. At the same time, successive iterates can exhibit wild variations in their performance impractical to evaluate the desired performance metric at multiple checkpoints, these oscillations imply that we have high risk of selecting and deploying poor policy. Thus, in order deter- mine best practices, we must first understand whether better training or better data will fix these instabilities. This leads to ask:
the task of interest. Because it is often

---

Published as a conference paper at ICLR 2024
Step 115000
Step
Walker2d-vd: 4-layer
Every-iterate evals (steps 40K-50K)
1.0
Clnne reundrd
Raward
Mean nomalized rewerd
0.3
Volldatio crror
0.0
0.0
EMA closer rewsrd
0.4
cloeet remerd cioner reward
0.3
policy
5.2
BC emarr
100K
200K
B00K
400K
500K
40K
53K
Training steps
Training steps
Figure 1: Typical reward instabilities over long-horizon (H 1000) rollouts of neural behavior cloners for the Walker2d-v4 MuJoCo locomotion task. Left: Rollout rewards (blue training curves) oscillate dramatically over the course of training (evaluated every 5000 iterations), while BC loss is stable. Center: Zoomed-in view of the highlighted region in (left). Large reward fluctua- tions are evident even between consecutive gradient iterates. Right: Exhaustive evaluation of small neighborhoods (in stochastic gradient directions) around iterates 115K and 120K, revealing fractal reward landscape - Jr (TO):: this jaggedness is invisible in the I-step behavior cloning objective (B((TT) ). Iterate averaging (EMA) drastically mitigates these effects (green training curves). Details are provided in Appendix C.I.I.
What causes instabilities in learning systems with feedback loops? To what ex- tent can they be mitigated by algorithmic interventions alone, without resorting to collecting additional data?
We explore this question in the context of behavior cloning (BC), technique for training policy to optimize a multi-step objective purely offline manner. This is achieved by introducing a surro- gate loss function fBC (behavior cloning loss) that measures the distance between actions generated by some expert policy Te. and those taken by the learner's s policy, and then minimizing bBC over an offline dataset of expert trajectories. BC is sufficiently broad as to capture important tasks ranging from robotics and autonomous driving (Pomerleau, 1988; Codevilla et al., 2018; Chi et al., 2023) to autoregressive language generation (Chang et al., 2023a), and is popular in practice due to its simplicity and purely offline nature.
Our starting point to observe that behavior cloning with deep neural networks exhibits training instabilities in which the multi-step objective (Jn), or rollout reward, of nearby checkpoints oscil- lates wildly during training, despite well-hehaved validation loss for {BC, even at the single iterate frequency; Figure I exhibits this phenomenon for sample training curve of behavior cloning pol- icy in the Walker2d-vy MuJoCo locomotion task (Towers et al., 2023; Todorov et al., 2012). This oscillatory behavior is clearly undesirable; we cannot differentiate between low- and high-quality iterates based on (validation loss for) ÊBC, and thus cannot reliably select a high-quality policy.
With regard to the final performance of BC, it is well understood that scarce or www-quality data can lead to statistical challenges and consequent performance degradation of imitator policies; unsurpris- ingly, better data often improves the quality of a learned policy. Unfortunately, existing approaches to obtaining better data require either interactive access to the demonstrating expert (Ross & Bag- nell, 2010; Laskey al., 2017) or additional side information (Pfrommer al., 2022: Block al., 2023a); these interventions may be costly or impossible in many applications. Thus, in this work we treat the data generating process as fixed and aim to investigate whether we can mitigate oscillations and improve the performance of BC solely through the application of better algorithmic choices.
1.1 CONTRIBUTIONS
In this paper, we aim to diagnose and ameliorate instabilities behavior cloning that arise from training on the surrogate cost alone in the purely offline setting. Our findings are as follows.
Diagnosis of rollout oscillations: gradient variance amplification. In Section 3, we conduct an extensive empirical study (278 distinct interventions) of BC in continuous control tasks and inves-
2

---

Published as a conference paper at ICLR 2024
tigate the effects that architecture, regularization, and optimization interventions have training instability. We identify the presence of training oscillations and attribute them to gradient variance amplification (GVA): the propagation of minibatch SGD noise through closed-loop dynamics, lead- ing to catastrophic error amplification resembling butterfly effects in chaotic systems. We ablate away much of the statistical difficulty, that the presence of oscillations suggests that GVA is an algorithmic rather than statistical pathology.
Mitigating GVA: stabilizers for unstable optimizers. In Section 4, we investigate mitigations for GVA. Because GVA is caused by variance in the stochastic gradients, it can be ameliorated with variance reduction. Indeed, we observe (Section 3.2) that i) aggressively decaying the learning rate, and ii) greatly increasing the batch size through gradient accumulation, both have positive effects on the stability of training. Unfortunately, both of these interventions come at a great increase in compute cost. As such, our most significant finding (Section 4.1) is that iterate averaging by taking an Exponential Moving Average (EMA) of the optimization trajectory (Polyak & Juditsky, 1992; Ruppert, 1988), stabilizes training and mitigates GVA across wide range of architectures and tasks, with essentially no downsides. While iterate averaging is popular in many deep learning research communities, this paper exposes iterate averaging as essential design consideration when training any deep model in the presence of feedback loops.
A preliminary study of GVA in language generation. In Section 4.2, broaden our focus by considering autoregressive sequence models. Our findings suggest that unstable optimizers, when stabilized with iterate averaging to mitigate GVA, do not need full learning rate decay, entailing potential computational and statistical benefits for training language models. For this reason, we suggest that EMA and related filters be designated stabilizers in their own right and incorporated into deep learning pipelines in the same vein as modern optimizers and schedulers.
The applicability of convex theory. In Section 4.3, we complement our empirical results with the- oretical vignettes. While the benefits of large learning rates cannot be explained in a convex setting, we demonstrate chat-conditional on using theoretically suboptimal learning rates stochasicc con- vex optimization provides useful intuition for the causes and mitigations of GVA in deep learning. With our empirical results, these findings add line of work on surprising near-convex behavior in deep learning (Sandler et al., 2023; Frankle al., 2020; Fang et al., 2022; Schaul et al., 2013).
1.2 RELATED WORK
Understanding and mitigating the effects of error amplification in behavior cloning has been the subject of much empirical work (Ross & Bagnell, 2010; Laskey et al., 2017), but most approaches use potentially impractical online query access to the expert policy; instead, we focus on a purely offline setting.
Complicated value function landscapes and their effect on training have been investigated in the context of planning in RL, with Dong al. (2020) investigating natural examples of fractal reward functions, Wang et al. (2021) examining the instabilities arising from poor representations, and Emmons et al. (2021); Chang et al. (2023a) observing the fact that fBC poor proxy for JH. To the best of our knowledge, there has not been a systematic study of training instability in the sense of rollout reward oscillation of nearby checkpoints.
In the context of stochastic optimization and optimization for deep learning, many previous works have attempted to reduce variance theory (Polyak & Juditsky, 1992; Ruppert, 1988) and practice (Izmailov et al., 2018; Busbridge et al., 2023; Kaddour, 2022; Kaddour et al., 2023). Of particular note is Sandler et al. (2023), which demonstrates (empirically and in a toy theoretical setting) a form of equivalence between learning rate decay and iterate averaging. Our focus is not on variance reduction per se, but rather on the propagation of vvariance through unstable feedback loops. We expand the relationship between our work and Sandler al. (2023) and discuss other related work in Appendix B.
PRELIMINARIES
MDP formalism. We let M= (S,A, Hv)) denote a finite- horizon Markov decision process (MDP), where S is an abstract state space, A is abstract action space, P S x A (SS) is
3

---

Published as a conference paper at ICLR 2024
Markov transition operator. We denote by 8 A + [0, 1] reward function and H N is the length of the horizon. Because we focus on continuous control tasks, we follow the notational conventions of control theory, denoting states by and actions by u. We let v/ E A(S) denote the initial distribution such that a trajectory from M consists of X1 ~ V and X4+11 ~ P(- Xx: u,) for all h.
The learner has access to a class of policies π. : S θ - ((A), where θ is the parameter space and Te : S + A(A) is the policy induced by parameter 0 E O. Given policy ma, we denote its expected cumulative reward by (TT) = - r(x², un)) where un ~ (-(- and the ex- pectation is with respect to both the transition dynamics of M and the possible stochasticity of the policy. Our experiments focus on MDPs whose transition operators P are deterministic, i.e., there exists a function f S A - S such that x/++1 f(x,, Uy) for all h. In this case the only stochas- ticity of the system comes from the sampling of the initial state X ~ v (and possibly the policy).
Imitation learning and behavior cloning. In imitation learning, we are given an offline data set (0)
of N trajectories Do = {(x, U JISHCH | I i S N) generated by an expert policy xe- interacting with the MDP MM. In this work, always consider deterministic policies, i.c. where for all x, (x) has support on a single action; in particular this holds for the expert тØ- The goal of the learner is to produce policy "g that maximizes the expected cumulative reward Ju (*r) over an episode. We focus on the popular behavior cloning (BC) framework, where we fix a loss function BBC AA A + R that measures the distance from the actions produced by sæ. and learn To by attempting to minimize the empirical risk of fBC over Dott: we abuse notation by denoting (T)( := ED... (x), u)] The basic premise behind behavior cloning that LBC should be chosen such that if (B)(TT) < 1 then JH (TT)) = Ju (To- ); that is, imitation of the expert is line with common practice in BC (Janner et al., 2021; Shafiullah et al., 2022; Chi et al., 2023), the imitator policies in our experiments augment the state with the previous action, i.c., Te : A + A. which can be integrated into the previous formalism by expanding the state space. For the special case of the first state X1, we always let ug O.
a surrogate for large cumulative reward.
Throughout the paper. we denote vectors by bold lower case letters and matrices by bold upper case letters. We reserve 0 for parameter of our policy and Ju for the cumulative reward over a trajectory, omitting H when it is clear from context. For conciseness, we often refer to Jır as the reward; the per-step reward function makes no appearance in the rest of the paper. Given set U, we let A (2) denote the class of probability distributions on U.
Notation.
3 DIAGNOSIS OF ROLLOUT OSCILLATIONS: GRADIENT VARIANCE AMPLIFICATIIO
3.1 INSTABILITIES IN BEHAVIOR CLONING OF MUJoCo TASKS
Experimental setup. We investigate instabilities in behavior cloning for the [Walerrdd, Hopper, HalfCheetah. Humanoid, J-v4 environments from the OpenAl Gymnasium (Towers et al., 2023), all rendered in MuloCo (Todorov et al., 2012). We focus on Wa Wal kellke-v for the discussion that follows, and defer detailed discussion of further environments (which exhibit similar behavior) to Appendix C. Our expert is a multilayer perceptron (MLP) trained with Soft Ac- tor Critic (SAC) (Haarnoja al., 2018) for 3M steps with able base leleiness (Raffin al., 2021), with out-of-the-box hyperparameters. The default imitator is a 4 layer MLP; details are in Appendix C. We examine several widths and depths, as well as Transformer (Vaswani et al., 2017) imitators.
Our first suite of experiments aims to isolate instability from statistical difficulties. We set up the experiments to make the behavior cloning problem as easy as possible. First, we focus on the "large-data" regime N
1000, in which overfitting with respect to the BC loss (BC(Tg) is not problem (see Figure I), and thus poor rollout performance for JH (T&) cannot be blamed on
In particular, we denote states by x and actions by in order emphasize that, in our experiments, they are vectors in Fuclidean space.
"By default, the Stable- Bable-Baselines3 SAC agent is stochastic, but we enforce determinism by selecting the mean action of the resulting policy. This results in negligible degradations to the rewards; see Figure 5.
4

---

Published as a conference paper at ICLR 2024
insufficient data; this removes a typical source of statistical difficulty faced in applying behavior cloning to domains such as robotics (Chi et al., 2023; Pfrommer et al., 2022; Ross & Bagnell, 2010; Laskey et al., 2017). Beyond focusing on the large-data regime, (i) consider only deterministic dynamics and deterministic experts, and (ii) we include within our default model the same class of MLPss that parameterize the expert policies, ensuring that expressivity is not issue. As such, have placed ourselves in perhaps the easiest possible setting for behavior cloning.
In Figure 1 (Left), we compare the evolution of the BC loss LBC (Tg) (on a validation set) and reward Jn (TT) for imitator policies in the Walker2d MuJoCo locomotion task. In this figure, we observe extreme oscillatory behavior in (TØ). juxtaposed with smoothly decaying (Ta). In Figure 1 (Middle), we zoom in on the training trajectory between iterates 40K and 50K and observe that the same instability persists even at the every-iterate level. Toward identifying what causes these instabilities, Figure I (Right) displays an experiment in which we independently sample two stochastic gradients of the training loss at a fixed checkpoint with good rollout reward. Policy weights are then perturbed by small steps in each of the two directions, and we evaluate the resulting reward J (Ta) over 20 rollouts, along with the BC loss (TT) on a held-out validation set. We see that nearby models vary erratically in terms of rollout performance, but vary smoothly in validation BC loss. These findings are reproduced consistently across the other environments and architectures in Appendix C; thus, we conclude:
(R1) The reward landscape is highly sensitive to small changes in policy parameters: small perturbations in model weights induce butterfly effects in the reward JH (Tg). In contrast, in the same regions, the BC loss landscape 0 + (BC( is
(nearly linear locally).
3.2 INSTABILITY IS CAUSED BY GRADIENT VARIANCE AMPLIFICATION
We now present compelling evidence that variance in stochastic gradients during training is sponsible for training instability, because gradient variance amplified through the sensitivity of the rollout rewards to fluctuations in network parameters. In Figure 2, we visualize evolu- tion of both (Tg) and Ju (#T) over training for a variety of potential algorithmic interventions. We find that neither changing the model architecture and scale (1st row) nor standard regulariza- tion techniques (2nd row) ameliorate the training instabilities observed. We do see, however, that aggressively decaying the learning rate and increasing the batch size (3rdl row) significantly reduces oscillations (at least when measuring mean rewards), the expense of substantially slowing down training. Thus, we conclude that fluctuations from stochasticity in the gradients are to blame for oscillations in rollout rewards, and term this phenomenon gradient variance amplification (GVA). To summarize:
(R2) GVA arises from algorithmic suboptimality rather than an information-theorrti limit. Even with "infinite" training data (i.e., fresh trajectories with i.i.d. initial conditions at each training step), rollout oscillations persist.
(R3) Training oscillations are not mitigated by many standard approaches to regularization, including architectural interventions and increased regularization. On the other hand, oscilla- tions are ameliorated by variance reduction techniques, such as large batch sizes, learning rate decay, and iterate averaging.
Appendix C shows that (R2) and (R3) remain true across environments and model architectures. In addition, we find that training instability is not the result of inadequate network architecture; we observe oscillations across model scales, and for both MLP and Transformer architectures.
While Figure 2 shows that it is possible to quell GVA using small learning rates or large batch sizes, this may not always be practical, both interventions can incur steep computational costs.? Even worse, the success of continuous optimization in deep learning depends on non-convex fea- ture learning mechanisms (Chizat et al., 2019), and too small learning rate or too large batch size can have deleterious effects on generalization.' 4 Thus, it is vital to seek interventions that are holistically compatible with existing deep learning pipelines. Among these, Figure 2 highlights that
As another unsatisfactory compromise, we also find that shallower models are less susceptible to GVA. We refer some theoretical and empirical accounts in Appendix B.
5

---

Published as a conference paper at ICLR 2024
Walker2d-v4: reward oscillations in behavior cloning
4L MLP (default)
2L MLP
6L MLP
2L Transformer
4L Transtorer
0.2
reward Individual rewards
0.1
0.0
Infinite
Dropout
Large decay
L2 loss
Augment
WM
0.2
8C error
0.1
16x batch size
Early decay
Large momentum
EMA
EMA (Transformer)
Ver. reduction
0.5
400k
400k
200k
400k
600k
200k
400k
600k
200k
Training steps
Figure 2: Highlights from a large suite of experiments, suggesting algorithmic (rather than sta- tistical) origin of reward oscillations. plots use the 4-layer MI.P architecture unless otherwise specified. Blue curves show mean rewards over 20 initial conditions, while teal dots show disag- gregated per-episode rewards (such that each point represents the rollout reward a fixed initial condition of the policy at the current iterate). These oscillations persist across dataset sizes, archi- tectures, model scales, and choices regularizers, and diminish toward the end training as the learning rate decays to O. They most strongly mitigated by variance reduction strategies. Here, we opt for direct visualizations, providing a qualitative demonstration of GVA and its mitigations. We accompany these with quantitative comparisons in Appendix C.1.2.
a large momentum coefficient mildly helpful, but taking an exponential moving average (EMA) of iterates (Polyak & Juditsky, 1992) is extremely effective. This motivates us to take closer look at the latter in Section 4 through another suite experiments.
3.3 UNDERSTANDING GVA: MISMATCH BETWEEN BC LOSS AND ROLLOUT REWARD
The disparity between behavior cloning loss LBC (TC(TT) and rollout reward JH (Tg) has long been ap- preciated in the imitation learning literature, and is understood to be caused by error amplification, the process by which mildly erroneous predictions, when fed repeatedly through feedhack loops, result in highly suboptimal performance (Chen & Hazan, 2021; Wang et al., 2020a). More pre- cisely, for given BC and JH well policy т&- and O, we define the error amplification constant at scale 8 to be the maximal value (with respect fo 0) for JH (Te-) Ju (me such that (B)(TT) BBC (
8. The following proposition provides a simple theoretical illustration for how small fluctuations in BC loss can be drastically amplified by feedback between imperfectly- imitated policies and system dynamics.
Proposition 3.1 (Example of exponential error amplification). Let Ba denote the set of 8-Lipschitz functions A : S > A with ((0) = O. For any > O, there exists deterministic MDP with horizon H and expert policy Tg- such that the dynamics are Lipschitz in both state and action and n is Lipschitz in the state, and such that
sup (mo( -HI + +)} 2(H) AEB,
yet supaes; ( + A) < O (H 8²), where LBC is the &g loss. Thus, the error amplification constant is exponential in the time horizon.
Working model for GVA. Proposition 3.1 shows that even when LRc is uniformly small in neigh- borhood around π the rollout loss can be exponentially large the same neighborhood. At the same time, there are good subsets of parameter space that do not experience this worst-case error amplification our construction. We therefore hypothesize that, when stochastic optimization con- verges to small neighhorhood around zero-BC error models, it bounces between low-BC emror

---

Published as a conference paper at ICLR 2024
Walker2d-v4: stabilizing effects of EMA (iterate averaging)
Architectures
Abletions
Fallure modes
Additional tasks
Low-GVA tasks
4L MLP (default)
No decay
Infrequent EMA
Hopper
Ant
0.2
EMA BC error
Reward EMA reward
0.5
0.1
4L Transfomerr
0.1x training data
No burn-in
HalfCheetah
Humanoid
0.2
BC error
0.5
200k
400k
600k
200k
400K
600k
200k
400K
600k
200k
400k
6OOK
0
200k 400k
600K
Training steps
Figure 3: Iterate averaging significantly mitigates GVA-induced reward oscillations, without need- ing to change the learning rate schedule or batch size. These improvements hold across ar- chitectures, dataset sizes, and some tasks. Column 2, bottom: Algorithmic instabilities are more pronounced at smaller sample sizes; thus, stabilization can lead to improved sample efficiency. Column 3: We recommend updating the EMA at every iterate, with an initial burn-in phase, and with tuned 7(()) = t- decay, to avoid divergence or slower progress. Columns 4-5: We verify that the benefits EMA are not exclusive to the Walker2d- task; for some other tasks (including the nigher-dimensional Humanoid-Hs oscillations are more benign.
models that experience large error amplification, and those that do not. To recapitulate: GVA the phenomenon in which gradient stochasticity leads to optimization Irajectories repeatedly visiting gions of parameter space with catastrophic error amplification. Because our MuloCo environments involve nonlinear contact dynamics (while the example in Proposition 3.1 is linear), oscillations in Figure are even more chaotic than this example may suggest. We elaborate on this point further by studying the advantages of EMA on discontinuous "cliff loss" problem in Section 4.3.
MITIGATING GVA: STABILIZERS FOR UNSTABLE OPTIMIZERS
In Section 3.2, we isolated GVA the primary cause of observed instabilities in BC (cf. Fig. 1) and identified iterate averaging with EMA (Polyak & Juditsky, 1992) as a promising remedy. In this section, we conduct an in-depth investigation of EMA as a mitigation. We start continuous control (Section 4.1), and find EMA works almost unreasonably well at reducing GVA in the experimental testhed described the prequel. Next, moving beyond continuous control (Section 4.2), we observe analogous effects in autoregressive language generation. In both settings, we find iterate averaging works so well as to eliminate the need for full learning rate decay; this leads us to recommend a conceptual reframing of EMA as stabilizer for training neural networks, akin to (and interacting with) conventional optimizers and schedulers. We conclude (Section 4.3) by exploring the extent to which intuition benefits of iterate averaging from the theory of stochastic convex optimization applies in our empirical settings.
4.1 THE OUTSIZED BENEFIT OF ITERATE AVERAGING
We recall the definition of the EMA method for iterate averaging (Polyak & Juditsky, 1992). Given an optimization trajectory (0()) Jost C Rdd and a sequence (7))ge C [0, 1). the EMA iterates (0%)
a
0 (00 and 0((+1) (1 (1-))) 0 + ((++1).
(4.1)
Many prior works have detailed the benefits of iterate averaging in stochastic convex optimization and beyond (see Appendix B). Here, we investigate its effect on GVA. We begin by considering the
CCommon heuristics include updating the EMA only after an initial "burn-in", and annealing 7 with a polynomial decay: (0)) max(t " 7min). is also customary to use (00) to denote 1 ((())

---

Published as a conference paper at ICLR 2024
same MuJoCo framework in Section 3. In Figure 3, we produce similar plots to those in Section 3, but this time juxtapose the vanilla trained models with an EMA of their iterates (further results and details are deferred to Appendix C). We observe the following:
(R4) EMA iterate averaging strongly mitigates rollout oscillations. In every setup we consider, across a variety of architectures and environments, EMA significantly reduces the oscillations in rollout reward; in no instance does it hurt performance.
We provide quantitative comparisons for wide range of interventions in Figures 8 to 11.
4.2
AUTOREGRESSIVE SEQUENCE MODELS AND LANGUAGE GENERATION
Every-iterate
EMA loss itertt loss (more nuns)
Valcation loos (lgg perleatt)
TinyStories LM: 12-layer Transformer
UK loss 2MA UK lost
LM
BMG no HM
Zepacts decyy
Reparts, deoy
eppacts, deay
EMA worse
BAA Untter
Repacts, conet
eppceee contt
Repects cont
0.50
0.55
245 5930 rat D806 12000 1160 2744 Tranning steps
Training steps
0.80
nd6
Validation less
nis
3.90
L.J One day, Tim and Spot were playing near the pit when mhey saw ball. They both wanted to play with...
15027-15035: the
they dinn't want shere. Tim said,
want the ball!" They started to light the can both play with the ball If you share." Tim Spot stopped fighting and [--]
ball!" barked,
bell. Tim's mom saw them fighting and said, "Stop fighting!
15038-15037: it. Tim said, "Let's share the ball, Spot! Spot barked happily, and they both played with the ball together. As they played, a litll giri named Sue came by. She saw Tim and Spot sharing the ball and wanted to join them. Tim said, "Yos, you can play with usi" So, Tim, [-.]
15038-15039; ...t Tim ssid, "I want the ball!" Spot barked and wagged his tail. They both ran the ball and started to play. As they played, the and Spot
bell went into
each other, knowing what do. Then, a friendly frog jumped out
pit the [-]]
15040-15049: the ball, they didn't want to share. Tim said, ball. As they tought, the ball rolied into the př. Tim and Spot stopped lighting and looked at the pit. They were sad because they couldint [...]
Spot barked,
want the ball!" They started to light over the
Figure 4: GVA in natural language generation, with 270M-paramett Transformer models trained on TinyStories. (Top row) Left: Validation loss curves with and without EMA. Center: Zooming in on (left), evaluations at every update demonstrate small per-iterate loss fluctuations, which are even smaller if EMA is applied; note that the green "lines" are also scatter plots. Right: Training paths in (model loss, EMA loss) space. EMA enables training without learning rate decay; this mitigates overfitting, resulting in the perrst-perplexity model. (Bottom) Examples of autoregressively gener- ated text (with argmax decoding), where nearby training iterates can bifurcate. See Appendix C.2 for full results, including quantitative evaluations of these "butterfly effects" " in generation.
We posit that GVA is generic phenomenon that can manifest in disparate settings: whenever model's 's predictions are applied within a (marginally stable or unstable) feedback loop. the closed- loop dynamics amplify small fluctuations in deleterious manner. A natural and timely setting with this tructure-which complements continuous control-is autoregressive language modeling. Here, network's parameters O are optimized on 1-step prediction loss, which takes the role of (Te( The network Ty is then used to generate a sequence of symbols WI-H by iteratively rolling in NLP, code synthesis, and
out Te Wi-k + Wh+1. Such models have been beyond. Motivated by the similarity of this pipeline to behavior cloning," " we perform a smaller set of analogous experiments on language generation. Our findings here parallel our findings for continuous control, and show (i) the presence of GVA, and (ii) substantial benefits of iterate aver- aging- In more detail, we train 270M-parameter 12-layer Transformer models on the Ttoyis dataset (Eldan & Li, 2023), which serves as inexpensive surrogate for full-scale pretraining pipeline. Highlights are shown in Fig. 4, while Appendix C.2 provides full documentation, includ- ing larger-scale training runs with a non-synthetic corpus (Wikipedia). We summarize our findings below:
MMan works have investigated GPT-style pretraining through the lens of offline IL (Chang et al., 2023a). There are many degrees of freedom in evaluating performance; thus, we do not commit to canonical notion of reward and measure GVA-induced oscillations via disagreements in long-horizon rollouts.
8

---

Published as a conference paper at ICLR 2024
(R5) Autoregressive LMs exhibit significant rollout oscillations throughout training. EMA sta- bilizes the trajectory, accelerates training, and improves generalization, complementing (and potentially obviating) standard practices in learning rate annealing.
4.3 To WHAT EXTENT DOES CONVEX THEORY EXPLAIN THE BENEFITS OF EMA?
We close by providing mathematical intuition as to why iterate averaging with EMA can reduce the oscillations caused by GVA. As discussed in Section 3.3, oscillations can occur when there is disparity between the BC loss (B)(TT) on which we train and the rollout reward function J J(Tr)) on which we evaluate. To study this phenomenon, we a consider simple, horizon-one behavior cloning problem with single action determined by the model parameter 0. We take the training loss to be quadratic (0)
110 ull? and the rollout reward J-) to be
((0) { 110 -- C,
ull
|0 - μب S
(4.2)
otherwise
where C' 6² are constants. Here the training loss is convex, but rollout reward is not; the latter exhibits "cliff," dropping sharply from - to -C once |00 - slll The pair (fpc, J) may be thought of as a discontinuous, horizon-one analogue of the example in Proposition 3.1, illustrating the contrast between extreme sensitivity of reward and insensitivity of the loss to the parameter of interest. The reward function encapsulates discontinuities arising in control tasks from, e.g. contract forces. In the MuJoCo walker, "cliff"-type behavior may come from an expert policy close to overbalancing the agent, with the learner's ' policy falling down if the parameter is "over the cliff."
We analyze SGD iterates ((++1) = g ()())) - 4 + w), where 17 > is a constant step size and w I). This corresponds SGD noisy version of the BC loss given by (0) O,
w/|²], which satisfies E [0)] = BC) + constant. We show that applying EMA to the resulting iterates achieves substantially higher rollout reward than vanilla SGD.
Proposition 4.1 (Informal version of Proposition D.6). Consider the setting in Eq. (4.2) for pa- rameters C >> 2 > 0 in dimension one, and let O(T) denote the SGD iterate with learning rate n > as described above. Let 0 denote the EMA iterate (4.1) with fixed parameter " E7 S n satisfying >> 1/T. Then, Bl(Bc(0(T)) scales as ((r), while E[[nc( (()))(()))) scales as O(7) S m. In particular, when 1j > C16, and 7 log(C/n) S e2c, for absolute constants C1, 61,02 > O, we find that
B[J(0*) J (O(T))) ç. but - O(7).
This proposition holds, which shows that the rollout performance for EMA can be arbitrarily small relative to that of SGD, holds even in the regime where SGD initialized at ( = la (so that both ((T) and 0.²²) are unbiased estimates of μ), and thus highlights that EMA can reduce the variance that arises from accumulation of SGD noise.? Notice that Proposition 4.1 requires 2 y >> 1/T, which is above the optimal step size of π) = 1/t." Indeed, in Appendix D we show that EMA, with the parameters we find empirically successful, only benefits optimization above these degeese theoretically optimal learning rate schedules. Thus, conclude that convex theory reveals the periance-reducíin benefit of either learning rate decay or EMA, but does not suggest which one is better. We defer further theoretical results Appendix and present an empirical study of a system motivated by the cliff loss in Appendix C.5; in particular, our analysis provides a simple example where GVA provably occurs, both theoretically and empirically.
The above example reveals the difference between the statistical and algorithmic difficulties of BC: with enough data, the empirical risk minimizer (sample mean) ① of BC loss exhibits (0) 1/T < 6, which ensures JH is small; on the other hand, with minibatch SGD and too large a learning rate, there is noise floor on how close the non-EMA'd iterate i will be to 0*, ensuring that JH large.
7 We compare to similar findings (Sandler et al., 2023) in Appendix B. 'Note that the 7/t step size schedule gives the sample mean, which is the maximum likelihood estimator for objective.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We are extremely grateful to Jonathan Chang for illuminating discussions at multiple stages of this project. We thank Daniel Pfrommer for independent verifications of GVA: that it is benign for the 2D quadcopter (a canonical smooth nonlinear control environment), and less so when training diffusion models imitate expert behavior on the PushT environment. We also thank Mojan Javaheripi, Piero Kauffmann, and Yin Tat Lee for helpful discussions about learning rate schedules, and Sadhika Malladi for helpful NLP references. greatfully acknowledges support from the National Science Foundation Graduate Research Fellowship under Grant No. 1122374.
REFERENCES
Emmanuel Abbe and Colin Sandon. Poly-time universality and limitations of deep learning. arXiv preprint arXiv:2001.0299 2020.
Emmanuel Abbe, Pritish Kamath, Eran Malach, Colin Sandon, and Nathan Srebro. On the power of differentiable learning versus pac and learning, Advances in Neural Information Processing Systems, 34:24340-24351, 2021.
Naman Agarwal, Brian Bullins, Elad Hazan, Sham Kakade, and Karan Singh. Online control with adversarial disturbances. International Conference on Machine Learning, pp. 111-119. PMLR, 2019.
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint arXiv:2002.11802 2020.
Maksym Andriushchenko, Aditya Vardhan Varre, Loucas Pillaud- Vivien, and Nicolas Flammarion. SGD with large step sizes learns sparse features. In International Conference Machine Learn- ing. pp. 903-925. PMLR, 2023.
Marcin Andrychowicz, Anton Raichuk, Piotr Stańezyk, Manu Orsini, Sertan Girgin, Raphaël Marinier, Leonard Hussenot, Matthieu Geist, Olivier Pietquin, Marcin Michalski, et al. What matters for on-policy deep actor-critic methods? large-scale study. In International Conference on Learning Representations, 2020.
David Angeli. Lyapunov approach to incremental stability properties. IEEE Transactions on Automatic Control, 47(3):410-43 2002.
Kushal Arora, Layla El Asri, Hareesh Bahuleyan, and Jackie Chi Kit Cheung. Why exposure bias matters: An imitation learning perspective of error accumulation in language generation. arXiv preprint XrXiv:2204.00 00000...117,1... 2022.
Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang, Hid- den progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 01221750-21766 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pp. 2397-2430. PMLR, 2023.
Adam Block and Max Simehowitz. Efficient and near-optimal smoothed online learning for gen- eralized linear functions. Advances in Neural Information Processing Systems, 35:7477-7489, 2022.
Adam Block, Daniel Pfrommer, and Max Simchowitz. Imitating complex trajectories: Bridging low-level stability and high-level behavior. arXiv preprint 1rXiv:2307.14619, 2023a.
Adam Block, Max Simchowitz, and Alexander Rakhlin. Oracle-efficient smoothed online learning for piecewise continuous decision making. In The Thirty Sixth Annual Conference on Learning Theory, pp. 1618-1663 PMLR, 2023b.
10

---

Published as a conference paper at ICLR 2024
Adam Block, Max Simchowitz, and Russ Tedrake. Smoothed online learning for prediction in piecewise affine systems. arXiv preprint arXiv:2301.11187 2023c.
Léon Bottou.
Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT 19th International Conference on Computational Statistics, pp. 177-186. Springer, 2010.
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learn- ing Research, 2:499-526, 2002.
Mark Braverman, Xinyi Chen, Sham Kakade, Karthik Narasimhan, Cyril Zhang, and Yi Zhang. Calibration, entropy rates, and memory in language models. In International Conference Machine Learning, pp. 1089-1099, PMLR, 2020.
Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Eeshan Gunesh Dhekane, Xavier Suau, and Russ Webb. How to scale your EMA. arXiv preprint arXiv:2307. / 13813, 2023.
Anthony Carbery and James Wright. Distributional and Lq norm inequalities for polynomials over convex bodies in Mathematical Research Letters, 8(3):233-244 2001.
Jonathan D Chang, Kiante Brantley, Rajkumar Ramamurthy, Dipendra Misra, and Wen Sun. Learn- ing to generate better than your LLM. arXiv preprint arXiv:2306. 1816, 2023a.
Jonathan D Chang, Qinqing Zhen, Brandon Amos, Wen Sun, and Mikael Henaff. A large scale study of deep imitation learning on the arcade learning environment. Personal Communication, 2023b.
Xinyi Chen and Elad Hazan. Black-box control for linear dynamical systems. In Conference Learning Theory, pp. 1 1114-1143. PMLR, 2021.
Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and Langevin processes. In International Conference on Machine Learning, pp. 1810-1819. PMLR, 2020.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy. End- to-end driving via conditional imitation learning. In 2018 IEEE international Conference Robotics and Automation (ICRA), pp. 4693-4700, IEEE, 2018.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, pp.
PMLR, 2022.
Sarah Dean and Jamie Morgenstern. Preference dynamics under personalized recommendations. In Proceedings of the 23rd ACM Conference on Economics and Computation, pp. 795-816, 2022.
Sarah Dean, Horia Mania, Nikolai Matni, Benjamin Recht, and Stephen Tu. Regret bounds for robust adaptive control of the linear quadratic regulator. Advances in Neural Information Processing Systems, 31, 2018.
Aaron Defazio and Léon Bottou. On the ineffectiveness of variance reduced optimization for deep learning. Advances in Neural Information Processing Systems, 32, 2019.
Kefan Dong, Yuping Luo, Tianhe Yu, Chelsea Finn, and Tengyu Ma. On the expressivity of neural networks for deep reinforcement learning. In International conference on machine learning, pp- 2627-2637. PMLR, 2020.
Simon S Du, Sham M Kakade, Jason D Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in RL. International Conference Machine Learning, 2021.

---

Published as a conference paper at ICLR 2024
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pp. 5793-5831. PMLR, 2022.
Ronen Eldan and Yuanzhi Li. TinyStories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.07759, 2023.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline rl via supervised learning? arXiv preprint arXiv:21 2021.
Logan Engstrom, Andrew Iyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep RL: A case study on PPO and TRPO. In International Conference on Learning Representations, 2019.
Cong Fang, Yihong Gu, Weizhong Zhang, and Tong Zhang. Convex formulation of overparameter ized deep neural networks. IEEE Transactions on Information Theory, 68(8):5340-5352, 2022.
Maryam Fazel, Rong Ge, Sham Kakade, and Mehran Mesbahi. Global convergence of policy gradi- ent methods for the linear quadratic regulator. In International Conference on Machine Learning, pp. 1467-1476. PMLR, 2018
Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning, In Conference on Robot Learning, pp. 158-168. PMLR, 2022.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint ar 13487, 2021a.
Dylan Foster, Akshay Krishnamurthy, David Simchi-Levi, and Yunzong Xu. Offline rein- forcement learning: Fundamental barriers for value function approximation. arXiv preprint arXiv:2111.10919 2021b.
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In International Conference on Machine Learning, pp- 3259-3269. PMLR, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp.
PMLR, 2019.
Udaya Ghai, Holden Lee, Karan Singh, Cyril Zhang, and Yi Zhang. No-regret prediction in marginally stable systems. In Conference on Learning Theory, pp. 1714-1757. PMLR, 2020.
Priya Goyal, Piotr Dollár, Ross Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 1oXiv:/706.000 706.02677, 2017.
Saan-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271-212844 2020.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International conference on machine learning, pp. 1225-1234. PMLR, 2016.
12

---

Published as a conference paper at ICLR 2024
Nicholas JA Harvey, Christopher Liaw, Yaniv Plan, and Sikander Randhawa. Tight analyses for non- smooth stochastic gradient descent. In Conference on Learning Theory, pp. 1579-1613. PMLR, 2019.
Aaron Havens and Bin Hu. On imitation learning of linear control policies: Enforcing stability and robustness constraints via Imi conditions. In 2021 American Control Conference (ACC), pp- 882-887. IEEE, 2021.
Elad Hazan and Karan Singh. arXiv:2211.09619 2022.
Introduction to online nonstochastic control. arXiv preprint
Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: Beyond sharp and flat local minima. Advances in Neural Information Processing Systems, 32, 2019.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Jorgen Hoffmann-lorgensen, Lawrence A Shepp, and Richard M Dudley, On the lower tail of gaussian seminorms. The Annals of Probubility, pp. 319-342, 1979.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint arXiv: 1904.0975, 2019.
Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging weights leads to wider optima and better generalization. arXiv preprint 1rXiv:1803.05407, 2018.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. Advances in Neural Information Processing Systems, 34:1273-1286, 2021.
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire, Con- textual decision processes with low Bellman rank are PAC-learnable. In International Conference on Machine Learning, pp. 1704-1713, 2017.
Chi Jin, Rong Ge, Praneeth Netrapalli, Sham M Kakade, and Michael Jordan. How to escape saddle points efficiently. In International conference on machine learning, pp. 1724-1732. PMLR, 2017.
Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of RL problems, and sample-efficient algorithms. Neural Information Processing Systems, 2021.
Jean Kaddour. Stop wasting my time! Saving days of ImageNet and BERT training with latest weight averaging. arXiv preprint arXiv:2209.1449 14981. 2022.
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt Kusner. No train gain: Revisiting efficient training algorithms for Transformer-bart based language models. arXiv preprint arXiv:2307.064400 2023.
Sham Kakade, Akshay Krishnamurthy, Kendall Lowrey, Motoya Ohnishi, and Wen Sun. Infor- mation theoretic regret bounds for online nonlinear control. Advances Neural Information Processing Systems, 33:15312-15325, 2020.
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.
Liyiming Ke, Jingqiang Wang, Tapomayukh Bhattacharjee, Byron Boots, and Siddhartha Srinivasa. Grasping with chopsticks: Combating covariate shift in model-free imitation learning for fine manipulation. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pp- 6185-6191. IEEF, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 6980, 2014.
13

---

Published as a conference paper at ICLR 2024
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska- Barkinska, et al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114(13):3521-332 2017.
Karl Krauth, Sarah Dean, Alex Zhao, Wenshuo Guo, Mihaela Curmei, Benjamin Recht, and Michael Jordan. Do offline metries predict online performance in recommender systems? arXiv preprint arXhv:2011.07931, 2020.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- off- policy q learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.04728 2020.
Simon Jacostte-ullen, Mark Schmidt, and Francis Bach. A simpler approach to obtaining (1/t) convergence rate for the projected stochastic subgradient method. HAL, 2012, 2012.
Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. In Conference on Robot Learning, pp. 143-156. PMLR, 2017.
Jean-François Le Gall. Brownian motion, martingales, and stochastic calculus. Springer, 2016.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. Advances neural information processing systems, 31, 2018.
Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms Mathematical foundations. The Journal of Machine Learning Research, 20 (1):1474-1520, 2019a.
Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activation.
Advances in neural information processing systems, 30, 2017.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in Neural Information Processing Systems, 32, 2019b.
Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora. Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate. Advances in Neural Information Processing Systems, 33:14544-14555, 2020.
Ernest Lindelöf. Sur l'application de la méthode des approximations successives aux équations différentielles ordinaires du premier ordre. Comptes rendus hebdomadaires des séances de l'Académie des sciences, 116(3):454 457, 1894.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Exposing atten- tion glitches with flip-flop language modeling. arXiv preprint arXiv:2306.00946 2023a.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. 2023b.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 171 2017.
Marlos C. Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora. On the sdes and scaling rules for adaptive gradient algorithms. Advances in Neural Information Processing Systems, 35: 7697-7711, 2022.
14

---

Published as a conference paper at ICLR 2024
Stephan Mandt, Matthew D Hoffmann, and David M Blei. Stochastic gradient descent approximate bayesian inference. arXiv preprint vrXiv:170..042200 2017.
Yurii Nesterov et al. Lectures convex optimization, volume 137. Springer, 2018.
Charles Packer, Katelyn Gao, Jemej Kos, Philipp Krähenbühl. Vladlen Koltun, and Dawn Song. 2018.
Assessing generalization in deep reinforcement learning. arXiv preprint
Paavo Parmas, Carl Edward Rasmussen, Jan Peters, and Kenji Doya. Pipps: Flexible model-based policy search robust the curse of chaos. In International Conference Machine Learning, pp- 4065-4074. PMLR, 2018.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy. Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-PPeformance Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, d' Alché Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019.
Juan Perdomo, Jack Umenberger, and Max Simchowitz. Stabilizing dynamical systems via policy gradient methods. Advances in Neural Information Processing Systems, 34:29274-29286, 2021.
Daniel Pfrommer, Thomas Zhang, Stephen Tu, and Nikolai Matni. Tasil: Taylor series imitation learning. Advances Neural Information Processing Systems, 35:20162-20174, 2022.
Daniel Pfrommer, Max Simchowitz, Tyler Westenbrock, Nikolai Matni, and Stephen Tu. The power of learned locally linear models for nonlinear policy optimization. arXiv preprint arXiv:2305.09619, 2023.
Boris Polyak and Anatoli Juditsky. Acceleration of stochastic approximation by averaging. SIAM journal control and optinization, 30(4):838 885, 1992.
Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Advances in Neural Information Processing Systems, 1, 1988.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselinss: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268): 2021. URL http: //tt:::/mmm////////////// org http://ppppas//22// 20-1134- html.
Alexander Rakhlin, Ohad Shamir, and Karthik Sridharan. Making gradient descent optimal for strongly convex stochastic optimization. arXiv preprint arXiv:1109.5644 2011.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The Annals of Mathemat- ical Statistics, pp. 400-407, 1951.
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 661 1668. JMLR Workshop and Conference Proceedings, 2010.
Mark Rudelson and Roman Vershynin. Hanson-wright inequality and sub-gaussian concentration. 2013.
David Ruppert. Efficient estimations from a slowly convergent bubbins-monre process. Technical report, Cornell University Operations Research and Industrial Engineering, 1988.
exploration. In Advances in Neural Information Processing Systems, pp. 2256-22566 2264, 2013. Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic
Mark Sandler, Andrey Zhmoginov, Max Vladymyrov, and Nolan Miller. Training trajectories, mini- batch losses and the curious role of the learning rate. arXiv preprint arXiv:2301.02312. 2023.
15

---

Published as a conference paper at ICLR 2024
Sunny Sanyal, Jean Kaddour, Abhishek Kumar, and Sujay Sanghavi. Understanding the ef- fectiveness of early weight averaging for training large language models. arXiv preprins arXiv:2306.03241, 2023.
Tuhin Sarkar and Alexander Rakhlin. Near optimal finite time identification of arbitrary linear dynamical systems, In International Conference on Machine Learning, pp. 5610-5618. PMLR, 2019.
Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International confer- ence on machine learming, pp. 343-351. PMLR, 2013.
Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through crowded valley- benchmarking deep learning optimizers. In International Conference on Machine Learning, pp- 9367-9376. PMIR, 2021.
Thibault Sellam, Steve Yadlowsky, Jason Wei, Naomi Saphra, Alexander Amour, Tal Linzen, Jasmijn Bastings, Iulia Ture, Jacob Eisenstein, Dipanjan Das, et al. The multiberts: Bert repro- ductions for robustness analysis. arXiv preprint arXiv:2106. 16163, 2021.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint arXiv: 108009908..7909 2015.
Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing sys- tems, 35:22955-22968 2022.
Max Simchowitz. Making non-stochastic control (almost) as easy as stochastic. Advances in Neural Information Processing Systems, 33: 13118888888329 2020.
Max Simchowitz and Dylan Foster. Naive exploration is optimal for online lqr. In International Conference on Machine Learning, pp. 1937-8948. PMLR, 2020.
Max Simehowitz, Horia Mania, Stephen Tu, Michael Jordan, and Benjamin Recht. Learning without mixing: Towards sharp analysis of linear system identification. In Conference On Learning Theory, pp. 439-473. PMLR, 2018.
Max Simchowitz, Karan Singh, and Elad Hazan. Improper learning for non-stochastic control. In Conference on Learning Theory, pp. 3320-3436. PMLR, 2020.
Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V Le. Don't decay the learning rate, increase the batch size. arXiv preprint arXiv:1711.00489, 2017.
Hyung Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake. Do differentiable simulators give better policy gradients? In International Conference on Machine Learning, pp. 20668- 20696. PMLR, 2022.
Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based RL in contextual decision processes: PAC bounds and exponential improvements over model-free approaches. In Conference on learning theory, pp. 2898-2933. PMLR, 2019.
Matus Telgarsky. Feature selection with gradient descent on two-layer networks in low-rotation regimes. arXiv preprint arXiv:2208.02788 2022.
Emanuel Todoroy, Tom Erez, and Yuval Tassa. Mujaco: A physics engine for model-based control.
In 20/2 IEEE/RSJ International Conference Intelligent Robots and Systems, pp. 5026-5033.
IEEE, 2012. doi: 10. 10.1
Mark Towers, Jordan K. Terry, Ariel Kwiatkowski, John U. Balis, Gianluca de Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Perez- Vicente, An- drea Pierré, Sander Schulhoff, Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymna- sium, March 2023. URL https: zenodo. org/ record/ B1 2127.25.
Stephen Tu, Alexander Robey, Tingnan Zhang, and Nikolai Matni On the sample complexity of stability constrained imitation learning. In Learning for Dynamics and Control Conference, pp- 180-191. PMLR, 2022.
16

---

Published as a conference paper at ICLR 2024
Arjan J Van Der Schaft and Hans Schumacher. An introduction hybrid dynamical systems, volume 251. springer, 2007.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Ruosong Wang, Dean P Foster, and Sham M Kakade. What are the statistical limits of offline rl with linear function approximatio?? arXiv preprint arXiv:2010. 1895, 2020a.
Ruosong Wang. Russ Salakhutdinov, and Lin Yang. Reinforcement learning with general value function approximation! Provably efficient approach via bounded eluder dimension. Advances in Neural Information Processing Systems, 33, 2020b.
Ruosong Wang, Yifan Wu, Ruslan Salakhutdinov, and Sham Kakade. Instabilities of offline rl with pre-trained neural representation. In International Conference on Machine Learning, pp. 10948- 10960. PMLR, 2021.
Tyler Westenbroek, Max Simchowitz, Michael Jordan, and S Shankar Sastry. On the stability of nonlinear receding horizon control: geometric perspective. In 2021 60th IEFE Conference on Decision and Control (CDC), pp. 742-749. IEBE, 2021.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lppes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learming, pp. 23965-23998, PMLR, 2022.
Yasin Yaz, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap. Georgios Piliouras, Vijay Chan- drasekhar, et al. The unusual effectiveness of averaging in gan training. In International Confer- ence on Learning Representations, 2018.
Dante Youla, Hamid Jabr, and Jr Bongiorno. Modern siener-hopf design of optimal controllers-pe -part ii: The multivariable case. IEEE Transactions Automatic Control, 21(3):319-338, 1976.
Andrea Zanette. Exponential lower bounds for batch reinforcement learning: Batch rl can be ex- ponentially harder than online rl. In International Conference on Machine Learning, pp. 12287- 12297. PMLR, 2021.
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. study on overfitting in deep reinforcement learning, arXiv preprint
2018.
17

---

Published as a conference paper at ICLR 2024
CONTENTS OF APPENDIX
A Discussion and conclusion
18
Additional related work
19
B.1 Stochastic optimization and deep learning
19
B.2 Stability reinforcement learning and control
20
Full experimental results
22
C.I MuloCo continuous control environments
22
C.2 Natural language generation
................................
30
C.3 Linear systems
35
C.4 LQR with marginally stable dynamics
C.5 LQR with cliff
36
Theoretical analysis of toy examples: details and additional results
38
D.I
Error amplification in linear dynamical systems
39
D.2 Benefits EMA for the cliff loss
40
D.3 For optimal convex learning rates, EMA does not help
44
E Proofs
46
E.1 Proof of Proposition D.1
46
E.2 Proof of Proposition D.2
47
E.3 Cliff loss for Gaussian random vectors (Proof of Proposition D.3)
48
E,4 Proof of Proposition D.5
49
E.5 EMA for continuous Gaussian processes
E.6 Instantiations of cliff loss in continuous time
57
E.7 Proof of Theorem
58
A DISCUSSION AND CONCLUSION
We investigated the causes of instabilities in training neural networks in tasks involving feedback loops, and empirically demonstrated that these instabilities often arise result of poor algorithmic choices rather than information-theoreti barriers. We hypothesize that the oscillations are caused by the combination of fractal reward landscapes and noise floor the stochastic gradients, and can be ameliorated by variance reduction-eilte either during training, or post hoc (via a stabilizer such as EMA). We now list some directions for future inquiry, containing theoretical and empirical myster- ies.
Unstable optimizers & trajectory stabilization. Despite many attempts to design faster- converging optimizers for deep learning, there has been no clear successor to vanilla recond-moment preconditioning (Duchi et al., 2011: Kingma Ba, 2014) with well-tuned learning rate schedule (Agarwal al., 2020; Schmidt et al., 2021; Kaddour et al., 2023). By decoupling the roles of opti- mizers and stabilizers, we open new avenue for innovation this space. When training deep neural networks in the presence of feedback loops, findings suggest: always try EMA whenever it is affordable, and train at higher learning rates without decaying to Investigations nuanced downsides, larger-scale settings, and new optimizer-stabllizer algorithms left for future work.
18

---

Published as a conference paper at ICLR 2024
Representation learning in RL. While this work focused fully offline data, online RL is also known to suffer from serious training instability, often thought to be caused by the statistical chal- lenges of interactive learning. On the other hand, GVA appears to be less pronounced in this setting," and the interaction between GVA and training with "better" data remains mysterious.
On related note, the inductive biases of iterate averaging for representation learning remain poorly understood. Many algorithms for representation learning benefit from EMA (Grill et al., 2020; Yaz al., 2018); to what extent can this be explained by a phenomenon related to GVA?
Finally, we observe that our results are consistent with the notion that SGD noise is only harmful, but this cannot be entirely true due to the fact that sufficiently large batch sizes hurt training; we leave a better understanding of when and why SGD noise is helpful for future work.
B ADDITIONAL RELATED WORK
In this section, we provide an extended discussion of related work. We split our discussion into two sections, beginning with our primary focus on stochastic optimization (especially as applied to deep learning) and ending with works relevant to one of our empirical domains, reinforcement learning and control.
B.1 STOCHASTIC OPTIMIZATION AND DEEP LEARNING
Stochastic optimization in deep learning. Minibatch SGD (Robbins & Monro, 1951) is the workhorse of deep learning, and its nuances are explained by mechanisms far outside the purview of stochastic convex optimization. Batch sizes and learning rates are known to control variance (Smith al., 2017); however, not all variance reduction techniques work in practice (Defazio & Bottou, 2019). Nonetheless, analyses of SGD in simplified settings have proven to be useful (Mandt et al., 2017: Malladi et al., 2022; Li al., 2019a; 2020). The theoretical literature contains numerous accounts of why finite-sample gradients might be beneficial, e.g., posterior inference (Mandt et al., 2017) and saddle point escape (Jin et al., 2017). On the other hand, we show that this stochasticity can lead to catastrophic error amplification in regimes of practical interest.
Iterate averaging. It known that returning the last iterate of even deterministic gradient de- scent is suboptimal, even when the function being optimized is convex (Harvey et al., 2019) (but non-smooth). Which averaging scheme optimal depends on whether the underlying function is strongly convex or not, and discuss the various tradeoffs in Appendix D.3. Importantly, there exists memory-efficient averaging schemes (storing at most 2 optimization iterates) which at- tain optimal rates (Lacoste-Julien et al., 2012). The optimality of EMA for stochastic optimization has not been studied in the literature, and indeed, we show in Appendix D.3.1 EMA is suboptimal even for deterministic optimization via gradient descent whenever the averaging parameter ? satisfies ~ T with 3 < 1. This theoretically suboptimal regime is precisely the one in which observe strong mitigation of GVA in our experiments. In fact, the case of 8 1, which is analyzed in Lacoste-Julien et al. (2012) in the convex case, is empirically suboptimal and leads to a drastic reduction in performance early in training, which is then washed out later iterations. The stabilizing behavior of iterate averaging of machine learning in practice is well-known (e.g., Bottou (2010), although it has received less emphasis in modern deep learning than other optimizer hyperparameters.
In addition to theory, iterate averaging in general, and EMA in particular, is extremely common in practice (Izmailov al., 2018; Kaddour, 2022; Kaddour et al., 2023; Izmailov et al., 2018; Busbridge et al., 2023; He et al., 2019). In contradistinction to these works, the present paper is not purporting to examine iterate averaging in and of itself, but rather is exploring the effects of iterate averaging in the presence of GVA.
Non-convex side effects of variance reduction techniques. Here, expand the discussion from Section 3.1, concerning why decreasing the learning rate and increasing the batch size are not satisfactory interventions for GVA:
In small anecdotal exploration, we found that even minimal online interventions to improve data coverage such DART (Laskey et al., 2017) appear to reduce oscillations significantly.
19

---

Published as a conference paper at ICLR 2024
Learning rate: is understood in practice that large learning rates promote faster and more robust representation learning in neural networks; theoretical understanding is an active fron- tier of research, which requires analyzing the trajectory of SGD in the rich (data- (dta--ddapted) regime: see (Li & Yuan, 2017; Li et.,, 2019b; Chizat et al., 2019; Barak et al., 2022; Damian al., 2022; Telgarsky; 2022; Andriushchenko et al., 2023) and references therein." 10
Batch size: It is also well-understodd in practice that the mechanisms enabling successful neural network training break down at extremely large batch sizes. This is especially relevant in the quest for massively parallel training pipelines (Goyal al., 2017). One way to analyze this theoretically is to note that smaller-batch SGD can more effectively memorize individual examples (Abbe & Sandon, 2020; Abbe et al., 2021), 1), often desired behavior.
Of greatest relevance to our investigation of iterate averaging, Sandler et al. (2023) presents some related empirical and theoretical results in the context of supervised learning and explores an equiv- alence between learning rate decay and EMA. Their empirical results focus on pure supervised learning and examine the relationships between training trajectories with different learning rates and how iterate averaging affects this, while we primarily focus on the phenomenon of GVA and the effect of iterate averaging therein.
SDE limits and analysis of training in deep learning. Taking the continuous time limit of adap- tive gradient optimization procedures and modeling the result as the solution to a stochastic differ- ential equation has seen lot of recent interest; see Mandt et al. (2017); Cheng et al. (2020); Li et al. (2020); Kunin et al. (2020); Li et al. (2019a); Malladi et al. (2022) and the references therein. Of greatest relevance our paper is Busbridge al. (2023), which proves the convergence of constant rate SGD along with constant rate EMA to system consisting of an SDE and an ODE. Proving rig- orous convergence guarantees in this limit is beyond the scope of this work; instead, for one of the results in Appendix D. we take the continuous time limit as given and precisely compute in several toy settings the benefits that EMA has.
Error amplification and closed-loop effects in language generation. The deep RL and deep NLP literatures are highly connected, and use common vocabulary to describe considerations in- volving feedback loops (e.g. training on the BC/autoregressive loss known as teacher forcing (Arora et al., 2022)). Long-range error amplification stemming from autoregression has been noted in prior work (Holtzman et al., 2019; Braverman et al., 2020) (sometimes called exposure bias), motivating the design of decoding heuristics.
B.2 STABILITY IN REINFORCEMENT LEARNING AND CONTROL.
Instabilities of online RL. Many works have empirically observed instabilities in online RL, often arising as a result of distribution shift (Kirkpatrick et al., 2017; Henderson et al., 2018; Machado et al., 2018; Packer al., 2018; Zhang et al., 2018; Engstrom et al., 2019; Andrychowicz et al., 2020). On the theoretical side. most existing work on online RL explores structural assumptions (e.g.. Bellman rank) that allow carefully designed algorithms control distribution shift (Russo & Van Roy, 2013; Jiang et al., 2017; Sun et al., 2019; Wang et al., 2020b; Du et al., 2021; Jin et al., 2021; Foster et al., 2021a), but does not directly address the behavior cloning setting we consider, in which instability arises despite the absence of distribution shift.
Instabilities of offline RI. While many works have empirically observed instabilities in online RL (Kirkpatrick et al., 2017; Henderson et al., 2018; Machado et al., 2018; Packer et al., 2018; Zhang et al., 2018; Engstrom al., 2019; Andrychowicz et al., 2020), it is generally understood that these instabilities arise as a result of the distribution shift caused by the complicated dependence structure of interactive data. In contrast, we observe instabilities in the offline setting, where the data is fixed and the dependence structure is significantly simpler. In offline RL, too, many works have observed instabilities (Fujimoto al., 2019; Kumar al., 2019), although several theoretical works suggest that the source of these instabilities may be information theoretic limits as opposed algorithmic issues (Wang et al., 2020a; Zanette, 2021: Foster et al., 2021b).
10The is also regime in which small learning rates promote generalization (Bousquet & Flisseeff, 2002; Hardt et al., 2016), but this prevails predominantly in online and convex problems.
20

---

Published as a conference paper at ICLR 2024
In particular, Wang et al. (2021) find that conducting offline RL using pre-trained representations often suffers from catastrophic error amplification. The authors there apply known planning algo- rithms top of pre-trained representations and find that even in the presence of mild distribution shift, severe degradation in rollout reward occurs. Dong et al. (2020) observe fractality in continuous MDPs as a generic phenomenon; like Wang et al. (2021), they are primarily interested in learning and planning as opposed to imitation.
(In-)stabilig in behavior cloning. Behavior cloning is known to suffer from compounding error, phenomemon in which the imitators's errors accumulate over time and place the learner in states not observed during training. As we confirm in this paper, these instabilities are readily mitigated by forcing the expert to demonstrate corrections from imperfect demonstrations, either by forcing experts to correct from noisy states (e.g. DART (Laskey et al., 2017)), or by sequential iteraction with the environment (e.g. DAGGER (Ross & Bagnell, 2010)); both interventions require online interaction with the expert and thus lies outside the scope of this paper. Direct data augmentation of the expert trajectories via noising is popular in practice (Ke et al., 2021), but does not mitigate training oscillations.
More recent theoretical work (Tu et al., 2022; Pfrommer et al., 2022; Havens Hu, 2021) has con- nected behavior cloning to the theory of incremental stability of the induced closed loop dynamics (Angeli, 2002). Leveraging this stability often requires stronger oracles, such as the ability to com- pute first- or zero- order gradients of the expert policy (Pfrommer al., 2022). Recently, Block al. (2023a) proposed a combination of probabilistic and control-theoretic stability to facilitate imitation of non-smooth expert trajectories in smooth dynamical systems, assuming access type of local stabilization oracle.
Comparison to Emmons et al. (2021); Chang et al. (2023b). These works investigate the fact that bBC on validation set is a poor proxy for JH- In Emmons et al. (2021), through a suite of experiments in a number of domains motivated by robotics, the authors probe how far applying only supervised learning can get in producing a high reward policy. They focus primarily on architectural and regularization interventions including depth and width of their imitator policies as well as the effects of weight decay and dropout on the resulting policies, noting the weak relationship between reward and fBC af the end of training in all of these examples. In contrast to this work, we focus on causes and mitigations of instabilities during truining as well as the more general phenomenon applied beyond continuous control tasks.
Chang et al. (2023h) finds empirically that BC evaluated on a held out validation set is an extremely poor empirical proxy for rollout rewards in many common benchmarks during a single training run, but were not primarily interested in the instability itself, and thus did not further investigate this phenomenon.
The role of stability in learning to control. Considerations of stability (in dynamical sense, rather than a training sense) have been most visible in the recent literature on learning to control linear systems. While stability is not always essential estimation in control systems (Simchowitz et al., 2018; Sarkar & Rakhlin, 2019; Ghai et al., 2020), much of the contemporary research on learning to control linear systems has assumed access an initial stabilizing control policy (Fazel et al., 2018; Simchowitz & Foster, 2020; Dean et al., 2018). Indeed, without access to such a con- troller, certain adaptive control settings require regret scaling exponentially in problem dimension (Chen & Hazan, 2021), until the system can be estimated to sufficient accuracy that a stabilizing controller can be synthesized (in some circumstance, however, these effects can be cireumvented one can construct simulators with reward discounting (Perdomo et al., 2021)). Access to an initial stabilizing controller forms the basis for both classical (Youla et al., 1976) and more recent (Agarwal et al., 2019; Simchowitz et al., 2020; Simchowitz, 2020) convex control parameterizations.
Stability appears more implicitly in the recent literature on nonlinear control. For example, assuming that the output of the dynamical map is bounded and that the system exhibits process noise amounts to a form of marginal stability in Wasserstein and total variation distances (see e.g. Kakade et al. impor-
(2020)). Lastly, stabilizability of linearizations of nonsmooth dynamics have emerged tant desideratum for optimization-bast and data-driven trajectory optimization in smooth nonlinear systems (Pfrommer al., 2023; Westenbroek et al., 2021).
21

---

Published as a conference paper at ICLR 2024
Stiffness in contact and hybrid dynamics. System stability is at least as bad system "stiff- ness", or largest Lipschitz constant dynamics, which can be arbitrarily large in contact and hybrid dynamical systems even if one does not observe the sort of exponential blow up witnessed in unstable linear dynamics. At the limit, dynamical discontinuities induce "one-step" instability. Contending with systems stiffness has been long-standing challenge in the field of robotics, where contact sys- tems are modeled non-smoothly (Van Der Schaft & Schumacher, 2007). Recent work has explored the use of stochastic smoothing to circumvented these challenges in trajectory optimization of known stiff system (Parmas et al., 2018; 8; Suh et al., 2022). Further work has shown that stochastic smoothing makes sequential prediction and planning problems tractable for piecewise affine sys- tems (Block & Simchowitz, 2022; Block et al., 2023b;e). Whether the dynamics of the noiseless system are well-approximmtee by those of the same system with process noise added remains an open question.
C FULL EXPERIMENTAL RESULTS
In this section, we provide additional experimental results as well as full details for the experiments included the main text. The appendix is organized as follows:
In Appendix C.1 by providing detailed description of the experimental setup we use in our empirical MuJoCo results.
In Appendix C.I.I, we provide an extended discussion of visualizations related Figure 1. We continue in Appendix C. 1...bbb by detailing our large suite of experiments investigating the effects that various algorithmic interventions have on GVA, both qualitatively and quantita- tively. These correspond to the findings highlighted in Figure 2 and Figure 3 in the main paper.
In Appendix C.1.3 provide several additional visualizations, including number of re- ward landscapes similar to Figure 1 (right) to ensure the robustness of our conclusion that the reward landscape extremely sensitive the parameter. We also demonstrate that the successful Transformer imitators learn Markovian policies; unlike the MLPs, the Trans- formers observe history (of length 32 our experiments) of many past states and actions; in Appendix C.1.3 we show that the Transformers are attending to previous states and not just the most recent one; this corroborates findings Janner al, (2021); Shafiullah et al. (2022).
In Appendix C.2, we give additional details and results relating to our autoregressive lan- guage generation experiments, described Section 4.2.
In Appendix C.4, we present synthetic experiments on 2-dimensional linear dynamical sys- tems, which serve as accompaniments to our thcoretical vignettes. We show empirically that GVA does not occur in marginally stable linear systems, and show that natural dynami- cal systems analogue of the cliff loss described in Section 4.3 indeed exhibits the predicted GVA in Appendix C.5; we also observe that, as predicted by the theory, EMA mitigates the oscillations.
All of our experiments are implemented in PyTorch (Paszke et al., 2019), and we use the x-transformers: implementation of the Transformer architecture. All of our MuJoCo tasks are in Gymnasium (Towers al., 2023). For our LQR experiments, we implement the dynamics directly in PyTorch.
C.I MuJoCo CONTINUOUS CONTROL ENVIRONMENTS
In all of our MuJoCo experiments, we run the following pipeline:
1. Train expert using online reinforcement learning.
2. Collect expert trajectories by rolling out the (deterministic) expert policy 1000 times.
3. Split the collected trajectories into a training set of size N = 900 and validation set of size 100 trajectories.
"https://g"" thub.
formers
22

---

Published as a conference paper at ICLR 2024
SAC expert training curves & final-checkpoint reward distributions
We -
Hopper-v4
HallCheetah-v4
Ant-v4
Humandid-44
Normalzed Rzed reward
1.0
M1
1.
www
0.5
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.9
0.0
Training steps
Training steps
Training steps
Training steps
Training steps
Frequency
5726.3
3710.4
μ 10694.5 5
6436.2
6515.9
58c0
10000
11000
8400
6500
Reward
Reward
Reward
Reward
Reward
Figure 5: Training curves of the SAC experts for the various MuJoCo continuous control agents, along with final-iterate mean rewards. Top: Online reinforcement learning training curves; these exhibit training instabilities, but not of the same nature as those encountered in our offline behavior cloning settings. Bottom: Unnormalized reward distributions (through out the rest of this paper, divide by these means). Outliers are marked by x symbols.
4. Train an agent to imitate the expert from these N offline trajectories, using standard deep learning pipelines optimize the behavior cloning loss BBC-
5. Evaluate the behavior cloner by rolling out the learned policy and evaluating loss BCC the validation set.
We now discuss some of the items above in more detail.
Training the expert. All of our experts are trained using the Soft Actor Critic (SAC) algorithm (Haarnoja al., 2018) implemented by default parameters for the environment under consideration. We train our experts for million steps; training curves and reward histograms (over the randomness of environment initial conditions) shown in Figure 5.
(Raffin et al., 2021) with the
Collecting expert trajectories. While the SAC online RL algorithm results in a stochastic policy, are interested in deterministic experts. Thus, derandomize the expert's outputs, via taking the mean the action distribution of the randomized SAC agent. We roll out trajectories of length 1000 for every environment. Interestingly, the presence of extreme outliers (Ant) and sighesiinnsional states (Humano id) appear to be uncorrelated with the magnitude of GVA (which is very low for both of these tasks).
Training the behavior cloner. This is the step that our work is interested in and is where all of our interventions occur. For every environment, unless specified otherwise, we train the cloner for 20 epochs. Our default hyperparameters included in Table 1, although each intervention below changes at least one of them.
Evaluating the behavior cloner. For most training runs, we collect checkpoint every 5000 gra- dient steps, with the exception of the in Figure where collect checkpoints every gradient step between iterates 40K and 50K in order to have more zoomed-in view the oscillations. For each checkpoint, we compute (nc the validation each seed, sampling trajectory beginning from the seed-induced initial condition and rolling out the learned policy. Note that the only randomness here in the initial condition of the agent, as determined by the seed.
well fixing 20 random seeds and, for
Choice of canonical environment. Overall, we choose to focus our experiments on Walker-vil for the following reasons: (1) reward oscillations are significantly more benign for cloners of Ant
23

---

Published as a conference paper at ICLR 2024
and Humanoid agents; (2) we could not get Hal @Cheetat behavior cloner to reach the max- imum reward. We choose to focus our intensive hyperparameter sweeps on Walker (rather than Walker and Hopper) to keep the number of runs manageable.
Name
Val uu
Optimizer
AdamW
Learning rate
3 x 10
Epochs
20
Batch size
8h (momentum)
0.9
Bg
0.999
Weight decay
0.1
(MLP) Numher of layers
(MLP) Width
1024
(MLP) Activation
ReLU
(Transformer) Number of layers
(Transformer) Number of heads
('ransformer) Embedding dimension
1024
(Trunsfonmer) Positional encoding
sinusoidal
(Transformer) History length
32
Table 1: Table of hyperparameters for behavior cloning.
C.I.1 REWARD SENSITIVITY, PER-ITERATE OSCILLATIONS, AND LANDSCAPES
Recall the first finding from the main paper:
(R1) The reward landscape is highly sensitive to small changes in policy parameters.
Figure provides closer look at this landscape, via additional visualizations and discussion for the example of GVA from Figure 1:
(a) Per-iterate fluctuations: Zooming into the training trajectory, extremely large fluc- tuations at the level individual gradient descent iterates (blue dots). Rewards can vary by 50% of the total maximum reward. Meanwhile, the behavior cloning loss varies at the scale of < 1% of the range of BBC.
(b) Weak correlation between JH and fBC: For many pairs of iterates, the behavior cloning objective and rollout rewards are negatively correlated. Note that cloners achieve high re- wards without converging to perfect action recovery (red star). Representation clearly not the bottleneck: the expert agent is a 2 2-layer MLP with hidden dimension 128, which is realizable by all of these networks.
(c) Exhaustive evaluation of cross-sections: For checkpoints (iterations 115000 and 120000), selected to exemplify negative progress as training progresses, we evaluate neigh- borhoods of these policies in 2-dimensional cross-sections of the reward landscape. We compute stochastic gradients g1 82 of fBc, for two randomly selected minibatches of size 32. For a choice step size r ing rate), we evaluate rewards and BC losses for the lattice of policies {0 + i 771 + j (EB)
10-4 (coinciding with the Adam optimizer's leam-
These rewards are averaged over 20 random seeds determining initial conditions of the environment, which are not re-randomized between evaluations. Thus, we conclude that these fluctuations arise from extremely high sensitivity to initial conditions. Aside from providing qualitative visualization that GVA is evident, this evalu- ation provides a lower bound of 2 x 10 on the Lipschitz constant of + (Te) (choosing the & norm for 0), even restricted to non-adversarial stochastic gradient directions. The loss landscape (depicted in Figure Ic, and omitted here) is nearly linear, and much smoother (Lipschitz constant 50).
In addition, in Figure we repeat Figure 6(a) but with each mean reward disaggregated such that each point represents the rollout reward of fixed initial condition of the policy at the current iterate. The points are colored by initial condition so as to illustrate that it is not the case that some initial conditions significantly more difficult than others.
24

---

Published as a conference paper at ICLR 2024
Every-iterate evals (steps 40K-50K)
Every-iterate evals (steps 40000-40500)
10
*
cloner reward Cloner reward
0.
policy
EMA cloner
02
0.8
0.7
0.0
0.0
0.6
Cloner
0.4
0.5
eerroo
0.2
40000
40100
40200
40300
40400
40500
0.00
0.04
0.06
Traning
(a) Per-iterate mean rewards and losses.
(b) Rewards Jr (To) vs. (TO).
Step (reward 0.85)
Step (reward 0.54)
1.0
Meam come
37
0.5
2.
5.5
200
400
d
-10
12
Unindaten grad
10 glad
50
50
(c) 2- 2-dimensional cross-sections (To i i- mg=)) of the reward landscape.
Figure 6: Additional visualizations for the experiments highlighted in Figure 1. (a) Rewards Ju (mo) and (Te) plotted at every training iteration in small range. (b) Scatter plot of J (TT) against (T((T)) for the full range of every-iterate evaluations displayed in Figure (c) Reward landscapes shown around good (left) and bad (right) iterates. Note the extremely high Lipschitz constant (local sensitivity) of the rollout reward in terms of policy parameters; furthermore, there exist good and bad policies in very close proximity along stochastic gradient directions.
Individual rewards (staps 40000-40500)
Cloner reward
0.8
0.2
1.0
EMA cloner reward
0.8
0.4
40000
40100
40200
40300
40400
40500
Figure 7: Disaggregated version of Figure 6(a). We visualize rollout rewards of each rollout for each iterate experiment in Figure with points colored by environment random seed (which only determine initial conditions). We exhibit both the non-EMA (top) and the EMA (bottom) runs.
25

---

Published as a conference paper at ICLR 2024
C.1.2 EFFECTS OF ALGORITHMIC INTERVENTIONS
Below, we list the exhaustive range of algorithmic interventions probed for in the Walker2d-vs environment. We begin with the suite of experiments focused on diagnosing GVA (highlighted in Figure 2):
(Ia) Architectures: For both (MLP, Transformer) architectures, we vary the model scale pa- rameters: depths {2, 4, 6) and embedding dimensions {128, 256, 512, 1024]. For all other experiments below, test interventions the 4-layer 1024-dimensional MLPs and Trans- formers.
(2) Regularization controls: We consider varied feedforward dropout coefficients ((0, 0.1, 0.5)) for both MLPs and Transformers, and embedding dropout coefficients the same range for Transformers. We also try weight decays (0.01, 0.3), losses (6, f2). and data augmentation scales {0.005, 0.01, 0.02].
(3) Optimizer configurations: We vary the momentum and adaptive preconditioner param- eters (2) of AdamW, noting that these are known to modulate training stability and generalization. We vary 31 E {0.7, 0.9, 0.95, 0.99) and 82 {0.9, 0.99, 0.999, 0.9999). 32 0.9999 leads to many non-convergent runs, and are omitted from Figure 9.
(4) Variance reduction controls: We vary the power in the learning rate schedule ne (1-4)- for E 0.5, 1, 1.5, 2, 2.5, 3); note that a = 0 corresponds to removing the learning rate decay. We also try batch sizes of 128 and 512 via 2x, 4x and 16x gradient accumulation with the default batch size of 32.
Next, we list the training runs selected to investigate the scope and hyperparameter sensitivity of EMA (highlighted in Figure 3);
(Ib) Architectures: For all of the architecture choices in (GI) above, we evaluate EMA cloners.
(5) EMA parameter ablations:
To check the robustness and hyperparameter sensitivities of EMA, we perform search in the hyperparameter space of update frequencies, burn-in du- rations, and annealing powers. Throughout these sweeps, we clip the annealing coefficient to minimum of Ymin parameter depends on the number of training steps; we recommend setting a low Yimin and tuning the annealing power instead.
10 10-4. Although this can be harmful, the optimal choice of this
(6a.b) A Ablations of sample size and learning rate decay: We downsample the 900 training trajec- tories, instead training agents on (10, 20, 50, 100, 200, 500} demonstrations. In all of these settings, remove learning rate decay, and find that this improves generalization (as long as EMA applied). We evaluate both non-EMA and EMA models.
(7a.) Additional environments: We perform a limited set of analogous comparisons between EMA and non-EMA cloners for other continuous control tasks provided by gymnasium. We evaluate both non-EMA and EMA models.
Quantitative evaluations. Stochastic optimization in deep learning is fundamentally non- stationary, in the sense that representations change over the course of the training trajectory. Thus, there is canonical choice of summary statistic which losslessly captures the magnitude of GVA. Nonetheless, to provide a digestible comparison of algorithmic interventions and their effect on GVA, record number of basic quantitative metrics:
Maximum and final rollout rewards Juax, Jünal- These serve as indicators of whether the BC training procedure ever arrives good policies, modulo stability considerations. GVA then entails lack of stable convergence to these iterates, in terms of rollout performance. We similarly record the minimum and final behavior cloning objectives Emin, Liiaa-
Reward statistics jmid, rangemid from iterates in the middle 50% of training. We track the mean and range, to summarize the amplitude of G GVA-induced reward oscillations through the course of training.
Indicators of early success and negative progress tearly tworse. We record the earliest iteration (as a percentage of the total number of training steps) at which the model visits a checkpoint whose mean reward is > 95% of the best one in its training trajectory. We also
26

---

Published as a conference paper at ICLR 2024
Architsccures, EMA
(1b)
Aue
Eee
Few
Fre
tangPen
Ewy
Taoe
foe
few
Fee
Zeet
rengP..
Guzy
Texe
/L.125d MIP.
0.84
0.86
0.073
0.075
0.49
0.51
0.66
0.83
0.73
0.075
0.073
0.30
0.52
0.57
4.125dMMLP
1.00
0.054
0.60
6L (L.128d MLP
1.00
1.00
0.051
0.051
0.77
0.55
0.33
0.99
1.00
0.97
0.050
0.050
0.86
0.24
0.24
0.96
2L 256d MLP
1.00
1.00
0.050
0.050
0.91
0.44
0.27
0.96
1.00
1,00
0.050
0.050
0.96
0.13
0.18
0.79
2568
1.00
1.00
0.030
0.030
0.94
0.72
1.00
0.030
0.99
0.06
0.09
0.39
21 512 MLP
1.00
1.00
0.031
0.031
0.98
0.20
0.10
0.90
1.00
1,00
0.031
0.031
1.00
0.05
0.07
0.25
512 MLP-
2.00
1.00
0.021
0.021
0.98
0.32
0.08
0.66
1.00
1.00
0.021
0.021
1.00
0.03
0.07
0.00
1024 MLP..
2.01
1.00
0.020
0.023
0.98
0.17
0.06
0.62
1.00
1.00
0.020
1.00
0.03
0.06
0.34
1024 MLP--
1.00
1.00
0.016
0.016
0.95
0.19
0.14
0.60
1.00
1.00
0.016
0.016
1.00
0.03
0.06
0.07
1024
1.00
0.015
0.015
0.015
0.99
0.06
0.11
0.05
0.06
AL 1286 Tt
1.00
1.00
0.027
0.027
0.90
0.59
0.10
0.94
1.00
1.00
0.026
0.026
0.99
0.08
0.06
0.45
GL 1211286T Tt
1.00
1.00
0.025
0.025
0.87
0.50
0.10
0.89
1.00
1,00
0.024
0.024
0.98
0.09
0.06
0.59
T1
1.00
1,00
0.019
0.019
1.00
0.05
0.07
0.45
AL. 2566 TL
1.00
1.00
0.017
0.017
0.91
0.39
0.13
0.53
1.00
1,00
0.017
0.017
0.99
0.07
0.06
0.44
11261661 Tt
2.00
1.00
0.019
0.019
0.93
0.10
0.90
1.00
0.97
0.018
0.018
0.99
0.05
0.06
5126 T(-
1.00
1.00
0.014
0.014
0.80
0.29
1.00
1.00
0.014
0.98
0.10
0.06
0.57
2L 10246 T[-
1.00
1.00
0.016
0.016
0.93
0.35
0.09
0.83
1.00
1.00
0.015
0.015
0.99
0.06
0.07
0.35
10240 Tt
0.013
0.013
0.84
1.00
0.013
0.98
0.30
0.07
0.06
Figure 8: Groups (la) and (1b): Non-EMA vs EMA rollout evaluations for a variety of architec- tures. The stabilizing effect of EMA is robust across these architectures and scales. This is most clearly seen in the rangemid column of (1a), and robustly confirms our empirical result (RI). Look- ing at the analogous column in (1b), we confirm (R4), that EMA uniformly stabilizes the training of these architectures.
record the latest iteration where the reward dips below this value. (In the table, denotes the absence of such iteration.)
For all runs, we report the median of these statistics over 3 random seeds. Using these quantitative results, we next restate and discuss the relevant empirical findings from the main paper, and discuss these in greater depth in the figure captions. Results are shown in Figures 8, 9, 10, and 11.
(R2) GVA arises from algorithmic suboptimality rather than an information-theorreiee limit. Even with "infinite" training data (i.e. fresh trajectories with i...d. initial conditions at each training step), rollout oscillations persist.
(R3) Training oscillations are not mitigated by many standard approaches to regularization, including architectural interventions and increased regularization. On the other hand, oscilla- tions are ameliorated by variance reduction techniques, such as large batch sizes, learning rate decay, and iterate averaging.
(R4) EMA iterate averaging strongly mitigates rollout oscillations, and its performance is robust across a variety of architectures and environments.
We close with a final visualization: Figure 12 shows scatter plot of JH vs. fB for the final iterates of all experimental conditions considered in this paper, analogous to Figure 6(b), showing that a generalization gap remains vs. the true (deterministic) expert policy, and that rollout reward can be uncorrelated with this gap in the regimes studied in this work.
C.1.3 MISCELLANEOUS VISUALIZATIONS
Additional reward and loss landscapes. We provide additional (lower-resolution) heatmaps around neighboorhoods of various cloner policies in Figure 13. In addition to stochastic gradi- ent directions, we visualize cross-sections of these landscapes in random directions, like the main figures in (Li et al., 2018).12 12 This confirms the generality of our finding in (R1), concerning the
Note that Li et al. (2018) (08) choose random directions because the loss landscape looks too flat in stochastic gradient directions. This provides further supporting evidence that our finding (R1) is generie: while the loss
27

---

Published as a conference paper at ICLR 2024
Jeu
[2] Regularizers
Auc
Aus
Fer
[2) Optiiitttttionn
Ale
lurglee
Two
[AL 10254 HLPI
1.00
1.00
0.016
0.016
0.39
0.19
0.34
0.60
infovic cola-
1.00
1.00
0.014
0.014
0.56
0.35
8.11
0.31
1.00
1.00
0.016
0.016
0.80
0.54
1.00
A. 07,0
1.00
1,00
0.016
0.016
0.36
0.24
6.09
0.52
R. 6. 07, 0.599
1.00
1.00
0.016
0.816
0.89
0.44
0.88
0.80
1378 sug. 9.005-
1.00
1.40
0.821
0.021
0.85
0.64
0.25
0.99
A, 0.522
1.00
1.00
0.016
0.016
0.59
0.33
zug.
1.00
1.00
0.018
0.16
coto zug 5.05
0.01
1.00
0:39
1.00
0:55
0.014
0.058
0.038
0.014
0.90
0.31
0.12
0.17
0.08
6.12
0.72
5u.: 0.599
1,00
1.00
2.00
1.00
6.017
0.016
0.17
0.816
0.54
0.31
0.2
0.08
0.15
0.69
0.66
wi. decay 0.3-
docay 0.2-
1.00
1.00
1.00
1.00
0.818
0.820
0.018
0.820
0.91
0.80
0.55
6.20
0.25
0.80
2..5 255,0.444
Aud:
8; 0.39..2
1,00
1.00
1.00
0.016
0.16
0.016
0.59
0.34
0.09
0.66
A.A:
dresout 0.2-
drcoout 0.3-
1.00
1.80
0.37
1.50
0.030
0.836
0.030
0.036
0.85
0.86
0.30
0.37
0.52
8.43
0.95
0.34
21.5 230 2.300.933
(4L 19249 MLPI
1.00
1.00
1.00
1.00
0.018
0.013
0.013
0.015
0.57
0.04
0.56
0.22
0.11
0.72
(4L 1924d 11.1
dressut
0.84
1.40
1.00
0.91
0.013
0.462
0.052
0.013
0.04
6.65
0.56
0.57
6.70
0.89
0.190
1.00
0:37
1.00
6.013
0.013
0.11
0.013
0.79
0.5
0.28
0.95
0.812
0.012
1.00
1.00
9.013
0.133
9.87
1.40
1.40
0.012
0.012
0.81
0.43
6.11
0.41
5: 1- 0.59
1.00
1.00
0.013
0.013
0.04
0.56
6.17
0.94
A.. 035.0.00
1.00
1.00
6.013
0.13
0.86
0.54
des sup 0.505
1,80
0.85
0.816
0.016
6.71
0.42
6.49
0.40
R.
1.00
1.00
9.013
0.13
9.81
calo 0.01-
0.29
0.50
0.919
0.019
0.44
8.85
0.99
2v.: -0 96 0.99
1.00
1.00
0.013
0.013
0.87
0.62
6.17
0.94
295.0.933
1.00
1.00
6.012
0.11
0.49
v.. Jecay 5.01-
1.00
1.00
0.011
0.011
0.96
0.37
9.15
0.51
R. 8 0.99.
1.00
0.014
0.14
A1.: 2050 0.99
1.00
1.00
0.014
0.814
0.51
0.53
0.12
0.93
decay 0.3-
1,00
1.80
0.821
0.821
0.05
0.73
5.61
0.35
1,00
1.00
6.013
0.013
0.60
032
dro 0.1-
1.00
1.40
0.817
0.017
0.86
0.32
0.27
(51 EMA parameters
If drop U.S-
1.00
0.76
0.825
0.525
0.80
0.56
5.26
0.4-
0.39
0.72
0.826
0.626
0.72
0.47
0:23
(MLP 2MAI
1.00
1.40
0.016
0.016
1.00
0.43
0.06
0.07
0.-
0.87
0.80
0.829
0.629
0.87
0.66
arbb drop 0.1-
0.87
0.82
0.820
0.020
0.60
0.53
0.13
1.00
1.50
0.016
0.016
1.00
0.03
6.06
0.25
Fres 10-
1.00
1.50
0.018
0.018
0.99
0.09
9.17
6.41
erb. crop
0.52
0.33
0.682
0.052
0.35
0.28
0.10
arb. 0.4-
0.42
0.20
0.095
0.305
0.14
0.08
Ires 50-
0.55
0.41
0.089
0.221
6.26
0.09
9.95
(4)
Varlance
reduction
bun-n
0.018
0.99
0.85
9.06
0.12
Jew.
Jne
Ion
fur
Hro
Auy
bumin 50
1.00
1.60
0.016
0.016
1.00
0.81
0.07
0.12
(4. 10208 HLP).
1:80
1.60
0.016
0.016
6.85
0.39
6.14
0
(r
1.80
1.00
0.010
0.018
0.91
0.30
0.08
0.49
tur 1000-
1.00
1.00
0.016
0.815
1.00
0.85
6.07
0.12
gl decag-
1.00
0.015
zum-n 20000-
1.00
0.016
0.016
of in rreeayy
1.50
1.40
0.815
0.015
0.08
0.16
6.08
0.56
1.00
0.018
0.99
0.88
decay-
no bubuinn in
1.00
1.00
0.016
0.015
0.99
0.04
0.15
0.26
of teray
1.80
1.50
0.815
0.016
0.99
0.88
6.12
10244
2x butch slue-
1.00
0.97
0.825
0.628
0.92
0.16
0.99
(TT. EMAI
1.00
1.80
0.013
0.013
0.90
0.10
0.07
froq
0.07
16% butch slee-
1.81
1.00
0.919
0.922
0.96
0.13
9.05
0.34
MLP- EMAA
hrep 20-
0.99
0.59
0.025
0.823
0.84
0.35
0.63
0.96
(424424 1T.1-
1,80
1.00
0.013
0.013
0.84
0.36
0.17
0.35
tree 20
decay-
1.00
1.00
0.018
0.016
0.81
0.42
9.16
0.98
fres, 20
0-42
0.20
0.007
0.272
6.21
0.05
0.01
secay-
fre. 100
0.31
0.21
0.091
0.713
6.22
0.85
0.01
of dusay-
1.80
1.00
0.812
0.012
0.35
0.44
6.25
0.71
(2
1.00
1.00
0.812
0.95
0.49
0.17
0.42
bumin 50
1.00
1.60
0.013
0.013
6.88
0.88
6.06
0.57
o' decay-
1,80
1,80
0.812
0.012
0.97
0.31
0.12
0.39
buminn 100-
1.00
1.90
0.915
0,915
9.95
0.07
9.07
0.57
2x batch stze-
0.99
0.76
0.824
0.032
0.80
0.58
8.09
sumni 20000-
1.00
1.90
0.015
0.815
0.55
0.05
0.06
0.37
10 batch stoe
1,80
0.57
0.815
0.628
0.92
0.37
8.06
0.38
TV. EMA
1.50
1.40
0.813
0.012
0.88
0.10
8.07
0.60
0.015
0.015
Figure 9: Groups (2) through (5): Effects of various algorithmic choices: regularization, optimizer hyperparameters, variance reduction, and EMA hyperparameters. (2) Standard regularizers do not reduce the oscillations (and often worsen them) (R2, R3). (3) Large momentum has a stabilizing effect on the rewards, presumably via variance reduction (R4). However, it can result in other (well- known) optimization instabilities, leading to non-convergence. (4) Variance reduction techniques do work (R4); among these. EMA is the most robust (see especially the range,uid column, measuring oscillation amplitude). (5) EMA tends to fail if updated too infrequently; an initial burn-in period is helpful for excluding poor initial iterates. EMA with suboptimal annealing (a, so that g(()) = t (-9) can delay learning (see fearly column), but still works.
28

---

Published as a conference paper at ICLR 2024
(da) Oetaset size & decay ablations, no EMA
(0b)
EMA
Jeur
buy
MLP)
1,00
0.016
0.13
0.00
1.00
1,90
0.016
0.026
1.00
0.08
10 engeecionns
0.37
0.082
0.085
0.30
0.08
0.13
0.32
0.080
0.035
0.31
0.08
20 Try actioiisss
0.52
0.47
6.060
0.062
0.38
0.18
0.57
0.54
0.44
0.060
0.062
0.42
0.14
0.78
6.72
0.49
0.63
0.58
0.79
100 actooies
0.85
0.80
0.030
0.030
0.65
0.51
0.82
0.94
0.03
0.030
0.030
0.79
0.27
0.99
0.96
0.16
500 Iajestoris.
1.00
1.00
0.018
0.018
0.92
0.26
0.25
0.83
1.00
1.00
0.018
0.018
0.99
0.05
0.10
0.32
20 me). const 7-
0.44
0.27
0.061
0.078
0.32
0.15
0.17
0.45
0.36
0.069
0.070
0.34
0.10
0.08
100 mg cond
0.73
0.66
8.636
0.041
8.48
0.43
0.37
0.99
0.79
0.58
0.033
0.635
0.68
0.26
0.10
const
0.051
500 1R) cona,
0.99
0.89
0.633
0.036
0.77
0.61
0.28
1.00
0.98
0.023
0.023
0.98
0.09
0.10
0.30
10 Uojectoris
0.33
0.28
0.091
0.090
0.29
0.06
0.05
0.37
0.26
0.077
0.090
0.29
0.06
0.05
20 ecionsa
0.41
0.33
5.050
0.063
0.32
0.15
0.39
0.46
0.31
0.056
0.063
0.33
0.10
0.07
tr -
0.62
0.011
0.21
0.48
0.011
0.48
0.15
100 ecionia
0.89
0.629
0.029
0.58
0.52
0.30
0.31
0.028
0.028
0.74
0.10
200 tn a
0.97
0.621
0.56
1.00
0.93
0.00
0.43
500 traestories
1.00
1.00
0.615
0.015
0.85
0.57
0.22
0.92
1.00
1.00
0.015
0.015
0.98
0.07
0.12
0.47
10
20 mm) const
0.39
0.32
8.060
0.068
6.31
0.12
0.16
0.44
0.32
0.056
0.057
0.31
0.09
0.06
50
100 ma. const
0.77
0.51
8.035
0.040
0.45
0.45
0.18
0.36
0.61
0.030
0.033
0.65
0.27
0.12
500 mm)
cond
0.03
0.60
8.630
0.033
0.68
0.62
0.22
1.00
0.87
0.021
0.622
0.94
0.13
0.17
0:48
Figure 10: Groups (6a) and (6b): Ablations of dataset size and learning rate decay. When learning from fewer trajectories, optimization instabilities can become even worse. In such regimes, the stabilizing effect of EMA results in end-to-end improvements in sample efficiency. Furthermore, EMA can replace learning rate decay. This refines (R4), showing that mitigating the non-statistical phenomenon of GVA can manifest data-efficiency in practice.
(7a) Other environments, no EMA
(7b) + EMA
Jun
Jea
Juer
Ino
Hezz
ranglowo
Cury
Guace
Jow
Jue
Tea
rangeez
Tuar
Seume
Hopper: MU.
1.00
0.018
0.018
0.87
0.37
0.28
0.98
1.00
0.97
0.018
0.015
0.12
0.15
0.98
Transfomar-
0.88
0.85
0.018
0.019
0.65
0.32
0.95
1.00
0.89
0.85
0.018
0.015
0.65
0.10
0.96
0.99
decay-
0.95
0.042
0.026
0.035
0.026
0.01
HalfCheetah: MLP-
0.67
0.60
0.011
0.011
0.39
0.29
0.97
0.70
0.54
0.011
0.011
0.42
0.25
0.88
Ttaraformer
0.62
0.013
0.013
0.39
0.30
0.97
0.011
0.011
0.34
0.23
0.99
MLP
0.36
Tt. no decay.
0.51
0.32
0.029
0.037
0.34
0.27
0.08
0.48
0.26
0.023
0.628
0.30
0.20
0.11
Ant: MLP
1.04
0.97
0.051
0.052
1.00
0.15
0.03
1.04
1.00
0.051
0.052
1.00
0.11
0.02
0.96
Tranffrmar-
MLP no decay.
1.04
0.97
0.075
0.083
0.99
0.14
0.02
1.04
1.01
0.059
0.060
1.00
9.11
0.02
0.98
decay
1.04
1.04
0.072
0.076
1.00
0.12
0.04
0.38
1.04
1.03
0.050
0.061
0.99
0.13
0.02
0.09
Humanoidt
Twenformer
1.01
1.00
0.018
0.99
0.03
0.01
0.00
1.00
0.99
0.018
0.018
0.02
0.01
0.00
MLP no decay 1.01
0.99
0.060
0.069
0.99
0.05
0.01
0.00
1.00
0.99
0.043
0.043
0.99
0.83
0.01
0.00
decay
0.051
0.01
Figure 11: Groups (7a) and (7b): Evaluation of non-EMA and EMA cloned agents on additional tasks, showing broader scope for (R4) than just the Walker2d environment. Improvements of an identical nature are observed for Hopper. Sometimes, behavior cloning does not fully work (HalfCheetaa) fhhe or the reward oscillations are more benign (Ant, Humanoid), but EMA stabi- lizes training in all of these circumstances, never resulting in a substantially worse cloned policy.
29

---

Published as a conference paper at ICLR 2024
Reward vs. BC error for all runs
1.0
expert policy
0.9
x MLP
x MLP + EMA
1.000
*
0.8
X Transformer
x
Roward
0.71
X Transformer EMA
0.998
0.6
0.996
0.5
0.4
X
0.3
0.992
X
X
0.2
x
0.990
0.00
0.02 0.04
0.06 0.08 0.10 0.12 0.14 0.16 0.18 0.20
0.00
0.01
0.02
0.03
0.04
BC error
Figure 12: Do BC agents succeed by perfectly cloning the expert? Global seatter plot of final- iterate rollout reward Jn vs. 1-step prediction error bBC for all kalker2d agents trained in this paper. The plot on the right is zoomed-in region (shown on the left). From (1) the lack of conver- gence to perfect agreement with the expert policy (including with access to infinite data) and (2) the lack of a strong correlation between the surrogate objective BC and true objective JH in the regime of high performance, we illustrate the following conclusion: although statistical convergence is sufficient for behavior cloning, not necessary. Deep imitation learning succeeds in a regime that is short of perfectly recovering the demonstrator, and successes and failures are instead gov- erned by fine-grained algorithmic choices.
fractality of the reward landscape in regions where the loss is extremely smooth and near-convex.
The setup is identical that described for Figures I(c) and 6(c), except the following: models are 4-layer dimensinns [MLPs. Transformers] [with, without EMA; checkpoint indices {1, 2 denote training steps 50000 and 205000; step size increments are varied by 3 x (10-3, 10-1, 10-5 (labeled as neighborhood sizes (10-2, 10-3,10-1)
Mon-Markoyianity of the Transformer cloner. Of interest (but somewhat orthogonal to the scope of this paper), our behavior-cloned Transformer agents are reparameterizations of the original MLP expert policies. They non-Mate non-Markov, in the sense that their architecture allows their actions to de- pend on a sliding window of state history, rather than only the most recent state. A natural question is whether the internal representations of Transformer agents correspond to the formation of mean- ingful non-Markov circuits (Elhage et al., 2021; Edelman et al., 2022; Liu et al., 2023b). Preliminary visualizations in Figure 14 suggest that this is indeed the case: attention heads appear to attend to the past, perhaps learning internal "in-context resets" (Liu et al., 2023a). Similar findings appear in the nascent "sequence models for continuous control" literature (Janner et al., 2021: Shafiullah et al., 2022).
C.2 NATURAL LANGUAGE GENERATION
In this section, describe the autoregressive language modeling experiments in full detail. Throughout this section, we use the same standard Transformer architecture, with 12 layers, em- bedding dimension 1536, context length 1024, and sinusoidal position encodings. We use BPE tokenizers (Sennrich et al., 2015) with vocabulary size 2000 for TinyStories (version 2, with some postprocessing to remove Unicode glitches; 598M tokens), and 32000 for English Wikipedia (5.3B tokens, with a random 99:1 train-validation split). We emphasize that the green lines in the bottom row are seatter plots.
The purpose of using the I.LM-synthesized TinyStories corpus (Eldan & Li, 2023), in accordance with the authors' intent, to provide a significantly cheaper proxy for all algorithmic considerations
landscape looks smooth, the landscape of the long-horizon rollout rewards can appear to be fractal in the same regions.
30

---

Published as a conference paper at ICLR 2024
Normalized reward landscapes
Meili:
T97:
Ti 119
T82: 9-10"
10"
T82:
T 1
Tedllgg 16
Teg1:910"
TH9 9-30-
1181:
1- 110
Teg2: 9-30"
Te22:910
Te07: 9
Teg2: #10
Te92 F-10
Te
Lass landscapes
MOT: 9-10
M11: 915-2
Mg1:-110
11
M87: p1p-1
Mg2: 9100
M82: r10-
ME2:
D0
Tatek
610
O
CELL 1
ToO 1
TecLir 10
9
TeeL: 010
0og 10
Figure 13: Additional visualizations of reward JH(T)) (Te) and loss (BO(Te) landscapes around neigh- borhoods of policies. (M, T) = (MLP, Transformer): e EMA; [1,2] = /early, late) checkpoints; [g. r) (gradient, random] directions. The number (e.g. 10-3 denotes neighborhood radius.
Attention weights (non-EMA model)
Attention weights (EMA model)
L1,
L1,
L1,
L1,
L2,
L2,
Figure 14: Heatmaps of attention weights of 4-yyer, 4-head, 1024 424--dimnsiooal Transformer cloners, evaluated on their own rollouts (evaluated step 200. where a non-Markov policy has access the past 32 states in its context). Cloned agents seem to learn policies which compute sliding-window filters in the first layer, and attend history in a systematic manner (which we not attempt to decode). This serves impossible for the MLPebaaed models to share. Weights are taken from the final checkpoint, and both agents incur normalized rewards of 0.99.
cursory check that the policies indeed non-Markov, property that is
31

---

Published as a conference paper at ICLR 2024
Every-iterate evaluations
0.9170
LM 1055
0.9830
0.9165
1.104
1.302
0.9620
0.9060
EMA LM loss
0344.0
0.9030
1.043
0.94350
1.042
0.44255
Switch fraction
1.00
1.00
0.75
0.75
0.75
original keerates
0.50
0.25
0.25
0.25
5000
5020
5060
5000
5080
5100
10000
10026
10040
10960
10080
10100
15000
15820
15940
15000
15080
15100
Training steps
Figure 15: Zoomed-in every-iterate evaluations from the 270M-parameter TinyStories training pro- cedure, enlarged from Figure 4 (top left) in the main paper. These evaluation phases occur for 100 iterations each, starting from checkpoints 5000, 10000, and 15000 (the total number steps is 18062). The first two rows show minuscule fluctuations in the "behavior cloning" (next-token- prediction) loss. The third row measures rollout disagreement between consecutive gradient iter- ates: the fraction of input sequences where the deterministic (argmax) autoregressive generations checkpoints steps and t 1. As the losses fluctuate, pairs of consecutive checkpoints disagree on >50% of autoregressively-eeeett gegerrtte deterministic completions. This is only for <5% of EMA iterates 0
relevant to the large language model pretraining pipeline. This mitigates the common occurrence of undertraining when benchmarking LLM training pipelines, while allowing for carefully controlled experiments at reasonable computational cost. Nonetheless, the fact that TinyStories is a proxy distribution comprises a limitation of this empirical study. We hope address this gap future work (and also encourage interested parties to incorporate EMA without learning rate decay into their LM pipelines).
Each model is trained on single node with 8 NVIDIA V100 GPUs, with global batch size of 64 (thus, 65536 tokens), using the Adam (Loshchilov & Hutter, 2017) optimizer, learning rates 8) 10-1, and weight decay 0.1. For the 2-epoch TinyStories runs, finer-grained grid of learning rates used: 17 = 3, 3, 4, 5, 7, 8. x 10 4, For Wikipedia, we only use x 10 (as a quick preliminary check that qualitative findings remain natural data). 1000 steps of learning rate warmup are used throughout these experiments. With this hardware setup, the TinyStories training runs take 4 hours per epoch, while single-epoch Wikipedia training run takes 50 hours.
We restate the empirical finding from the main paper, then discuss each aspect in depth.
(R5) Autoregressive LMs exhibit significant rollout oscillations throughout training. EMA sta- bilizes the trajectory, accelerates training, and improves generalization, complementing (and potentially obviating) standard practices in learning rate anncaling.
Small fluctuations in the "BC" loss. In the training run with 2 epochs (18K steps) learn- ing rate decay, and learning rate 3 10-4 we perform zoomed-in evaluations every iterate t (5001,... ... 5099} {10001, ... 10099} U (15001, 15099). Figure 4 (top and middle terms of the validation log loss. Strikingly, upon
rows) shows that there are small fluctuations performing the same evaluations the EMA iterates, progress is extremely smooth in time (nearly indistinguishable from linear). Note that this is not to illustrate GVA- this setting, the next-token- prediction log loss is BCC Indeed, these fluctuations minuscule.
32

---

Published as a conference paper at ICLR 2024
Stories parameter Transformer
Wikipedia: parameter Transformer
1.00
EMA warse EMA better
2.45
EMA EMA
EMA non-EMA
2 epochs, decay
epoch, decay Ir
EMA worse
4 epochs, decay
epoch, const Ir
EMA beiteer
0,95
8 epochs, decay
2.40
2 epochs, const ir
EMA validation loss
4 epochs, const
0.90
8 epochs, const ir
EMA validation
2.25
epochs, opachs,
16 opochs, 0.5x
0.75
0.80
0.85
0,90
0.55
1.00
2.20
2.25
2.30
2.35
2.40
2.45
dalidation loss
Validation loss
(a)
(b)
Figure 16: Training paths of autoregressive LMs in (model loss, EMA loss) space, for (a) TinyStories and (b) English Wikipedia language models. In both cases, we find that with a suitable 3 schedule (1000-iteration burn-in and 7 t 0.67 annealing). EMA never hurts validation perplexity. With linear learning rate annealing (down to 0), the optimizer and EMA iterates arrive at the same final validation loss. (In a limited other hand, training at a constant learning rate, EMA iterates are better, even though the optimizer iterates are worse. Training pipelines which decouple optimization from stabilization can reduce overfitting and produce lower-loss models.
of experiments, findings are identical with cosine decay). On the
Large rollout oscillations due to GVA. Figure (bottom) provides a quick check that while local oscillations in the loss are small, the generations of these models diverge over long-horizon autoregressive rollouts. Figure 4 (third row) makes this quantitative: for 1000 examples from the TinyStories validation set, we take prefixes, of length (10%, 25%, 50%, 75% ), and query each of the 300 model checkpoints for their argmax completion. For each checkpoint in the intervals above, record switching cost: the fraction of prefixes ("in-distribution prompts") for which generations from 0'-1 and O¹ differ; these are the orange and yellow dots per iterate. The mean switching costs on the respective intervals are (89.2% 76.5%, 53.0%) for the optimizer iterates, and [4.9%, 2.3%, 1.1%) for the EMA iterates. In short, these numbers can be interpreted probabilities that the greedy step-wise MLE sequence consecutive LM training iterates disagree over long rollouts, conditioned on in-distribution prefixes.
Improvements in validation loss. In the continuous control settings, we do not pay much atten- tion analyzing the effects of trajectory smoothing on the behavior cloning loss, because obtaining a high-reward policy is the clear primary objective. In language generation, {BC has a special role as the main quantity interest in text compression (as opposed to closed-loop generation): namely, it is the statistical language model's log-perplexity. Figure 16a shows that on TinyStories, in terms of validation set perplexity, the EMA iterates are always least as good as the optimizer iterates. Furthermore, beyond 4 training epochs, when models overfit when the learning rate is decayed to O. without decay, models do not overfit and keep improving, even up to training epochs. We find that a partial annealing strategy (decaying linearly from 8 10-4 to 4 x 10-4) produces the best of all models, with a token perplexity of 2.21 (compared to 2.31 without EMA). While the margin of overall improvement for epoch over the larger Wikipedia dataset (Figure 16b) is smaller (9.54 9.39 token perplexity), the finding that EMA never hurts is upheld.
Examples of divergent generations. Below, we provide some examples of argmax completions from these TinyStories language models. This serves as a quick check that (1) the autoregressive generations produced by these language models are qualitatively fluent in a restricted linguistic domain, and (2) minibatch inise-induced fluctuations induce distinct autoregressive generations. The robust benchmarking of end-to-end generation quality presents methodological ambiguities; resolving them is outside the scope of this paper.
33

---

Published as a conference paper at ICLR 2024
Prompt: Once upon time, in warm and sunny place, there was pit. A little boy
Continuation from data: named Tom liked to play near the pit. One day. Tom lost red ball. He was very sad. Tom asked his friend, Sam, help him search for the ball. They looked high and low, but they could not find the ball. Tom said, "I think my ball fell into the pit." Sam and Tom went close to the pit. They were scared, but they wanted to find the red ball. They looked into pit. but was dark see. Tom said, "We must go in and search for my ball." They went into the was dark and scary. They could not find the ball. They tried to get out, but the pit was too deep. Tom and Sam were stuck in the pit. They called help, but no one could hear them. They were
search. it
and scared, and they never got out of the pit.
All models agree: named Tim and his dog, Spot, liked to play near the pit. They would run, jump, and laugh all day long. One day, Tim and Spot were playing near the pit when they saw a big ball. [...]
15027-15035: [..] They both wanted to play with the ball, but they clidn't want to share. Tim said, "I want the ball!" Spot barked, 'No, want the ball!" They started fight over ball. can both play with the ball if you share." Tim and Spot stopped lighling and started to play logether. They had so much lun playing with the ball near the warm pit. From that day on, they always shared their toys and played happily together.
mom saw them fighting and said, "Stop fighting!
15036-15037: [..] They both wanted to play with it. Tim said, "Let's share the ball, Spot!" Spot barked happily, and they both played with the ball logether. As they played, little girl named came by. saw Tim Spot sharing the ball wanted to join them. Tim said. "Yes, you can play with us!" So, Tim. Spot, and Sue played together near the warm pit, and they all had a great time. They learned that sharing is fun and makes everyone happy.
15038-15033 [..] They both wanted to play with It. Tim said, "I want the ball!" Spot barked and wagged his tail. They both ran to the ball and started to play. As they played, the went into the pit. Tim and Spot looked at each other, not knowing what do. Then, a friendly frog jumped out of the with the ball in its mouth. The frog said, "Here your ball, Tim and Spot! saw you playingg and wanted to join in too!" Tim and Spot were so happy and thanked the frog. They all played together and had great day.
15040-15049: [..]] They both wanted to play with the ball, but they didn''t want share. Tim said, "I want the ball!" Spot barked, "No, want the ball!" They started to fight over the ball. As they fought, the ball rolled into the pit. Tim and Spot stopped lighting and looked at the pit. They were sad because they couldn't play with the ball anymore. They learned that fighting was not good and they should share their toys.
15000-15099 (EMA): [...] Tim said, "Spot, let's play with the ball!" They played with the ball, but was too big for them. They tried push but would not move. Then, Tim had an idea. He said, "Spot, let's reverse and push the ball back where we found it." They pushed the ball back and it started to roll. They were so happy and played with the ball all day long.
Figure 17: Argmax-decodee generations from a sequence of consecutive training checkpoints (15027-15049), seeded by prefix of examples from the TinyStories dataset. These are the full generations corresponding to Figure 4 (bottom). Note that we do not attempt to evaluate generation quality systematically in this work; we only note that (1) the argmax generations oscillate between semantically distinct modes, and (2) the EMA iterates, aside from having better losses, switch their rollout argmax trajectories (conditioned on in-distribution prompt prefixes) far less frequently.
34

---

Published as a conference paper at ICLR 2024
We close with few remarks larger-saale experiments, in light of the flurry of interest in under- standing and improving the training of large language models (LLMs).
Prior work. Our preliminary NLP experiments are certainly not the first to note the benefits of averaging language models (Kaddour, 2022: Wortsman et al., 2022; Sanyal et al., 2023; Sandler et al., 2023). In contradistinction to these works, work contributes set of controlled experiments to isolate the phenomenon of GVA.
Iterate averaging at the LLM scale. Towards understanding considerations which may appear at the scale of our study, it is difficult to perform completely analogous experiments with open-source pretrained models. Model releases which target scientific analysis (e.g. Sellam et al. (2021); Biderman et al. (2023)) only publish a small number of intermediate checkpoints (if at all), so that frequent EMA cannot be emulated from these artifacts." 1 Thus, we encourage interested parties to explore trajectory stabilization at the scale of frontier LLMs, and publish both EMA-filtered and frequently-saved checkpoints for evaluation and analysis.
Preliminary recommendations. We suggest updating the EMA at every iteration, setting 7= 10-4 (3 0.9999), employing burn-in that is roughly the same length as the learning rate warmup, and tuning the annealing parameter in the range 0.5 S S (larger a acts like smaller learning rate). We also note (based on limited, informal tests) that the stabilizing benefits of EMA appear to remain when finetuning pretrained LLMs; in this use case, we recommend setting burn-in 0, and carefully tuning the annealing parameter based on the finetuning dataset size and number of passes over the data.
C.3 LINEAR SYSTEMS
In order to generate intuition about the cause and mitigation of GVA, we study a simple linear system as well as a simple 2-piece affine system. We recall the setting of the Linear Quadratic Regulator (LQR) (Kalman, 1960). Let x, E gd* and u, E Rdu be state and control vectors and suppose that
X++1 Ax, + Bu, + We:
W. "0,N(0,
(C.1)
where A E xdx and B E Rdu For horizon H, and policy TK X Kx with K e Rddudd the reward given by
H --1
H-1
(H( =
r r(x,, uz)
= Σ x (IA+
(Q KR) (A BK) x(C.2)
r(X,, u,) -x, Qx, - u, Ru,
where Q, R are positive semi-definite matrices.
Remark C.1. In most treatments of LQR, JH is taken be positive and thus corresponds control cost rather than a reward. In order to remain consistent with the rest of the paper, we take Ju to be negative and thus reward.
In the infinite horizon case, it is known that an optimal policy K* can be found by solving the Discrete Algebraic Riccati Equation (DARE) (Kalman, 1960):
K (R B T (BSA),
S = Q + A A SA - A T SB (R+B BTB T STS..
In this section, when the system matrices (A, B, Q, R) are clear from context, reserve K* for this optimal policy. The closed-loop system is called marginally unstable if | + BK' llap = 1 and marginally stable if |A + BK* llop < 1. Stability is helpful condition for analysis as it implies that trajectories converge toward zero and thus we do not expect GVA to appear. Marginal stability does not cause convergence, but we show in Appendix E.2. when the imitator is sufficiently close to the Ricatti policy, we also do not expect to see GVA. In both experiments, we follow a similar pipeline as in the MuloCo experiments as described Appendix C.I, with the main difference that we hard-code our expert. We now describe our two experiments in detail. In our experiments, we let dx = d for ease of visualization of trajectories.
Nor would it be practical store every training iterate.
35

---

Published as a conference paper at ICLR 2024
0.1
0.1
800
0.0
Time step
0.0
600
X -0.1
400
-0.1
0.2
200
0.3
0.2
0.0
0.1
0.2
0.3
0.4
0.2
0.0
0.2
X1
(a)
(b)
Figure 18: Depiction of marginally stable LQR setting. In (a) we show a contour plot of the reward (for fixed action u). which simply a quadratic and in (b) show sample expert trajectory in state space. Each point is the state at a given time step, with the color corresponding to the time.
C.4 LQR WITH MARGINALLY STABLE DYNAMICS
In our first experiment, we let
1.0025
0.0043
-0.0026
A
Q Q=R=l
1.0025
0.0026
-0.0043
To generate these matrices, we chose a small a > 0 and let e H We then sampled a rotation O uniformly at random from the orthogonal group. We let A (1 + e)] and B -00. Note that this system is clearly marginally stable. The rotation is included to make the learning problem slightly more challenging for the MLP. We let K* be the Riccati policy of this marginally stable system,
1.3867
0.8250
K
0.8250
-1.3867
Steps (2-5) the pipeline in Appendix C.1 are identical to the MuJoCo experiments. A visualization of the reward landscape, the expert policy, and sample expert trajectory be found in Figure 18. As demonstrated theoretically in Appendix E.2, we expect little to no oscillation in this setting, and indeed we observe this. In particular, we consider imitator function classes:
Linear. This is simply linear regression with dependent features, where optimize the class of functions (x + Kx E Rddxddu)
MLP. We also investigate the effect of introducing nonconvexity into the optimization by optimizing over depth neural networks with the dimension of the hidden layer being 32. For optimization, we use Adam with learning rate of ,0003 and default parameters. We also use linear decay with a warmup of 50 steps.
We present the result of this experiment in Figure 19. As expected, do not see GVA in this setting; indeed, both the linear imitator and the MLP are able to recover the expert to sufficiently high degree of accuracy so as to obviate the lack of stability, as predicted by Proposition D.2. We also exhibit a scatter plot comparing {BC and J in Figure 20 to emphasize the lack of GVA in this setting.
C.5 LQR WITH CLIFF
Unsurprisingly, we did not see GVA in the previous experiment. Indeed, we predict GVA to oceur when the rollout dynamics function is highly unstable and possibly discontinuous. To demonstrate that this is indeed the case even with extremely simple dynamics, we consider a system motivated by spring falling off cliff. Again suppose that dx d,, 2 and think of x1 position
36

---

Published as a conference paper at ICLR 2024
LQR: linear behavior cloner
LQR: MLP behavior cloner
validation error LOR reward
2500
No EMA
-500
EMA
-7500
-10000
0.010
No EMA
EMA
0.000
5000
10000
15000
20000
5000
10000
15000
20000
Training steps
Training steps
Figure 19: Training curves imitator marginally stable LQR setting. We show both linear imitators (left) and MLP imitators (right), comparing the reward curves (top) with the LBC curves on a validation set (bottom). We show the results both of EMA and of no EMA.
LQR experiments: rollout reward
Cliff LQR, linear cloner
LQR, linear cloner
LQR, MLP cloner
1000
800
Reward
600
No EMA
400
EMA
200
0.01
0.02
0.03
0.04
0.05
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
BC error 1e-3
1e-3
Figure 20: Scatter plots comparing {BC to the reward J at different checkpoints of single in the LQR experiments. We show the results both of EMA and no EMA. We exhibit linear imitator of LQR with cliff setting (left) well as a linear (center) and MLP (right) imitator of the marginally stable LQR setting.
0.1
Cliff
Cliff
0.0
0.0
1im stop
X -.1
-0.1
-0.2
0.0
0.1
0.2
0.3
0.4
0.0
0.2
0.4
X[
Xi
(a)
(b)
Figure 21: Depiction the LQR with a cliff setting. (a) we show contour plot the reward
landscape; note that the reward is negative infinity to the left of the cliff, marked in red. In (b) show an example expert trajectory in state space. Each point is the state at given time step, with the color corresponding to the time. Also marked is the cliff in red.
37

---

Published as a conference paper at ICLR 2024
LQR with cliff: linear behavior cloner
Mean episode length
800
600
No EMA
400
EMA
200
BC validation error
0.04
0.03
EMA
0.02
EMA
0.01
0.00
5000
10000
15000
20000
Training steps
Figure 22: Training curves of linear imitator in LQR with cliff. We show scatter plots of both the reward eurves (top) and the LBc curves on a validation set (bottom). We show the results both of EMA and of no EMA.
coordinate in one-dimensional space and velocity coordinate, We let
0 A exp 17 [ ! ))
B
Q R
for a time parameter 77 that set to be 0.1 in experiments. As it stands, this system acts as a discrete time approximation of a spring oscillating in one dimension, where the control is such that the learner can only affect the velocity directly (perhaps by applying some small force). To this system, we introduce a 'cliff" parameter x < O such that if x1 < K, then the agent 'falls off a cliff" and the episode ends. We modify the reward function by letting
{
x >
r(x, u)
-0 otherwise.
in our experiments, we let к 0.05. For visualization of the reward, see Figure 21 (a). We again use the Ricatti policy K* (with Q = R = I) as the expert. Note that this policy does not take into account the infinite negative reward incurred by falling off the cliff and thus is not necessarily optimal in the above setting.
Here we only train a linear policy and examine the effect that EMA has on training. We use large, constant learning rate of 0.3 and apply EMA post hoc filter, making this direct MDP analogue of the situation considered in Proposition D.5 in Appendix D. We display the results of this exper- iment in Figure 22. Note that GVA clearly occurs in this setting, although we observe that the Furthermore, we see that EMA is able to mitigate the oscillations and allow the agent to converge an least mediocre policy. We emphasize that the expert policy here is suboptimal because the Riccati policy does not take into account the infinite negative reward incurred by falling off the cliff and thus is not surprising that the EMA'd imitator does not perform optimally. We also include scatater plot of LBC vs. in Figure 20 to emphasize the fact that, even though there remains some small oscillation in LBC due the large constant learning rate, GVA is still very much observable in the complete lack of relationship between the two quantities.
predicted by Proposition D.5,
D THEORETICAL ANALYSIS OF TOY EXAMPLES: DETAILS AND ADDITIONAL RESULTS
In this section, we provide details and frrmal statements for the theoretical vignettes in Section and Section as well additional results.
38

---

Published as a conference paper at ICLR 2024
We begin in Appendix D.I, which studies error amplification in linear dynamical systems. Proposi- tion D.I, which is a formal version of Proposition 3.1, provides a linear example extreme sensi- tivity of rollout reward to parameters despite the lack of sensitivity in the training loss. We proceed by recalling that for (marginally) stable linear systems, this is not an issue, which is in line with our empirical results in Appendix C.4. We then shift our focus to exploring the effect of EMA in several toy examples.
Appendix D.2, we introduce and analyze a variant of the cliff loss example studied em- pirically in Appendix C.5, and demonstrate in Proposition D.3 that, if iterates are distributed as Gaussians, then the cliff' loss can be characterized in terms of the BC loss. For the sake of simplicty, we assume Gaussianity in our iterates when conducting computations.
Next, in Appendix D.2.1, we consider the effect of SGD on square loss and its interaction with EMA; Proposition D.5 demonstrates that the benefits of EMA can be heavily dependent on the choice of learning rate. We then use this result to give further consequences for the cliff loss. This constitutes a formal version of Proposition 4. 1. Proposition D.7, we provide a similar analysis in the continuous time limit for a variety of more complicated learning rate schedules, with the caveat that this result is restricted to the setting where EMA is started after sufficiently large warmup period so as to saturate the population gradient.
Finally, in Appendix D.3, we consider the extent to which convex theory explains the em- pirical benefits of EMA. We begin by surveying number of alternative iterate averaging schemes in Appendix D.3.1 and demonstrate in Theorem 2 that EMA is provably subopti- mal for stochastic convex optimization when the EMA parameter is chosen with the scaling that we find necessary in practice. These results apply to the step size choices which are op- timal for the convex optimization tasks. To summarize, while the large learning rates we find empirically effective cannot be explained by the theory. conditional matically suboptimal rates, stochastic convex optimization does provide useful intuition.
using these mathe-
All proofs are deferred to Appendix E.
D.1 ERROR AMPLIFICATION IN LINEAR DYNAMICAL SYSTEMS
In this section, study the error amplification and the GVA phenomenon in linear dynamical systems. We begin with formal version of Proposition 3.1, which demonstrates that, even with a Lipschitz dynamical system, the imitation loss may have to quite small in order to guarantee that the rollout reward is close to optimal. This is the origin of GVA: small perturbations in parameters that do not appear in BC the loss landscape lead enormous changes in the rollout reward.
Proposition D.1 (GVA in linear dynamical systems). For any e >>> 0 and d E N, there exists a linear function f IRd x Rd > Rdd and linear policy no- RRd - Rdd such that the following statements hold:
(a) There exists an MDP M with deterministic transitions such that X4+1 = f(x,, un) for all t.
(b) The functions x + [(x, no- (x)), (x, u) + [(x, u), and x + To- (x) are globally 1 - e, +I+c, and e/c Lipschitz, respectively.
(c) Let the reward function r(x, u) - [|| be quadratic, so that Jir E Σ/h [|x=13 and let Be denote the set of e' e'-iipschizz linear functions A Rd IRd satisfying A() 0. If 8 = cé' e Then,
fmi (H,13) Hnb
8 8 otherwise
sup (JI(TT-)) JII(TO + A)} 2 Cd. AEB
for universal constant C O.
(d) Let imitation loss be defined
(BO(TT) = Σ [| - xg- (xx
39

---

Published as a conference paper at ICLR 2024
Then
sup (BC(TT) I H ||x||| (c')² AEB,
(e) Moreover, there exists subset B3 C Bg such that, if is sufficiently small. then for A B'g ( + A) 2 H82 and yet, (To=) JH([T= + A) I O.
Note that Proposition D.I provides an existence result, providing a simple example where GVA may occur. In effect. (d) ensures that small perturbations around the expert Te. result in small changes to the next step prediction loss BC. while (c) shows that such small changes lead to exponential blowup in the rollout reward of the policies Ju; statements (a) and (b) demonstrate that such phenomenon can even occur in otherwise well-behaved (in the sense of being Lipschitz) systems.
We complement Proposition D.I with negative example, showing that if a linear system is sufficiently stable, then GVA does not occur. For more intuition and a number of references relevant to these linear examples, recall the dynamics given in Eq. (C.1), as well as the LQR loss given in Eq. (C.2). We consider linear policies in which there is some K Rdx such that u, T (x, Kx,, which implies
Simchowitz et al. (2020); Hazan & Singh (2022). To state the result,
that the expected dynamics are given by a linear function: E [x = (A BK) Xg. We have the following result.
Proposition D.2 (GVA does not occur in sufficiently stable linear systems). Consider the linear dynamical system in Eq. (C.1) and suppose that the expert policy given by ne- (xt) = K* is such that || A + BK I 1, ie., the closed-loop dynamics are marginally stable. Suppose that K is an imitator policy trained so that | -Kllap I HIIBI.,, for c S and ||Q + KR|lop S C. Then
JH(TR) || K inf
(n( 100 (CH²² [|R|loo) ||xo1|ee.
(TK)(TT) EE E II(K K*) xxll] then
{BO(TK) S CH (|s011l + dx e²
sup K
K Il.pp
Furthermore, if there exists some such that | + BK op
S 1 8, then as soon as K
K
up S 5
2| we have
Ju(AR) [k S
Jn( < 100 (CH*+H||R),) inf
|
Note that in both cases, the rollout ervor JH does not grow significantly more quickly than LDg, in contrast to Proposition D.I, suggesting that GVA will not occur.
Note that in the data regimes we consider empirically for the LQR setting, the assumption that ||K K lop < is reasonable, and SGD will converge relatively quickly this regime (Nesterov al., 2018). Thus, for this setting, we do expect GVA to persist training, even for marginally stable systems. It follows that the conclusion of this proposition is in line with our empirical results for LQQ Appendix C.4.
D.2 BENEFITS OF EMA FOR THE CLIFF LOSS
In this section, we state formal version of Proposition 4.1, which shows that EMA reduce the variance of SGD for the cliff loss in Section 4.3. Recall that we consider a simple parameter estimation task, where we minimize the "bchavior cloning" loss
LDC 0
- 0 1ul
(D.1)
40

---

Published as a conference paper at ICLR 2024
for mean vector μ R while the "reward" function is given by a cliff loss:
J(0) 10 s=l
10-14 S
(D.2)
otherwise
where 1 2 0 fixed scale parameter. This cliff loss simplified version of the "LQR- with-a-cliff" dynamical system given in Appendix C.5, corresponding to trajectories of length H = 1, which simplifies calculations. We begin by showing that despite its simplicity, this setting is nonetheless nontrivial: when 0 is drawn from Gaussian distribution, the expected cliff loss indeed exhibits cliff-like behavior in the regime where [[] @(e²). ). To state the result, let 0 l denote the optimal parameter. Note that the Gaussian assumption is used mainly for the sake of simplicity, as it allows for exact calculations; addition in the continuous time limit of SGD as applied to {BC, the iterates approach Gaussian process (Mandt et al., 2017), fact which we use in the sequel to apply this result.
Proposition D.3. Let ~ N (p', Σ) be a Gaussian random vector in Rd. Then, there exists universal constants cı, C2. C3 0 such that the following hold. First,
C
J( B[J(O)]
whenever E[(bc(0)] 2 c
(D.3)
2
On the other hand, suppose that BB(Bc (0) S C. Then,
Cgc²
J(Ø") B[J(()) < 2E/Ebc + C. c2 exp
2E((Bc(())
In particular, e 12200 cg
then.
J(0)) - [[)( - 3E({bc(0)]-
In the sequel, we show (via Proposition D.3) how small improvements in BC loss can translate to major improvements in the cliff loss, and how EMA can induce these improvements.
We shall also show how the SGD iterates will exhibit large probabilistic fluctuations in their cliff loss, even if the expected cliff loss large. To so, we require the following.
Lemma D.4. Let 21 NV Then,
2 PlZz| S se 2 EV T
22
In particular, for 1, P[|zi at) 2 e/3.
Proof. By rescaling, we can assume σ = 1. We then bound
Plzzl $ sd== - e e-2112 du 2 20
c 2/22 212=
22 =/ -
2012
22
2T
D.2.1 ANALYSIS OF STOCHASTIC GRADIENT DESCENT ON SQUARE LOSS AND CLIFF LOSs
In this section, we study the benefits of EMA for the square-loss Le defined in (D.1) just above. This objective function is equivalent to estimation of the mean parameter μ.. We demonstrate math- ematically that in this problem, iterate averaging effectively decreases the learning rate. Because rapid O(1/t) learning decay is optimal, we find that for larger learning rates, EMA yields a ben- efit. Note that similar qualitative observation in more restricted setting was made in Sandler et al. (2023). While this setting is too simple to observe interesting fractal behavior that we observe in GVA on more sophisticated imitation learning tasks it does provide some theoretical intuition for the empirical finding that aggressively decaying the learning rate and iterate averaging result in similar qualitative behavior.
41

---

Published as a conference paper at ICLR 2024
We consider stochastic gradient updates on the fac defined in (D.1) defined as follows,
(+) μ1 + w(), 0 0, ((++1) ((0) 7:yx @ = 70 + - 20)
Ø(0) (
(D.4)
where ((() is a scaled, isotropic Gaussian and T)e 7t > 0 are the learning rate and EMA parame- ters, respectively. Because w(f) is isotropic, the problem tensorizes across coordinates. Hence, the following result considers only the dimension one case.
Proposition D.5. Consider the process in (D.4) for dimension d = 1, with constant learning rate
and EMA parameters "It = r and T) = 17. Let = |0(0) Hl), and We 1.1. ~ NN(0, o(), Then
40 - 4b2 (1 - 7))2
7 2 21,,
2EI(no(0 = - )= s 26 7)
2602n n + 32b2 (1 - ))²²
2 $ S 27
20 + 4b2 r²
n/2 (1 7)
n 2 27.
Similarly, the BC loss is lower bounded as
1
021]+ (((-
72
[E[[Bc ((() B(( 82 7)
+ 4 02ny + b² 21 2)
92?
Moreover, in the "no-EMA" setting in which = 1, we have
2E(( = T)E(((()) [0) 0² 1 ((1 -1
+ b2 1)229
To understand when this result reveals benefits of EMA, let us first consider the regime where >>> 1 1/t. Here, the b² term capturing transient dependence on initial condition can be neglected. We then observe that, as long as Z 1). EMA with parameter increases the BC loss by at most a constant factor relative to no EMA. On the other hand, when 7 < 17. the BC loss BC (0 )] for EMA scales linearly in Y, while the no-EMA BC loss E((Bc((())) is linear in n > y. In other words, when the step size is large, a small EMA parameter significantly aids variance reduction. This is intuitively clear, as smaller EMA parameters correspond to averages over longer windows.
Consequences for the cliff loss. By combining Proposition D.5 with Proposition D.3 and Lemma D.4, we now state that EMA is beneficial for the cliff loss setting considered Ap- pendix D.2 discrete-time SGD, complementing the results for continuous-tinn time SGD in the prequel. This amounts to a discrete time analogue of Proposition E.3.
Proposition D.6. Consider the cliff loss with parameter >> O. consider the iterates produced in Proposition D.5 with o² = 1 and b < 1. Then, there exists constants C1, 02, >> 0 such that, if the step size n, and EMA parameter n satisfy
1) 2 c 2 2 22 27, (1-7)2T S 7
then
(a) It holds that
J(0)) =()))))
C
yea -(0)) - E[J(O s cs (7+ Ce
2
(b) Moreover, cnn S BB (0())) S con and es < E(Erc (0,0) co7
(c) e c7 (og(C/7), then it holds that
err
C 2' yet J(0') - -(J') s 0s7.
J( - B|J(
(d) Lastly, even though ((0) - [[((T())] 2 G. we have that
PP(((') - [[ < 7) 2 097/n-
42

---

Published as a conference paper at ICLR 2024
This proposition reveals that an logarithmic difference in the magnitude of and can lead to ((C")) -iC")-differences in cliff loss. Since C' is an arbitrarily large problem parameter, and since y and 7 are proportional to the magnitudes of the respective BC losses due to Proposition D.5, this shows how minor differences in BC loss due to EMA translate into substantial differences in evaluation performance.
Despite these separations in expected rollout reward, part (d) of the proposition states that with probability /2(7/1), the SGD iterate has performance comparable to that of the EMA'd iterate. This reproduces the second feature of GVA: that test rollout performance of the SGD iterate is not uniformly suboptimal, and some SGD iterates can indeed have loss comparable that of the EMA'd weights.
D.2.2 EMA FOR CLIFF LOSS IN THE CONTINUOUS LIMIT
In what follows we analyze the effect of EMA on SGD for the cliff loss. To simplify calculations, analyze SGD in the continuous-timg limit (Mandt al., 2017: Malladi et al., 2022; Busbridge et al., 2023). In particular, we consider the limit as the learning rate tends to zero at an appropriate rate such that the iterate trajectory becomes continuous semi-martin martingaatng,,, following e.g. Busbridge et al. (2023); see Appendix B for further references that consider similar limits.
Formally, as in Section 4, we consider a fixed sequence of positive learning rates (m=) as well as an adaptive sequence of (potentially stochastic) estimates )220 C IR of the gradient of (O). The discrete time iterates (()) E Rd are defined recursively, with the initial iterate (()) fixed and subsequent iterates given by
( = O(() - 7) (0)
Our primary objective of focus is Exponential Moving Average (EMA) of the iterates 0(), defined recursively by 0 (00) and
0 (+++) (1 7e) 0 + 0(++1)
(D.5)
where (70)220 [0, 1] is given, non-adaptive sequence of parameters.
Continuous time limit of SGD. To study the continuous time limit of the process above, we take Te J. 0 and K at fixed rate such that limx-+0 Σk= Th (0, oo); concrete choices for the sequence 7/ and stochastic gradient process g()) are given in the sequel. We then abuse notation by letting g(t), 04)) 7t denote the continuous time analogues of ((t), 0(t) and %t. With this limit, Ø!!)
is the solution the following differential equation:
d (ol - x(0) dt.
(D.6)
In what follows, we always suppose that a unique solution to (D.6) exists; this is always true by the Picard-Lindelöf theorem (Lindelöf, 1894) if 7/ are uniformly bounded and both 7t. and g( are continuous in 14 In the case where Jk is stochastic gradient with mean the population gradient and finite second moment of its Euclidean norm, 1) is constant, and 84 = 33 is constant, Busbridge al. (2023, Theorem D.I) shows that as 77 4 O, with the correct scaling. 0.5) is the correct limit of the 0. Ø((+1) Similar limits with adaptive gradient updates can be found in Busbridge et al. (2023); Malladi al. (2022). In lieu of re-proving such limits here, we appeal these results and instead consider a fixed process ( and compare the behavior of the EMA process Ø to the vanilla SGD process (() in order to provide intuition as to the effect of EMA on the iterates.
Benefits of EMA for continuous-time SGD after saturation. In practice, EMA is often only applied to the tail iterates of SGD, i.c., after some amount of warmup. If the warmup is sufficiently large, then the iterates may already have approximately found population-level stationary point, and thus the observed gradients are non-zero only due to the stochasticity therein. Our most basic
"For hackground terminology and results from stochastic calculus, we refer the reader to the excellent exposition of Le Gall (2016).
43

---

Published as a conference paper at ICLR 2024
result, Proposition D.7, abstracts this limit driftless stochastic process. assume that
More precisely,
0
n.dB,,
(D.7)
where B, is the standard Brownian motion in Rd which corresonds to the limit when E [g]0]] | 0. In this simpler setting, examine the benefits of EMA with different schedules of learning rates.
Proposition D.7. Suppose that Th is chosen via one of the following schedules:
(i) π) 1j for all t (constant learning rate);
(ii) Т/t '1(1 + t) (inverse square root schedule);
(iii) T) - 17 + 1) ))² (inverse schedule);
(iv) = 1 (linear decay schedule).
Let ( be as in Eq. (D.7) and (1,0 as in Eq. (D.6) for fixed Let the behavior cloning loss {BC and the reward function J be as in Eq. (D.1) and Eq. (D.2) for parameters C > E >> and μ O. If (()) 0, then there exists an 1) > O sucht that for all sufficiently large t, there is some y > 0 such that
J(Ø*) - [ 2 yer J(O') - E J
1
C S 2e < 2
In addition. for the inverse learning rate, we may choose ((4)))==0 =
this 17 such that
max [E[C (A(+)) E [BC (44)
This shows that in the continuous-tin limit, EMA can substantially improve regret relative to vanilla SGD. While the three learning rate settings above are obviously not exhaustive, they pro- vide number of classic examples. Constant learning rates are frequently analyzed in theory due their simplicity. while square root learning rate are typically used in convex optimization to attain the optimal rate of convergence. Linear decay learning rates are common in deep learning and many of our experiments use them. Thus we see that in many natural settings, iterate averaging can help alleviate EMA. order to broaden the applicability of the result, we also consider an analogue of Proposition D.5 in continuous time, which is similar to Proposition D..7 but allows for drift and thus does not assume we begin EMA in a stationary regime. We state this result in Appendix E.5.2.
D.3 FOR OPTIMAL CONVEX LEARNING RATES, EMA DOES NOT HELP
In both Propositions E.3 and D.7 we saw the benefits of EMA in convex settings. We now con- sider several ways in which the convex theory does not predict improvement due to EMA, which complements the empirical results in Section 4. We begin by surveying number of iterate aver- aging schemes from the literature as applied to stochastic convex optimization in Appendix D.3.1. 1. Then, in Appendix D.3.2 we prove that EMA with the parameters that we find work empirically (see Appendix C) is provably suboptimal for convex optimization. This is attributed to the fact that aggressive learning rate decay far more than desirable for neural network training - optimal for convex optimization.
D.3.1 COMPARISON OF EFFECTS OF ITERATE AVERAGING SCHEMES
In Appendix D.2 we showed the potential benefits of EMA when there exists a discrepancy between the training loss and the reward function. such settings, both the first and second moments of the distance between the returned parameter and the optimum this section, we survey some prior work and observe that the benefits are often less clear when such discrepancy does not exist. To do this, compare various existing iterate averaging schemes and show that some versions of EMA are minimax optimal, while others are not. We focus our discussion strongly convex but non-smooth stochastic convex optimization. Specifically, we let O, IRd,
relevant to the expected reward. In
44

---

Published as a conference paper at ICLR 2024
and consider optimizing function F(0): Rd - IR on convex domain K. C R We assume that F is p-strongly convex and L Lipschitz w.r.t. the Euclidean norm, and that diam(XC) I D. Lastly, assume that gradients given according to the following oracle:
Definition D.I (Subgradient Oracle), We assume that there exists an oracle 9(0) and a a L which return random vector g such that (a) lg G and (b) Eg gg OF(O), where aF(-)) denotes the convex subgradient. We say ((0) is deterministic oracle if is deterministic. In this case, we can always take G =.
We consider stochastic optimization algorithms of the form
(1+1) = 11 (0(0) m):
C(0
(D.8)
0 0() (t)
7/.) 0(4)
(D.9)
and let F+ minøek F(0). We denote "/- 1, we call this final iterate gradient descent, because _h = 0, When 7 = 1. we call this full iterate averaging because 0 = E'= it Oave Finally, when
1 O,, and denote is fixed, we call this fixed emponential moving averuge, denoted @()) Lacoste-Julic et al. (2012) show that the averaging parameters
are optimal, and we denote the resulting parameter 0 after the authors. Finally, we consider scheme that cannor be represented in the form (D.9): the a suffix averaging analyized in Rakhlin et al. (201 1) and defined be
sufaa
Tat
s=t-[al]+1
known that no setting of step sizes or EMA weights, and in fact, no stochastic optimization algorithm making at most 7I' queries to the gradient oracle G(O), can do better than the following information theoretic optimal rate for
G² (TT)
E[F(Or)- F*) =
(information-theorriic optimal)
(D.10)
The following theorem summarizes the performance of these various schemes relative to information-theoretie optimal benchmark.
Theorem 1. With optimal step size, the final-iterate and suffix averaging schemes described above suffer the following suboptimal lower bounds:
In dimension at least T, exists an F which is 1-strongly convex and 3-Lipschitz, and a deter- ministic gradient oracle such that the final (unaveraged) iterate Or under update (D.8) with the standard step size np. = 1/t 1/ satisfies
F(O(T)) - F* T/To with probability one.
More generally, the lower bound on a-suffix averaging Qatta F(O F 2 2 log 2),, which is suboptimal when a = o(T). This is due to Harvey et al. (2019).
There exists an F which is Otr)-strongly convex and O(1) dient oracle with G O(1) such that the average iterate produced by grudient descent with step size 1/cl for any constant 2(1) is shown by Rakhlin et al. (2011)
and stochastic gra-
F 2vg
F*) 2 eS2(log T/T)
On the other hand, let E (0,1) be a constant bounded away from either zero or Then the a-suffix average g(t) or as well as the Lacoste-Julien average 0.¹ satisfy
max BI
suf - F B(F(O))) -
The guarantee for suffix averaging and d are due to Rakhlin et al. (2011) and Lacoste-Julien et al. (2012), respectively.
45

---

Published as a conference paper at ICLR 2024
We note that the Lacoste-Juiie Juliin scheme 0 is unique among those discussed that attains the information-thew theorretic optimal rate (D. (D.10) and can also be efficiently computed in a streaming fashion if I' is unknown in advance In the regimes of interest, however, the functions we minimizing are not necessarily continuous, let alone smooth. In the following section, we demonstrate that the smoothness assumption dropped, then EMA becomes supoptimal.
D.3.2 EMA IS SUBOPTIMAL FOR NONSMОOTИ CONVEX GD, PROVIDED THE STEP SIZE SCHEDULE OPTIMAL
Given the above discussion of EMA weighting schemes, a natural question to ask is whether optimal choice of in our empirical results is borne out by our theoretical analysis. We show that this is not the case and thus a convex mental model the benefits EMA we see empirically is insufficient, absent discrepancy between the loss Throughout, we let ((0) denote gradient oracle, and again consider updates given (D.8) and (D.9). Our main result is follows.
optimizing and the reward care about.
Theorem 2. Let E N be given and let K denote the unit radius Euclidean ball in IRT There exists 1-strongly convex, (1)--iischiizz function F(O) on K. and a deterministic gradient oracle such that, for any 3 E (0, 1) the updates (D.8) and (D.9) with either fixed = T or time varying Tt 8 suffers lower bound of
(
F(Or) F F
Moreover, there exists an O1))Liipshitz but not strongly convex function F(0) on KC such that gradient descent update with step size The V and EMA with either */t 7 T or 7t. t- and T 2 21-0 suffers
3 log VT
F(OT)= - F* 222 Ω
Finally, running EMA with step size y clog for a sufficiently small universal constant e in either example suffers F(Or)= F 2 0 ( In particular, for all fixed EMA parameters "
log
strongly convex case
F(Or)
T
Ω
log log vT
weakly convex case
We emphasize that Theorem demonstrates that in the regime that find EMA empirically works, the convex theory is lacking in explanatory power. This is due to the fact that the ideal step size choices for convex optimization decay far more rapidly than those in the non-convex optimization of neural networks. By contrast, if the step size decays more slowly than the EMA parameter, we can achieve benefits in convex problems, as shown in Proposition D.5. And is precisely the slow- or no-step size decay regimes that we empirically evaluate for the nonconvex optimization of deep neural networks.
E PROOFS
In this section provide rigorous proofs of the statements from Appendix D.
E.1 PROOF PROPOSITION D.
Proof of Proposition D.I. Consider the dynamical system A = Ia. B = cls. and expert policy parameter K L-g We let f(x, u) = Ax + Bu and o ( Kx. is clear that these satisfy the requisite Lipschitz constants. For any other policy A which is at most e' Lipschitz, we have that
|| 1x+1 || < [x,, (xx, + A(x))) II(A - BKjx, + cxxx)] < (A BK +<))x]
< (1 - e [x+ll-
15Oseerve that if T is known the learner, then one can implement tail averaging special case of full iterate averaging, but only begin the averating time step [o]]
46

---

Published as a conference paper at ICLR 2024
Defining p ; c + cé' = 1 8, we have
H
H
(T((TT)) Jr + +1)
X
S IE
Σ
2h
- 1
x11
Σ hml
p
2( 1a1--191
DE
||| 112)
A=1
H
d
2 1)
h=1
Moreover, by selecting Axx) = e'x, this upper bound is attained. As p + 8, for S 1.- it is standard bound d 1 11 h=1 0² (h- d1 + 8)2( 11-1 52 exp( and similarly,
Σ (hh-1) < Hp2(B-1) < H exp(O(H)). For ā s 0, then Σ 22 is (min }
by a standard computation. The first result follows.
For the second statement, note that
I
[[(K - K)
xi|
[1]
K - K
|xx,|
|
K
||x
S 2² |x]
op
where the second inequality follows from the fact that ||A + BK op S summing over
The result follows by
Lastly, we show the existence of A's with relatively large imitation cost but that do not suffer expo-
nential error amplification. Define An == & I. and define the set B, = (A cll - All s é /6). This implies that for every A E Ber. Define Á ( Ao). Then,
(A + B(KKAY A)) (1 (1-E é ++
Note that if ||xo|| S 1 almost surely, then by the preceding analysis, it is clear that IE and thus by construction, if we let C À A, then
||xn|| S H
H
||x = E
X
H Xo + Σ CH CCH-sw w.
|| E [||xo||] +
Σ
H-
we suppose that c < H-1, then the above (2(H); thus [C)c(K]] 2(H8²) as desired. the other hand, because ||A + B(K + A)l < ||A BK* the rollout reward of the perturbed policy is greater than that of the expert. The result follows.
llop
E.2 PROOF OF PROPOSITION D.2
The triangle inequality tells us that
A BK BK
SIA BK lop + |B[KK -R)
S1 H'
lop
by assumption. In particular, for t E [H]. we have that
A + BK BK
Se 1+
47

---

Published as a conference paper at ICLR 2024
for 1. For each t, we have
x ((++BR) (Q+ KR) (A BK) Xo
(AABB))) (Q+ (QRKR) (A + BK) X0
2 xi
-(A))-(A+BK))) (Q + KR) (A + BK)* Xn
xo
((A+BR) (kK (K-KK) K)R)
xo
(A (K-K)
x ((A+DR)) (Q + KR) ((++ (A -K)K xo
CHc |xol| ||x0| R op (1 20)²c CHo||xol1
2 100C | 1 (H +
[R| pp
The first result follows by summing over t. The second result follows similarly after observing that
||A + BKll S for all
K
K
E.3
CLIFF LOSS FOR GAUSSIAN RANDOM VECTORS (PROOF OF PROPOSITION D.3)
In this section, we prove Proposition D.3, which establishes the direct but discontinuous relationship between the the square BC loss and the cliff loss. Recall imitation losses and cliff loss for j fixed:
2 -,100 mll?, J(Ø) = -1 1ll2 C-
00 اlll S
BBC (0c
otherwise
The statement (D.3) Proposition D.3 is direct consequence of the following lemma:
Lemma E.1. Let 0 be a random Gaussian vector with covariance matrix Σ. Then there is some constant e such that for all e >> 0,
inf P (1 pll () 21
BillO 1
Proof. By the classical Hurbery-Wrighh inequality (Carbery Wright, 2001, Theorem 8), for any degree N, there is constant 0 such that for any nonnegative polynomial P(-) of degree vector- valued log concave random variable z E Rdd such that
1/0
P (P(z)) S E c
B[P(z)]
=c
E [P(z))
Noting that 0 μ μ ja + D1/22 where (O. I), that 0 - 1all2 can be expressed as as a degree 2 polynomial in the log-concave standard Gaussian variable. The result follows. O
_
To prove the second part of Proposition D.3, we recall the following corollary of the classical Hanson-Wrigee Wright inequality:
Lemma E.2. Let 0 be a Gaussian random vector with mean μí' and covariance Σ. Then there is constant such that for all t. 0, it holds that
P (110 2al 2 2- 11æ0
1'112 2(r(E) + 21) sde
-emin
2 Salp Tr
48

---

Published as a conference paper at ICLR 2024
Proof. This follows immediately from the classical Hanson- Wright inequality (Rudelson & Ver-
shynin, 2013) by considering the random vector 0 = 11 + // 7 with Z isotropic Gaussian and
applying Young's inequality to upper bound ||0 -1l2 2 1²'1 2 mll?
Translated into the language of Proposition D.3. set 2= k 11'1 Σ) +2t (B]{+c 21. Suppose that e² 2 8B](/c((O)], so that / = _ - (B]]]c(()]] 2 T Then,
(11e ull 2 2) ce
-mmin
|
Morever, we can upper bound | I tr(z)² Σ (0]) and |S S (l) S
(0) Thus, using e² 2 8E[ min
2 (0)( By reassigning || 1
ppllop)
constants, we conclude that, if 2 2 (E)))((O)),
IP (110 ul( 2 2²) S Cze
c3 c
Finally, we lower bound
E[J(O)] 2 CP (10- 11 10122 22) B[-(nc(0)] 2 (O)] C oge 20
This concludes the proof of the second statement of Proposition D.3.
E.4
PROOF OF PROPOSITION D.5
We now turn to the proof of Proposition D.5. Introduce the error 5(t) -(() - 1 and 8.¹) =64 1. We readily compute
(0)) = (1 - (L-1))((11) - 771(1-1), 8(0) (1 - 1-))8-1111 + +
Let's 's quickly compute the No-EMA setting:
8(() (1-
(1-m1) (1l-1)T-10 T-1(( +
(1-
1=0
and thus, with |5(1) = and EF [w??] = σ²,
T-1
B(( = =1 11) 2T
(1 11) 2(T 11-1)
+
T-1
= =12 1) 2T
+ 11000 0²
(1 - 11) 21
1=0
- ((1 17
=,,b(11
+ 2T
1 (1 1)
((1 27
11 1) 2T
0²
1/2 17)
= n) 2T
0²
1
((1 11) 2T)
2 -
For the EMA setting, defining z (8(t), {((t)) we have
(1)
(1
17)
( (t-1) 1)
(e-11)
7)
=A
49

---

Published as a conference paper at ICLR 2024
Let e e2 denote the canonical basis vectors for R2. Using z0)) (811), 8:00 (81))
(c) + (2)) 84 = oj (t)), and Bg = 1/c1, we compute
T-1
8
Σ
7 A
Bwze1 + +e Ae(e) + e2)
And thus, by independence of W1..... WT-1 ~0 and 8
E B((8 1 = 0² T-1
m A -2 (el + +z))²
T-1
(nA (6 [2 + A T- 22,2))
(2, 11)2 (800) (2, 1] + A 21)²
(nA (2, 11)² 22 2. 112 + 12)
(2, +1) 1|" + (1 - 7)2(2)) = V1 + 25 - (12(1-1) 222 (E.1)
where above, [[i is the i, j-the element of matrix X, and where in the last line, we use that the diagonal elements of a power of triangular matrix are the powers of its diagonals, as well as (6 = b².
To compute the powers A [2. 1] lets first assume Y # 17, and moreover, that either 7 2 or 2 2y. We can address the other cases at the end. Using a formula for diagonal matrix exponentiation with # 17, we obtain
7) (1-n)
A'
(1-)))
(1 7)-
We then have
nA'(i m (1 (1- - (1-7)²
(1 (1-1)) (1 7)
5 { 17(1 11)²
72 21
177
(E.2)
(1 11-11 11) 10 (1 11-11 7)
7-11 17
17(1 7)
17 7 27
Hence,
11 3 T T-2 (1 - 7))2
//-0 1=
2 21j
(1 ( 11)3((-11
72
V 402
+ 4b2
72 "T= (1 22
T-2
2 27
[1 - (1)--)2224 2/((-1))
2 27
1-11-p) 72 21}
2²
(1 1))
72 21
S 2
+ 412
(1
1- 7)2 n 2 27
| (1 - 2)
17 7 27
02 1 + (2(1 (T))
72 2 21)
2-
02 2 62 11 2)1
17 2 27
7 2 21) 00' + (2 - ())((--))
22n + 12 3/2 (1 2
17 2 27
Lets now hand the case 7 E [ 27). In this case, observe that the entries
(1 -1-10 n)
A
(1 117)) 7)
50

---

Published as a conference paper at ICLR 2024
non-decreasing in 7). Hence, we can upper bound
A $ (A')* entrywise,
(1 - 1/) A'
7 (1 (1-7)] n/= 7
(E.3)
We can now apply the same upper bound with parameters 7; = 7. Defining V4 as corresponding Vr but with r/' instead of Vr, we have that
Vi - 4(021// + (2 - -/)=²) 0²n 1-4b2(1 7 4
4
Moreover, its easy to see from (F.1) and E.3 that
Vr S 16V%
yielding V < 1602))+32b2(1 2)ar. We conclude with lower bound (((6,1)) Repeating the computation from (E.1), and using non-negativity of the entries of the matrices A can lower bound
T-1
Σ
1)) +
+ (1- = = VU + + - -) (E.4)
Following the computation in (E.2), we can also lower bound
7 > 17 7 - 7) 7 > 7
( - 7))
A V,
Thus,
11 T-2
11= (1 1)2t
7> 21 17 77?
( (1 - -1)
9>11
VT 2 0²
+ 63
322 - (1 7)26
(1 1 (1 7)2(0-1)
979)
n2 ((1)(11))((111)
(s(1 ((1--1)2(1-1)
7>11
ア>11
< 40
+ 1/2
211-(119) 11
2/ (1 2)
1-(1-7)
?7?
11 7>?
( o'n + (1 - 17) (((
7 721
4 2 + b² 1p2 7)
729
where in (i) we use that (1 (1- (1- x))) s continuity of V and our lower bound on it, we directly extend to the 1) 7 case. Combining with Eq.E.4.
2r for 1, and the assumption that n; 7 S 1/2. Using
1
/ 62 ( 1)
72 n 02 + 12 12 # - 2(T
62
+
0
11 11
E.5 EMA FOR CONTINUOUS GAUSSIAN PROCESSES
In this section, we prove Proposition D.7 using stochastic calculuss. We begin by stating our main computation, used in tandem with Proposition D.3 to prove Proposition D.7.
Theorem 3. Let t + "it be continuous and nonegative, and ler g(t) be (a) Gaussian process with almost surely continuous paths, (b) continuous martinaatin,,,, and (c) have pointwise finite second moments, Introduce the function
G(L) = %/¹ (s)dss
and consider the random process 0%00 defined by
d ((0) = (0 - 00(0 dt, σ ())
dt
Then,
51

---

Published as a conference paper at ICLR 2024
(a) 0 Gaussian for each l.
(b) The following identity holds:
E
O =e -G()
(--G())BB E [000
/ c(1))-(0)) E [0l5) ds
(c) It holds that
E 0 0 G0()
0 000 L/'
(00)) E (0()) @ 0 ds,
where the outer product and denotes the Loewner order.
E.5.1 ANALYSIS OF THE DRIFTLESS PROCESS
In this section, we study the variance reduction of EMA on the Gaussian drift process
( 1// nddB,
To quantify this, define
H(s) In' ngdu.
Corollary E.1. Suppose that ((0) So ndBB,, where ne is some deterministic process and B, is
Brownian motion in Rd. Then 0 O are Gaussian for all and it holds that
E
1 0 and
Cov (050) = (000000 (s)) ds I,
where the identity matrix and H()) = So d2 as above. On the other hand,
Cov (0()) H()) I.
Proof. Note that (() satisfies the conditions of Theorem Furthermore, E [0(0]] = and Cov (0(0) (Jo r/ ds) I. The result follows.
As stated above, Corollary E. toy model for the setting where EMA is only applied after the training loss has saturated, with a) denoting the continuous analogue of the learning rate in the discrete time setting of (D.5). We now consider several instantiations of this result, all assuming that g(0) is deterministic and 7/ = y constant.
Example E.1 (Constant Learning Rate). We first model the constant learning rate setting where 7, = 1) > 0 for some fixed n7. We first observe that (t) 172t and thus Cov (0(+)) 172 I. On the other hand, we have that
(1-11-100 I - 2 I Cov (8
Cov (0() 3 3n²
?
Note that as 1 0. Cov 0 tendds to zero, while as 7 1 20, the covariance tends to Cov (e(!)).
While it demonstrates that EMA can be effective in variance reduction, the first example is not very realistic, partly because constant learning rates rarely used in practice. For convex optimization, it is common to let 77, scale inversely with the square root the iteration. In this case, the separation similarly pronounced:
52

---

Published as a conference paper at ICLR 2024
Example E.2 (Inverse Square Root Learning Rate). If we let 7 n11 + +t)- as is commonly done in convex optimization, then may that (t) n/ log(1 + t) and so Cov (0())
r" log(1 + +t) I. To compute the covariance of 0 we apply Jensen's 's inequality to conclude that
L 707(6-10 log(1 + 5) ds (1- (1- S. 1026-111 log(1 8) ds
(1 e x) llog 1+ 1- S 767(5-0) ds
(t- 1 (t-1-c-7)
) llog
1 (1 e--)
Thus,
t - (---)-))
Cov ((()) イ
3 (1 - (1-e-- --)) .log 1+ (1 - e x0-
?
< 1 log(1 + +1).1 II
= Cov (0!1).
As in the previous example, while the variance of the iterate process is unbounded, the variance of the EMA process remains bounded. In our empirical experiments, however, we often consider a linear decay schedule (after warmup)) In this case, have the following:
Example E.3 (Linear Decay Learning Rate). Here we suppose that 7) = =1 - for E [0 and we compare the final iterate (()) = / different processes (indexed by t) to the EMA versions
0 0 n.7- We compute directly to get that
2/2
1 1 - 1
(x1+))) I I=
L I Cov (BlO). 2
Cov
3
2
72t
In all of our computations, we observe that EMA leads to potentially significant variance reduction, depending on the value of 7. While the previous examples are illustrative, the driftless assumption limiting.
E.5.2 ANALYSIS OF THE ORNSTEIN- CONSTEIN-UILENBECS PROCESS
We now relax the assumption that the (( have saturated in the sense that the population gradients all vanish and demonstrate that even in this relaxed setting, EMA can substantially improve the cliff reward in the continuous time limit of SGD. This analysis acts as a continuous time analogue of Proposition D.5. To proceed, we first recall that Mandt et al. (2017) shows that with the appropri- ate scaling of the learning rate, the continuous time limit of SGD for the quadratic loss LC is an Ornstein-Uhlenberr process:
d0 = - (A -4 11) dt EdB,,
(E.5)
where Σ > is positive definite and, in the context of LnC, A and E are scaled identity matrices. with the scale depending on the precise learning rate. Above μ E Id is the vector that minimizes LBC, and B, is d-dimensional Brownian motion. We now show that this limit, corresponding to a constant learning rate, EMA can offer substantial improvement for the cliff reward.
Proposition E.3. Suppose that g and d are as in (E.5) and (D.6) respectively, with the dimen- sion 1 and A = a > 0 fixed; further suppose that Σ 1. Let S ca for some small constant e> >> 0 and let BBC, J be in (D.1) and (D.2). Furthermore, suppose that ( E TR is deterministic. If ( 11 < K -1,,, then all > 0, there is some choice of y > 0 such that
J(O
(0())
2 yes J(0') -
E (
S <
C
2
In addition, max E [[nc (0(1)] ,E [enc 0. )1 < Oc)..
53

---

Published as a conference paper at ICLR 2024
The above result shows that as long as ( is not too far from μt, EMA can lead to substantial improvement the cliff reward. Conversely, using a similar computation. it is possible to show that ((() l4| > Ke for some fixed K independent of 6, then our approach cannot show similar gains.
We now characterize the covariance the EMA and non EMA'd processes under the OU process. We first recall the classical fact that ( can be expressed as
g( = ( + EI-E-A)) pe + ((AAl--) EdB,,
cAll-t)
(E.6)
where the last term is an Ito integral (Le Gall, 2016). The following corollary specializes to the univariate case; the more general case is handled in Corollary E.3 below.
Corollary E.2. Conside the univariate case of (E.5) with scalar A = a a #0 0, Σ = 1, and = =Y >
O fixed, and ay $ {1,2}. Then, it holds for almost every 7 that 0 0(0) are Gaussian for each and
E (())
= E [o(t)] + -e----
7 (0 11)
Cov (1:1) 3 e rt Cov (B81))
1 1
(6-7- -2at)
2a
2a(1 - -2a/7)
On the other hand,
E [[((t)] aat (0 (1))
11).
Cov (0(*) (00) ((()))
e -2at
2a
Note that if Cov (0()) 0, then it holds by rearranging that
( < Cov (0())
(e - c < Cov (0())
Cov
and thus EMA leads to variance reduction. We also have the following, more general computation for higher dimensional processes:
Corollary E.3. Consider (E.5). Suppose that E > is a positive definite matrices and μ E R and g satisfies (E.5) and ((9) is independent of the Brownian motion B,. Suppose that Y r fixed, and 7 A is nonsingular. Then,
E μ rr (0 (B())-1) -w) (C-TA - cc (I-7-1A) (B 11
Furthermore, ifE = oI, AAT = ATA, and A + A - AT) are nonsingular,
Cov (0))) -t Cov (D1O) + o² (1-
o² (e-(A-ATyy - (I-1 (A+AT))
On the other hand,
E [(()
1e + c e (B(w) -H)
and
Cov (0 r 0(0)000)
0² -02 (A+AT)-
Proof. We begin by observing that by (E.6), it holds that
E 0 (1)
r + At (0 14)
54

---

Published as a conference paper at ICLR 2024
Thus by Proposition E.4, we have
G(t) ll0) + oo (f cQlo)-o1tyy E [0(0) ds
- " ( e1e-1ly () + - ( -p) ds
x ( + (1 - -7) p + me Sf el( (010) - u) ds,
where the final equality follows because I commutes with A. A simple calculation then tells us that
L/ c(I-T-1 (000)-) - -1) ds = (c (71 (II-A) (010) -H)
and thus
E 0 = ne. 00()) + (1- μ + (e-rA (e-rA-e (1 (I- (010) -m).
Thus the expressions of the means of both (()) and 0 hold.
For the covariances, we first observe that by (F.6), if A commutes with A and E II, then if
A + A is invertible, we have
Cov (o0) - ( (00)
= 0². =0² f" (
J
o² (I-E-(AHAD) (A+AT)*
Now, by Proposition E.5, we have that
Cov (0())
3 7t Cov (010)
+ ( e2(0=1)y00² (1- (1-0-AAAA7) (A + ATT ds
0
= Cov (10)
+ o² (1- (A+AT)
E2 (((6-0]I-(A+AA)) ds) (A+AT)-
= Cov (B(0)
+ 0² (1 (1-6-79 (A+AA)
0² (c-(A-ATy) ( ( (A+A))
The result follows.
E.5.3 PROOF OF THEOREM
Theorem 3 follows from three constituent propositions, corresponding to items (b), (c), and (a) of the theorem, respectively.
Proposition E.4. Let ((t) be continuous semi-martingale (see Le Gall (2016, Definition 4. 4.19) and let σ(0) be defined such that 0 = (()) and (D.6) holds. Suppose that Yt is continuous and let G(t) = So r(s)ds. Then,
6 7 = e -() ((0) + [' ( c(o) g(+) ds.
(E.7)
Furthermore, if [1 has finite first moment for all S t, then
E (0) + S/ é ((0)-o1) B [[o[o]] ds.
0 c GG()
E
(E.8)
55

---

Published as a conference paper at ICLR 2024
Proof. By the Picard-Lindelöf theorem (Lindelöf, 1894), 0,5) exists and is unique. Noting that 0. given in (E.7) is differentiable by the fundamental theorem of calculus, we have
d0
g (0) + +2000) - - l/" (0 o ds
Ye G(T)
dt
V/ (0()) ((()
Because G(0) = O, we have 0 (()) and thus by the uniqueness the solution the expression given in (E.7) is correct.
For the second statement, we may apply Fubini's 's theorem and the first moment assumption.
We remark that if (E.7) seems unmotivated, one construct guess for this solution informally in much the same way as the final part of the proof of Proposition E.5 below proceeds. Our second main computation is to upper bound the covariance matrix of 0. by that of ((t);
Proposition E.5. Let ((*) and " be as in Proposition E.4. Further suppose that E for all t. Then,
| < co
I -C()) E ( & (0) + ('' (0s)) E =((lt) @ ((o) ds, (E.9)
E
0 4
0(0) 7
where u & v T is the outer product and 3 denotes the Loewner order.
Proof. We begin by applying the chain rule to get
d 0
8 0.0)
d0
dt
dt
where we are using Einstein notation to express the contraction of the third order tensor coming from the Jacobean with the first order tensor arising from the time derivative of 0.. By the definition we have then that
of 07
@ 2 040
d
do
Ø
+
@(*)
dt
dt
dt
0
(0()
(
(0)))
θ(0)
x(0)
7
3 2a (t())
0
0
0
where we used the fact that for any two vectors u, Rd it holds that
u x v + & 3 u @ u &
We may take expectations both sides of the preceding inequality and observe that the expec- tation exists by the finiteness of the second moment of the norm. Let Σ(t) E [(()) o(() and
Σ(t) E
0
0(0)
then, by Fubini's theorem and the above computation, have
dE(t) 3 [t() -Et)))((t))
dt
rearranging gives Multiplying by exp(G(tt) on both sides (which is positive definite and thus preserves the order) and
dt ( (00) G(0 dÉ(())
+ G(0, 0(t) O())
dt
Integrating and again rearranging tells that G(()
ε(1 1 e G(t) [ (((t)
The result follows.
56

---

Published as a conference paper at ICLR 2024
Note that both (E.8) and (E.9) should make intuitive sense. Indeed in the two limits, as Y 0 (corre- sponding to only keeping the first iterate) and y 1 oo (only keeping the last iterate), both expressions behave as expected. We remark that in the case where (()) E IR a scalar, then Jensen's 's inequality proves the result directly, without the need for the more complicated computations. As corollary of the above, letting Cov(-) denote the covariance matrix of a random vector and combining (E.8) with (E.9), we that
Cov @ 3 e G(t) Cov (000) (00)) f/ c(t)-C((ty. Cov (0(+)) ds.
In order to simplify our computations, we will suppose that 0 is Gaussian for all t. In this case, the following lemma shows that 0 is also Gaussian for all t:
Lemma E.6. Suppose that (0(*) ocess is Gaussian process with almost surely continuous sample
paths. Then, 0/0) - g is also Gaussian.
Proof. We begin by observing that for any partition 0 = < < < 8n-1 = t, the random vector
O
Ysg e G(x,) ((t)))) (8j 83 8j-1)
is Gaussian with
Σ j=l
Yoj G((1)) [o(so]] (s} - 8
and
11
O 2 On]
Ys Ys (G ((ss))-G()))((000-0000) (plx)) ol+))) (s) 83-1) 1))))) (st
st-1),
j=1==1
by the fact that (0(*)) 1 oc, it holds that E [0,] and E [O,, & O,] approach limits given by the Riemann integral. Thus, because the limit of Gaussian sequence Gaussian (in LP for all 1 S p < oc by, for example
Oeaase a Gaussian process. Now, by the continuity assumption, we see that as
Le Gall (2016, Proposition 1.1)), we see that 0 (0))) is Gaussian, as desired.
E.6 INSTANTIATIONS OF CLIFF LOSS IN CONTINUOUS TIME
We now prove the main results. We begin with the simpler result (presented second), where we do not concern ourselves with drift.
Proof of Proposition D.7. By Lemma E.6, 0 is Gaussian. Let >> such that (te(Cor(Oc)))
2. which occurs because in all regimes, except the third, Cov (0()) t oo; in the case of inverse fearning rate, simply set 7) large enough so that this holds. By Lemma E. it holds that
P (J (0 cc)) 2 c
V tr (Cov (0(:))
2
Thus E [J (o(+)) < C. On the other hand, if we let such that the upper bound on Cov 0 is
at most d e) then by Lemma E.2 it holds that with probability at least 1 - - that 0 S
By the construction of J. the first result holds. For the second result, note that fBc J whenever I the result holds by letting 7) +(c).
||00)|
Finally, we prove the second result, for OU processes.
57

---

Published as a conference paper at ICLR 2024
Proof of Proposition E.3. We apply Corollary E.2 to compute the first and second moments of ((t) and 0 Let R O(0) ul| We now note that for > O, by the assumption that c S ca ca= for
some small c, it holds that Cov (Ø(9) 2 4c2e2. Thus, by Hoffmann-Jorgensen et al. (1979) and the convexity of the norm followed by Lemma E.I, it holds that
P ( u <<) IP (000) E [
<< Se. 4 2
1
Thus, by the construction of the cliff loss,
IE (010)
C
2
@(0) have
On the other hand, we note that if then by the computation of E
b 0 m) P1 S 3R.
On the other hand, for 7 4 O, we see that Cov 0 + and thus we may apply Lemma E.2 to conclude that as long as 3R 1, there is some small 7 such that
050)
P 2) I P
0
(1 3R)c - C
The the first statement then follows from the construction of the cliff loss as long as K > 3. The second statement follows in the same way in the proof of Proposition D.7.
E.7 PROOF OF THEOREM 2
First, we show a reduction to lower bounding quantity involving the k-length suffix average.
Lemma E.7. Let F satisfy F(Ø) F- 2 (h, 0), and let (0(1) be a sequence of iterates with (h, ((t)) 2 0. Finally, let (w,) be a sequence of weights which are non non-increeaing. Then, for any
k S K, we have
T
F
Σ wr((t)
-F max WT-K+1
Σ
(0
h),
w w,hToo) (
I=T k+1
Proof of Lemma E.7. Under the assumption of the lemma, we lower bound
F
Σ
(+)(t)
h
Σ w,
1=1
W. (0)
we (0()) h)
k+1
Σ
(0(t) h)
WT-k41
1=
= kwT k+1
Σ
(0(0, h).
1=T-k+1
It is straightforward to also bound the above expression by wi h T (1)).
We can now prove our desired lower bounds. We can write
Ø
wi (1)
58

---

Published as a conference paper at ICLR 2024
where the weights induced by the EMA averaging
W1 II
wk 7 II (1 71) t=k+1
t=2
and which, for fixed 7, satisfy
W1 = (1 (1-7)²,
w = 7(1 - -)²=
Harvey et al. (2019, Theorem 3.4) exhibit function F which 1-strongly convex and 3-Lipschitz,
satisfies F(0) F* 2 (h,0), and sequence of iterates @()) corresponding gradient descent with step size ?/: 1/1 satisfying (0(9, h) 2 0 for all and
T Σ 1+T-k+1
log (og(T/k))
( h) Ω
(h, 01) 2(1).
Similarly, it is shown that there exists an (1) Lipschitz function inducing iterates that satisfy
[oog (log(T/k)
e h) 2 2
VT
k+1
Thus, for any non-increasing weighting scheme (w;), we have for all k
Ω [((VT
strongly convex case
F(O(T)- F NWWT-K4 log(T/k)
non-strongly convex case
We quickly note that for both these functions, we also have the lower bound F(O -
F* 2 w 0
due to Lemma E.7. For fixed, the expression wi = 71- shows that F F(O,,)) F* 2 T provived that 7 colong T for universal constant. Selecting p 1/4 gives the last statement of our theorem.
remains quantify the worst-case behavior of
(kWT-R+1 (oog(///))) w1)
under EMA.
Consider EMA with parameter - and choose h: = [1/7] = 1 Then, KWT- 2 k> 7)*h ((7k) 9(1).
Moreover, log(T/k) = 2(log(TT)) = 12 T). Thus, kwT-k41 (Tg( = 223 log 7").
Consider EMA with decaying parameter "tt = - Then, we have for k > T/2
WT-k+1 (T 1+
II (1- 22 (T k k+ (1 (T k +) Byt T=+1
(E.10)
Again select A: = 8 and assume that k < T/2, which requires T < T/2, T >
21's. Then, we can lower bound (E.10) by at least 2 2-8T-8(1 2T - Aggain, for k == @) and B E [0, 1J, this expression is ((T- Hence kwt k+1 ?(1) and kwT- ++1 log(T//L /k) = 2(3 log
59


Reviews:
[{"review_id": "RUQKaUifDw", "reviewer": "Reviewer_1ahA", "summary": "The paper examines a known problem in behavior cloning. This problem occurs during the training of such agents, where the validation criterion (return of an episode under long horizon rollouts) has large variance at all training steps, while the training surrogate loss has small, making it hard to perform model selection.\n\nThe authors perform a theoretical and empirical study of the origins of this problem. They argue that the observed variance is due to training instability, instead of insufficient dataset size, and they name the phenomenon as *Gradient Variance Amplification (GVA)*. They suggest that alternative training algorithms might not have the same issue. For this reason, they propose a very simple fix, which is to track the exponential moving average (EMA) of the parameter iterates of the SGD trained model. They demonstrate the effectiveness of the approach by performing experiments under various environments and by ablating design choices of the proposed EMA intervention.\n\nIn addition, they argue that GVA generally exists whenever agents are expected to operate by conditioning themselves in their past output (or effects of them). This definition aligns with conditional language modelling, and they argue with experiments that language generation quality in LLMs is similarly affected and that EMA can also help mitigate the problem there.", "questions": "1. In **Proposition 3.1**, I guess that $\\Delta$ refers to two different smooth “error functions” in each of the two inequalities.\n\n### Typos\n\n2. **Section 4.3**: “Here training loss is convex, but rollout reward is not.” The training loss is indeed convex, but the rollout reward is (still) concave, even if it is discountinuous. Is that right?\n3. Later in **Section 4.3**: “SGD iterates:” $\\theta_{t+1} = \\theta_{t} - …$ instead of $\\theta_{t+1} = \\theta_{t+1} - …$\n4. **Appendix A.1**/**Role of SGD noise**: “It is now well appreciate*d* that gradient noise facilitates the escape *from* saddle points.”", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* The paper is well-organized, well-written and well-argumented. It introduces and motivates the problem sufficiently and discussed related work in depth.\n\n * They pose a single clear question at the heart of the problem: Do we just need more data? Or is there in the training algorithm that amplifies variance for the validation criterion? Their ablations convince that the second is the case (section 3.1/figure 2), and they provide with theoretical insight that the return in horizon H can be exponentially large for some suboptimal policy in a smooth neighborhood around the expert, even though for all of those suboptimal policies in the neighborhood the surrogate training loss (of behavioral cloning) is small.\n\n * They introduce various decision choices around implementing EMA as a candidate solution, and they ablate many of them, demonstrating that the problem is effectively mitigated in many RL environments.\n\n * They make connections to the relevance of the problem to other topics in ML, such as autoregressive language modelling.", "weaknesses": "1. Ablations regarding the design choice of $\\gamma_t$ scheduling are missing. In the paper, a polynomial decay is used and ablated, but other schedules have also been used with EMA (like cosine). It would be nice to understand a bit better why this decay is important and how to design one which is tailored at the problem at hand.\n2. Middle of Figure 4 misses y-axis values, which is important in order to know at which scale are we seeing the zoom at.\n3. In the context of autoregressive language generation: Validation perplexity perhaps does not show the existence of GVA problem here in the most clear way. Ideally, some equivalent to a metric of generation quality of horizon H autoregressive rollouts should have been used instead.", "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": "", "extra_scores": {}}, {"review_id": "jRRYSJnrX8", "reviewer": "Reviewer_xe8u", "summary": "The paper examines how noise in gradients affects systems with feedback loops. It argues that SGD error accumulates over longer horizons, leading to Gradient Variance Amplification (GVA), explaining sharp reward oscillations. GVA is shown to be the main cause of these oscillations, surpassing the influence of statistical and architectural factors. Empirical evidence shows that the empirical moving average (MVA) mitigates these oscillations, ensuring stability. Additionally, the paper revisits the theory of stochastic optimization for convex functions, assessing its explanatory power for empirical observations of EMA and various step size schedules.", "questions": "It is mentioned in comments after proposition 3.1 that there is a good subset in parameter space that do not experience this worst-case error amplification. Does such good neighbourhood exists around any element in the parameter space or only around elements with specific properties ?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "a) The paper conducts meticulous experiments revealing that the variance in stochastic gradients is the true source of instability. It introduces exponential moving average as an effective solution to mitigate this issue. Additionally, this phenomenon is also demonstrated for other tasks with feedback loop, i.e., auto regressive processes for language generation. \nb) The paper is well-written and the main message is clearly presented.", "weaknesses": "a) The main problem of the paper in my opinion is the explanation provided for the benefits of EMA.  Proposition 3.1 says that land scape is very intricate and for every $\\delta$ there is a separation between $J$ and the behaviour cloning loss. However, the *cliff*-type loss framework used to study this does not capture this behaviour, as it is small in a neighbourhood of radius $\\epsilon $ and is very large outside. In my opinion, this framework it too tailor-made to study the benefits of EMA and does not reveal the real reason behind its working mechanism.", "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": "", "extra_scores": {}}, {"review_id": "ey7a4KrgEf", "reviewer": "Reviewer_fhjo", "summary": "The paper proposes that the training instability in the policy network is due to the gradient variance amplication during training", "questions": "See weakness", "limitations": "", "rating": 3, "confidence": 2, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "The idea is novel and the problem being dealt with is sufficiently important, namely, during training in reinforcement learning, the model does experience strong instabilities, and good methods to mitigate them would be of great importance", "weaknesses": "I am quite confused by the main claim of the work. \n\nThe main claim is that the gradient variance is responsible for the instability and that it amplifies throughout training, but no numerical result in the paper really plots the gradient variance, and, of course, not a single experiment shows that this variance is amplified. It should not be difficult to plot the variance of the gradient at all, and the absence of such evidence makes it impossible for me to recommend acceptance.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)