Summary
The exponential moving average (EMA) of model parameters is used in many contexts, including to improve robustness in supervised learning, and as an important component of self-supervised learning. Given model parameters $\theta_t$ at iteration $t$, the exponential moving average $\zeta_t$ is updated as $\zeta_{t+1} = \rho \zeta_t + (1 - \rho) \theta_t$, where $\rho$ is the EMA coefficient. This paper proposes a scaling rule that dictates how $\rho$ should be modified to preserve training dynamics when the batch size is modified. In particular, when the batch size $B$ is scaled by a factor of $\kappa$ yielding $\kappa B$, the EMA coefficient should be exponentiated to the power of $\kappa$, yielding $\rho^\kappa$.
The authors justify this rule by analogy to the learning rate vs batch size rule from (Goyal et al., 2017), under the assumption that the gradients change extremely slowly between successive steps in training. They also give a justification using an SDE interpretation.
Empirically, they evaluate the EMA scaling rule on several real-world tasks, including supervised learning for image classification and speech recognition, semi-supervised learning for speech recognition, and self-supervised image representation learning. In each case, they compare the loss curves of: 1) a baseline model trained with batch size $B$ and EMA coefficient $\rho$; 2) a model trained with batch size $\kappa B$ and the same EMA coefficient $\rho$; and 3) a model trained with batch size $\kappa B$ and modified EMA coefficient $\rho^\kappa$. They show that using the scaling rule $\rho^\kappa$ produces more similar loss curves to the baseline model than the un-modified $\rho$.
Strengths
* The authors present a toy example where they apply the EMA to a noisy quadratic problem, where the "batch size" is simulated by adjusting the gradient variance. On this task, they compute the optimal EMA coefficient $\rho^*$ for various scaling factors $\kappa$ up to 1024, and show that the $\rho^\kappa$ coefficient nearly matches $\rho^*$ with respect to the approximation error. I appreciate the inclusion of such a toy example.
* The empirical evaluation considers diverse tasks, including both vision (ImageNet1K classification with a ResNet) and language (speech recognition with a Transformer) tasks. These experiments also consider different optimizers (SGD and Adam).
* The authors also evaluate progressive scaling, in which they increase the batch size during training and simultaneously scale the LR and EMA coefficient. They found that in one case (training a BYOC Vision Transformer), this was necessary to obtain decent performance.
* The plots are generally well-made and clear, with good labels.
Weaknesses
* One of the main drawbacks of the paper is that it is not clearly written. The core idea is very simple, and should be simply explained. However, most of the paper is needlessly obfuscated, making it hard to read. The majority of the definitions (2.1, 2.2, 2.3, and 3.1) are not novel, and are basic concepts that would not typically be presented as formal definitions. This reads as though the authors are trying to add math and formality to the paper, when it is not necessary.
* The plots in Figure 1(a) are not clearly presented, as the legend overlaps with more than half the figure.
* While the experiments use several architectures and datasets, they all measure the same thing, comparing the loss values of models trained with either $\{ (\kappa B, \rho^\kappa), (\kappa B, \rho), (B, \rho) \}$. There is very little diversity in this respect. Are there other things that can be done with the ability to match scaling behaviors, like doing hyperparameter optimization in the small-batch regime (which could be run in parallel on multiple devices rather than requiring a large amount of compute for a single experiment) to find the optimal $\rho$, and then scale it to the large-batch setting using $\rho^\kappa$?
* For the Transformer experiment in Figure 3, it is hard to see the long-term behavior of the loss curves. It would help if the y-axis was log-scaled. Here, the red curve ($\rho$ unmodified) seems to do just as well as $\rho^\kappa$, especially in terms of final performance.
* Throughout the paper, the notation $\times$ is used for simple scalar multiplication. This is non-standard in ML literature, and confusing to the reader, as $\times$ looks like a cross product. There is no reason to write $\eta \times \frac{1}{B}$ rather than simply $\frac{\eta}{B}$.
* In Eq. 6, it is unclear what $\Theta$ and $Z_t$ are, as they are not defined.
* Probable typo in L163: "scaled additive $b > 0$" --> "scaling $b > 0$"
* The base results (black dashed lines) in the real-world experiments do not seem to have bands for the mean and std over three runs (only the red and blue curves have bands).
**Minor**
* The math font looks non-standard.
* I do not think that the abstract does a good job describing the exponential moving average.
* The term "model EMA" should be more clearly defined at the beginning, as "the EMA of the model parameters."
* In the abstract, it is not clear what is meant by "Prior works have treated the model EMA separately from optimization."
* In the abstract, what is meant by "optimally a $6\times$ wall-clock time reduction"? Is this the actual wall-clock reduction achieved in practice, or under some idealized scenarios? This should be clarified.
* L32: "These weights are updated through by a momentum hyperparameter." --> Bad grammar "through by," and it does not make sense to say that the weights are updated "by a momentum hyperparameter."
* To improve clarity, the introduction should introduce the mathematical form of the EMA, $\theta' = \alpha \theta' + (1 - \alpha) \theta'$. The expression for the EMA needs to be introduced before Definition 1.1 because otherwise it is not clear what $\rho$ corresponds to and how it is used.
* Definition 1.1 should simply be labeled Definition 1, because there is no Definition 1.2. In Def. 1.1, it is confusing to say "scale other optimizers." Should this be "scale other optimization hyperparameters"?
* On Line 64, write out "self-supervised learning" instead of saying "SSL" because in this context it could be misinterpreted as standing for "semi-supervised learning."
* What is the purpose of including the third row in Eq. 4, which keeps the gradient $g$ unchanged? Is this just to highlight the assumption that the gradient changes very slowly, so that the first two rows hold?
* Typo: L115, "typically use to" --> "typically used to"
* L278 Typo: "a model EMA alters" --> "a model EMA that alters"
* L307 Typo: "only when by the EMA" --> "only by using the EMA"
* L210 Typo: "In summary, ASR" --> "In summary, this ASR"
Questions
* In many of the experiments, $\rho$ is very large: in Figure 1(a), $\rho=0.9999$, in Figure 2 $\rho = 0.9999$, in Figure 3 $\rho = 0.9995$, and in Figure 4 $\rho=0.9999$. Why are these EMA coefficients so large? Are these optimal values at the base batch size, and how were they chosen? How do these experiments behave when $\rho$ is smaller, like $\rho=0.9$? As a more general point, because $\rho$ is exponentiated, it seems that there could be numerical overflow or underflow issues depending on the initial value of $\rho$ and the scaling $\kappa$.
* In Eq. 5, what form does the gradient noise distribution $\mathcal{E}\_{\sigma}(\theta_k, \zeta_k)$ take? Is it Gaussian? If so, it would be better to use standard notation $\mathcal{N}$. What does it mean that the gradient noise distribution is "assumed to be zero-mean and variance $\Sigma(\theta_k, \zeta_k)$ independent of $\sigma$"? If the variance is independent of $\sigma$, why use the notation $\mathcal{E}_{\sigma}$?
* See also questions from the weaknesses box.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Limitations
* For the real-world experiments, it would be important to report scaling behavior beyond the factors that are shown in the paper. Why is the maximum scaling factor $\kappa = 8$? Why not continue the scaling to see if the scaling rule continues to hold, and determine where it breaks down? Does it work if $\kappa \in \{ 16, 32, 64, 128, \dots \}$? If the reason for not using larger factors is that the base model already used a large minibatch size ($B=1024$), then why not train a base model with a small batch size $B=16$ such that large $\kappa$ can be evaluated? In addition, because the scaling rule is meant to hold for both increasing and decreasing the batch size, why not have an example reproducing the results of a model trained with $B= 2048$ using a much smaller batch size like $B=16$? This would be a factor of $128$ decrease in batch size.
* This also relates to the question about numerical overflow/underflow, because if the original EMA coefficient is $\rho = 0.9999$ and $\kappa = \frac{1}{128}$ then the adapted coefficient would be $\rho^{\frac{1}{128}} = 0.99999922$. What are the limits to the approach, and what are its limitations?