TL;DR. We compare NeXus with plain SGD inner steps to gradient accumulation. Both the river–valley and linear stability analyses suggest a bias toward flatter minima in NeXus under suitable assumptions.
NeXus is a counterpart to gradient accumulation. Both use several mini-batches. Accumulation computes all gradients at the same model. NeXus updates an inner model after each gradient, and the next gradient is computed at the new position.
This reminds me of local SGD. Each worker takes several steps before the models are averaged. Under some assumptions, these steps help the model move toward flatter regions (Gu et al., 2023). Could NeXus's inner steps do the same?
Notation. Let $\boldsymbol\theta\in\mathbb R^D$ be the parameter vector. Let $\ell_i(\boldsymbol\theta)$ be the loss on sample $i$. We consider the average batched losses and gradients, where each mini-batch $\mathcal B$ contains $b$ samples:
$$ \mathcal L(\boldsymbol\theta):=\frac1n\sum_{i=1}^n\ell_i(\boldsymbol\theta), \qquad \mathbf g_{\mathcal B}(\boldsymbol\theta) :=\frac1b\sum_{i\in\mathcal B}\nabla\ell_i(\boldsymbol\theta). $$
Modified NeXus. We simplify NeXus to study the role of its inner steps. We use plain SGD for the inner steps and a plain outer update. We also use ordinary gradients for accumulation.
<aside> 💡
Algorithm 1. (Modified NeXus).
Take $q$ inner steps of size $\gamma$. Start with $\boldsymbol\theta_0=\boldsymbol\theta$. At step $k$, draw mini-batch $\mathcal B_k$ and update
$$ \boldsymbol\theta_k =\boldsymbol\theta_{k-1} -\gamma\mathbf g_{\mathcal B_k}(\boldsymbol\theta_{k-1}). $$
Return
$$ \boldsymbol\theta^{+} =\boldsymbol\theta-\eta(\boldsymbol\theta_0-\boldsymbol\theta_q) =(1-\eta)\boldsymbol\theta+\eta\boldsymbol\theta_q. $$
Algorithm 2. (Gradient accumulation).
Draw $q$ mini-batches. Compute all gradients at $\boldsymbol\theta$. Average them and update:
$$ \boldsymbol\theta^{+} =\boldsymbol\theta-\eta\gamma \left(\frac1q\sum_{k=1}^{q}\mathbf g_{\mathcal B_k}(\boldsymbol\theta)\right). $$
</aside>
The methods are the same when $q=1$. When $\eta=1$, modified NeXus is just ordinary SGD grouped into $q$ steps.
Kaiyue Wen et al. use a river–valley landscape to explain learning-rate schedules. The river follows a relatively flat direction. The valley walls are steep. Inspired by this picture, we use a two-dimensional river–valley model to study NeXus's bias toward lower sharpness.
A river between steep walls. Let $x$ be the coordinate along the river. Let $y$ be the coordinate across the valley, where the loss rises steeply.
$$ \mathcal L(x,y)=C+\frac12\lambda(x)y^2, \qquad \lambda(x)>0. $$
Every point $(x,0)$ has the same loss. The Hessian there is
$$ \nabla^2\mathcal L(x,0) =\begin{pmatrix}0&0\\0&\lambda(x)\end{pmatrix}. $$
Here $\lambda(x)$ is the largest Hessian eigenvalue. A smaller $\lambda$ means a flatter valley.
Mini-batch noise. Model the mini-batch losses by
$$ \mathcal L_s(x,y) =C+\frac12\lambda(x)y^2+s a y, \qquad s\in\{-1,+1\},\quad a>0, $$
Each sign is drawn independently. Both signs are equally likely. The scalar $a$ sets the noise size.
<aside> 💡
Theorem 1 (Two-step sharpness reduction). Assume $\lambda>0$ and its first two derivatives are continuous. Start at $(x,0)$ and take $q=2$ inner steps. For every pair of noise signs, modified NeXus gives
$$ \Delta x:=x_{\mathrm N}^{+}-x =-\frac{\eta\gamma^3a^2}{2}\lambda'(x). $$
If $\lambda'(x)\ne0$, choose a small enough $\eta>0$. Then $\lambda(x_{\mathrm N}^{+})<\lambda(x)$. For averaged gradient accumulation, $\Delta x_{\mathrm A}=0$ in this iteration.
</aside>