ZEMA: confidence-calibrated signed descent with only one state

Introduction

Sign descent dynamics explain a considerable part of Adam’s success Bernstein et al. [4]; Kunstner et al. [14]; Kunstner and Bach [13]; Tomihari and Sato [23]; Zhao et al. [31], alongside recent geometric accounts that are outside the scope of this paper Xie et al. [24]; Bernstein and Newhouse [3]. However, sign descent methods degrade at smaller batch sizes relative to Adam Semenov et al. [20]; Kunstner et al. [14], which suggests a greater sensitivity to gradient noise: even extensively tuned Signum underperforms Adam Orvieto and Gower [18].

Adam as uncertainty-aware adaptation One explanation for this gap is that Adam scales its step by the uncertainty in the sign. If the sign of the stochastic batch gradient, $m_i = B^{-1}\sum_{n} g_i^n$ for per-example gradients $g_i^n$ with respect to the $i$th parameter, is positive with probability $\rho_i$, then the $i$th component of the Adam update is approximately $2\rho_i - 1$. When the sign is deterministic, i.e., $\rho_i = 0$ or $\rho_i = 1$, this reduces to sign descent; when both signs are equally likely, i.e., $\rho_i = 0.5$, the update is zero. Prior work interprets Adam as a three-stage procedure for estimating $\rho_i$: (i) assume a normal distribution for the stochastic batch gradient Ballés and Hennig [1]; Orvieto and Gower [18]; (ii) estimate its mean and standard deviation, $m_i$ and $\sigma_i$, by variational inference, implemented as an exponential moving average (EMA) Orvieto and Gower [18]; and (iii) compute the update $2\rho_i - 1$ from these estimates using an approximation to the error function, since $2\rho_i - 1 = \operatorname{erf}(z_i/\sqrt{2}) = \operatorname{erf}(m_i/(\sqrt{2}\sigma_i))$ for the $z$-score $z_i = m_i/\sigma_i$ Balles and Hennig [1].

Contributions. We take the magnitude of a signed update to reflect how strongly the batch supports that sign, and we measure that support per coordinate with a $z$-score. ZEMA computes $z_i$ from the spread of the per-example gradients in the current batch, converts it to the probability $\rho_i$ that the full-dataset gradient is positive, and smooths the resulting update with a single exponential moving average (Algorithm 1). Adam forms this probability after smoothing, whereas ZEMA forms it before, and therefore needs only one state per parameter. The per-example variance this requires costs 4.6% more per step (Table 14, Appendix E), and can be estimated from a small fraction of the batch (Appendix E.7). ZEMA matches AdamW on LoRA fine-tuning to within 0.04 perplexity (Table 1) and improves on it by 0.6 perplexity at 160M-parameter pretraining, at roughly a quarter of AdamW’s seed-to-seed standard deviation and at the learning rate chosen for AdamW (Table 2).

Table 1: LoRA fine-tuning: best validation perplexity over epochs, mean of three seeds. The first three blocks are a Llama scale ladder; Qwen3-1.7B is a single-scale check on a second family and its perplexities are not comparable across families. $\Delta$ is the paired difference from AdamW at the same seed, with a 95% Student-$t$ interval; positive favours AdamW. Learning rates are selected per cell (Appendix C.§ 1).

Model	Optimizer	E2E NLG ppl	E2E NLG $\Delta$ vs. AdamW	SAMSum ppl	SAMSum $\Delta$ vs. AdamW
Llama-3.2-1B	AdamW	2.959	—	8.134	—
Llama-3.2-1B	ZEMA-CLT	2.971	+0.011 ± 0.038	8.173	+0.039 ± 0.017
Llama-3.2-1B	ZEMA-Adam	2.975	+0.015 ± 0.026	8.172	+0.038 ± 0.017
Llama-3.2-3B	AdamW	2.911	—	7.157	—
Llama-3.2-3B	ZEMA-CLT	2.932	+0.021 ± 0.012	7.197	+0.040 ± 0.008
Llama-3.2-3B	ZEMA-Adam	2.931	+0.020 ± 0.011	7.194	+0.037 ± 0.007
Llama-3.1-8B	AdamW	2.878	—	6.595	—
Llama-3.1-8B	ZEMA-CLT	2.889	+0.011 ± 0.018	6.622	+0.026 ± 0.026
Llama-3.1-8B	ZEMA-Adam	2.887	+0.009 ± 0.017	6.621	+0.025 ± 0.024
Qwen3-1.7B	AdamW	2.953	—	8.297	—
Qwen3-1.7B	ZEMA-CLT	2.968	+0.015 ± 0.022	8.344	+0.047 ± 0.006
Qwen3-1.7B	ZEMA-Adam	2.967	+0.014 ± 0.028	8.343	+0.045 ± 0.009

Optimizer	states/param	learning rate	val. ppl	Δ vs. AdamW
AdamW	2	$8\cdot10^{-3}$	$22.79 \pm0.46$	—
ZEMA-CLT	1	$8\cdot10^{-3}$	$22.18 \pm0.12$	$-0.61 \pm0.55$
ZEMA-Adam	1	$8\cdot10^{-3}$	$22.36 \pm0.22$	$-0.43 \pm0.36$
Signum	1	$1\cdot10^{-3}$	$25.25 \pm1.02$	$+2.46 \pm0.60$
RMSProp	1	$1\cdot10^{-3}$	$26.71 \pm0.11$	$+3.93 \pm0.39$
Lion	1	$1\cdot10^{-4}$	$29.10 \pm1.51$	$+6.32 \pm1.35$
Adagrad	1	$1\cdot10^{-2}$	$38.31 \pm1.23$	$+15.53 \pm1.10$

Configuration	ms/step	vs. AdamW	Peak GiB
AdamW	2109	1.000×	22.20
Direct implementation of Algorithm 1	3359	1.593×	21.79
AdamW	2091	1.000×	22.20
Rank-local pooling + fused kernel, no overlap	2247	1.075×	21.15
All three changes	2188	1.046×	21.15

Related work Previous adaptive methods store only the first moment (Signum Bernstein et al. [4], Lion Chen et al. [5]) or only the second (AdaGrad Duchi et al. [7], RMSProp Hinton et al. [10]), whereas ZEMA keeps a single statistic, the $z$-score, that combines the information both provide. Muon Jordan et al. [12], which works in a different geometry, and GaLore Zhao et al. [30], which projects gradients into a low-rank subspace, are also used for training large language models, but their connection to signed descent is currently unclear and outside the scope of this paper; we include both in Appendix C.6 for a broad comparison. Variance-reduction strategies such as MARS (Yuan et al. [26]) could be applied on top of ZEMA, which we leave to future work. Finally, ours is not the first work to use per-example variance, or to motivate Adam in terms of a $z$-score-like statistic (Kunze et al. [15]; Balles and Hennig [1]; Patitucci and Mokhtari [19]), but these works do not use it to reduce the number of optimizer states, nor do they evaluate extensively in the LLM regime, where large-batch pretraining involves non-stationary gradients from rare tokens.

ZEMA: Z-score Evidence Momentum Adaptation

Adaptivity as evidence and confidence. Adam estimates $\rho_i$ indirectly, through two moments smoothed across steps. We estimate it directly from the current batch. The batch-level evidence is the $z$-score $z_i = m_i/\sigma_i$, which measures how far the batch mean sits from zero relative to its sampling error. Under the Gaussian assumption, the probability that the full-dataset gradient is positive is $\rho_i = \Phi(z_i)$, where $\Phi$ is the CDF of $\mathcal{N}(0,1)$; we call this probability the confidence in the sign. The confidence-adapted sign descent direction is $u_i = 2\rho_i - 1 = \operatorname{erf}(z_i/\sqrt{2})$, which we smooth over time to obtain the direction $\hat{u}_i$ we actually apply. Two changes relative to Adam follow: the evidence comes from statistics of the current batch rather than from moments accumulated across steps (§ 2.1), and the smoothing acts on the confidence rather than on the moments it is computed from (§ 2.2). The second is what leaves ZEMA with one state per parameter instead of two.

Because $\mathbb{E}[u_i]$ vanishes when the evidence does, coordinates with weak evidence take proportionally smaller steps than under Signum (Proposition 3); and because $z_i$ grows as $\sqrt{B}$ for any fixed non-zero gradient, the step saturates to the sign of the full-dataset gradient as $B$ grows (Proposition 2). Adam’s denominator retains the variation of the gradient across steps and so does not saturate the same way (Appendix A); Appendix B.§ 2 develops the evidence view further.

Estimating the $z$-score from batch statistics

At iteration $t$ we estimate the standard error of the batch mean, $\boldsymbol{\sigma}_t$, from the per-example gradients (Algorithm 1, line 5). This adds computation on top of the batch gradient $\boldsymbol{m}_t$ alone, but most of it can be shared with the usual backpropagation ([6 BackPack]). The overhead is largest in the pretraining setting of § 3.2, where a direct implementation costs 1.59$\times$ AdamW per step; the implementation of Appendix E reduces this to 1.046$\times$, a 4.6% increase in per-step time (Table 14). Appendix E.7 discusses a subsampling approach that reduces the cost of the dispersion further.

The evidence is then $z_t = m_t/(\sigma_t+\sigma_0)$, clipped to $\pm z_{\max}$ (line 6), where $\sigma_0 > 0$ prevents division by a vanishing standard error. It is a property of the current batch alone, with no dependence on earlier steps.

Smoothing the confidence over time

Two links turn $z_t$ into a confidence. Equation (5) gives $\rho_t^{\mathrm{CLT}} = \frac{1}{2} + \frac{1}{2}\operatorname{erf}(z_t/\sqrt{2})$ (Algorithm 1, line 8), and the rational approximation of Equation (7) gives $\rho_t^{\mathrm{Adam}} = \frac{1}{2} + \frac{1}{2}(z_t/\sqrt{1+z_t^2})$ (line 10), the form in which Adam’s own update can be written.

Since $\rho_t$ is estimated from a single batch, we smooth it over time. Our variational inference formulation is

whose solution is available in closed form as the exponential moving average

with $\beta= \frac{\lambda}{1+\lambda}$; a proof is given in Appendix B.§ 1. The update direction is $\hat{u}_t = 2\hat{\rho}_t - 1$. Because the smoothing and the map $\rho\mapsto2\rho- 1$ commute, Algorithm 1 equivalently forms $u_t = 2\rho_t - 1$ first (lines 8 and 10) and smooths that instead, so $\hat{u}$ is the only state it carries.

Experiments

We focus exclusively on LLM settings, where Adam has consistently been shown to be of particular importance ([14, 23, 29]). In the main paper, we present both fine-tuning and pretraining settings, to show the advantages and the limitations of ZEMA; experiments on non-stationarity (rare tokens and task mixtures) are in Appendix C.§ 3. Further experiments with different optimizers, with the choice of $k$ in batch subsampling, and on computational overhead are in Appendix C. These Algorithm 1: ZEMA, in its CLT and Adam-shaped variants.

Input: step size $\eta$; smoothing $\beta\in[0, 1)$; weight decay $\lambda$; $\sigma_0$, $z_{\max}$; link $\in\{\mathrm{CLT}, \mathrm{Adam}\}$

Output: parameters $\theta_T$

1. $\hat{u}_0 \leftarrow0$ // the only state, one scalar per parameter

2. for $t \leftarrow1$ to $T$ do

3. sample a batch of size $B$; form per-example gradients $g_t^{(1)}, \ldots, g_t^{(B)}$

4. $m_t \leftarrow\frac{1}{B}\sum_n g_t^{(n)}$ // batch mean

5. $\sigma_t \leftarrow\left[\frac{1}{B(B-1)}\sum_n(g_t^{(n)}-m_t)^2\right]^{1/2}$ // standard error of $m_t$

6. $z_t \leftarrow\operatorname{clip}(m_t/(\sigma_t+\sigma_0), \pm z_{\max})$ // z-score, this batch alone

7. if link = CLT then

8. $u_t \leftarrow\operatorname{erf}(z_t/\sqrt{2})$ // $2\rho_t^{\mathrm{CLT}}-1$

9. else

10. $u_t \leftarrow z_t(1+z_t^2)^{-1/2}$ // $2\rho_t^{\mathrm{Adam}}-1$

11. end

12. $\hat{u}_t \leftarrow\beta\hat{u}_{t-1}+(1-\beta)u_t$ // smooth the belief, not $z$

13. $\theta_t \leftarrow(1-\eta\lambda)\theta_{t-1}-\eta\hat{u}_t$ // decoupled decay, then step

14. end

include a probe of how often each optimizer recovers the sign of the full-dataset gradient, where ZEMA is ahead of AdamW by 9–12 points on the output head in pretraining, but not separable from it on the other tensors we probed (Appendix C.§ 4).

Fine-tuning

We compare both ZEMA variants against AdamW on LoRA (Hu et al., [11]) fine-tuning across three model scales—Llama-3.2-1B, Llama-3.2-3B and Llama-3.1-8B—and two generation tasks, E2E NLG (Novikova et al., [17]) and SAMSum (Gliwa et al., [8]). We add Qwen3-1.7B (Yang et al., [25]) at a single scale as a check that the result is not specific to one model family. Each optimizer is tuned on its own learning-rate grid, wide enough that the selected rate is interior, and the selected rate is then rerun with two further seeds. The setup, the rate-selection procedure, and the full sweeps are in Appendix C.§ 1.

Table 1 reports the best validation perplexity over epochs, averaged over three seeds. Both optimizers see the same seeds, so each difference is formed per seed; the reported interval is a paired 95% Student-$t$ interval.

The two optimizers are separated by at most 0.05 validation perplexity. Across all sixteen comparisons the gap lies between +0.009 and +0.047 perplexity, against absolute values of 2.9–3.0 on E2E and 6.6–8.3 on SAMSum, or 0.3% to 0.7%; six of the sixteen are indistinguishable from seed noise at the 95% level. We read the fine-tuning result as practical parity, with ZEMA carrying one state variable per parameter instead of two.

The seed-to-seed spread is also smaller for both ZEMA variants in four of the six Llama cells, most clearly on SAMSum at 8B (0.003 and 0.004 perplexity against AdamW’s 0.011); it is comparable at 3B/E2E and larger at 8B/E2E.

Pretraining

For pretraining, we replicate the 160M-parameter setting of Orvieto and Gower [18]: a 162.2M-parameter decoder-only transformer trained from scratch on 3.2B tokens of SlimPajama at context 2048 and global batch 256, three seeds per arm. AdamW and Signum run at the peak learning rates that paper selects, $8 \cdot10^{-3}$ and $10^{-3}$. The two ZEMA arms are not tuned and both run at AdamW’s. Table 2 adds the three one-state coordinate-wise baselines named in § 1—Lion, Adagrad and RMSProp—each swept on its own grid. Appendix C.§ 2 gives the architecture, the schedule and the corpus, and Appendix C.§ 2.1 the rate grids and the configuration choices behind these rows.

Both ZEMA variants finish ahead of AdamW, by 0.61 and 0.43 validation perplexity; both paired differences resolve at the 95% level. ZEMA therefore matches or exceeds AdamW at AdamW’s own learning rate, with a smaller seed spread and one state variable per parameter instead of two. A second pretraining setting supports reading this as parity or better: at 9.31M parameters on TinyStories, run to the same 3.2B tokens, the two optimizers finish level (Appendix D). Both settings show the same shape over training. AdamW leads early in validation loss—by 21.5 perplexity at 0.26B tokens here—after which ZEMA catches up, takes the lead by about 1.1B tokens, and holds it to the end, finishing level on TinyStories and 0.61 ahead here (Table 4). Which optimizer is ahead therefore depends on where training stops, which is a caution for comparisons made at a fixed update budget.

Table 4: SlimPajama, validation perplexity against token budget, three seeds per cell. $\Delta$ is ZEMA-CLT minus AdamW at matched seed, with a two-sided 95% $t$-interval over the three paired differences; negative favours ZEMA. The final column counts seeds favouring ZEMA. Bold marks intervals excluding zero. ZEMA-t2 is shown for reference and does not enter the paired test. The final row is the endpoint reported in Table 2.

Tokens	AdamW	ZEMA-CLT	ZEMA-t2	$\Delta$	95% CI	seeds
0.26B	94.76	116.22	115.84	+21.45	[+8.18, +34.73]	0/3
0.52B	47.49	49.94	50.57	+2.45	[−0.60, +5.50]	0/3
0.79B	38.49	38.73	39.24	+0.24	[−1.46, +1.94]	1/3
1.05B	34.50	34.20	34.64	−0.29	[−1.64, +1.06]	1/3
1.31B	31.90	31.37	31.70	−0.52	[−1.50, +0.46]	3/3
1.57B	29.89	29.26	29.52	−0.62	[−1.54, +0.30]	3/3
1.83B	28.09	27.48	27.71	−0.61	[−1.44, +0.22]	3/3
2.10B	26.49	25.85	26.07	−0.64	[−1.36, +0.08]	3/3
2.36B	25.02	24.40	24.59	−0.63	[−1.29, +0.03]	3/3
2.62B	23.84	23.22	23.41	−0.62	[−1.21, −0.03]	3/3
2.88B	23.07	22.45	22.64	−0.62	[−1.18, −0.05]	3/3
3.15B	22.80	22.19	22.37	−0.61	[−1.16, −0.06]	3/3
3.20B	22.79	22.18	22.36	−0.61	[−1.16, −0.06]	3/3

ZEMA is the only one-state method in Table 2 that gives up no validation loss for the reduction; the others trail AdamW by between 2.5 and 15.5 perplexity (Appendix C.§ 2.1).

Keeping one state instead of two also halves what the optimizer stores: 84.9M values on the bound weights against AdamW’s 169.9M (Appendix C.6).

Conclusion

Discussion & Limitations ZEMA trades computation for memory at little to no cost in quality: it halves the optimizer state for a few percent more time per step (Appendix E), and at every endpoint we measured it is at parity with AdamW or ahead of it. Which to prefer depends on which resource is scarce—ZEMA when memory is, as in training a large model on limited VRAM, AdamW when step time is. Two limitations qualify this: the $z$-score is formed within a single batch, so ZEMA does not accumulate the across-step gradient variation that Adam’s second moment carries, though this costs no validation loss in our experiments; and ZEMA converges more slowly early in pretraining, with AdamW leading until about a third of the way through our longest run (Table 4, Appendix D).

Conclusion We introduce ZEMA, an uncertainty-aware sign-based algorithm for training transformer-based language models. It forms the confidence in the sign of the gradient from the statistics of the current batch rather than from moments accumulated across steps, and because that confidence is formed before smoothing, one state per parameter does the work of Adam’s two. In our experiments this costs little or nothing in quality: parity with AdamW on LoRA fine-tuning, to within 0.05 validation perplexity, and 0.61 perplexity ahead at 160M-parameter pretraining at AdamW’s own learning rate, with a smaller seed-to-seed spread. Where the optimizer state is what limits training, dropping one of its two tensors for a few percent of step time is a favourable trade.

Future work We see three directions of future work. First, our work highlights that the choices inside the confidence adaptation of sign descent itself are still open. The link function assumes a Gaussian distribution for the stochastic gradient; the Student-$t$ link noted in Appendix B.§ 2 is one alternative, and heavier-tailed per-example distributions Battash et al. [2] are another. The target is a choice as well: the step $2p_i - 1$ minimizes the mean squared error to the gradient sign (Theorem 1), and a different loss—a more robust one, for instance—would map the same confidence to a different step.

Second, the construction is not tied to the sign geometry. Operator-norm optimizers such as Muon ([12]) and Shampoo Gupta et al. [9], which can be read as steepest descent under a norm ([3]), have no built-in confidence adaptation and are in that respect closer to Signum than to Adam. Adding it requires deciding how per-coordinate evidence should scale a matrix-valued step rather than one coordinate at a time; attempts exist [21, 16, 28, 27], but we are not aware of a principled derivation.

Third, differential privacy. Tang et al. [22] show that in DP-Adam the privacy noise dominates the second moment, so that the optimizer behaves like DP-SGD unless that noise is corrected out of the denominator; from the confidence view, such a correction discards the uncertainty the noise itself introduces into the sign. A confidence formed directly from batch statistics, i.e., from the $z_i$ computed in ZEMA, is a more natural target for privatization.

Adam as an uncertainty-aware sign descent algorithm

In this section we review the approaches from previous literature that treat Adam as an uncertainty-aware sign descent algorithm, together with the proofs and derivations they rest on.

Uncertainty adaptation of the sign descent algorithm

The proof is given in Appendix A.§ 2 below. ²

Proof of Theorem 1

While this theorem generalizes Balles and Hennig [1] as it does not assume that $g$ is normally distributed, it is itself a special case of the well-known fact that the mean of a random variable is the unique minimizer of the mean squared error to that random variable. For the first statement, since $\operatorname{sign}(g_i)$ takes the value $+1$ with probability $\rho_i$ and the value $-1$ with probability $1-\rho_i$,

Applying this argument coordinate-wise gives

For the second statement, let

and consider the objective

Expanding the squared norm,

Since $x$ is deterministic,

The last term does not depend on $x$. Completing the square gives

Therefore the unique minimizer is

Using the first part,

Distribution of sign of stochastic batch gradient

Assuming gradients are normally distributed, $g_i\sim\mathcal{N}(m_i,\sigma_i^2)$, the probability of the sign of $g_i$ being positive is given by

A derivation of Equation (5) is given in Appendix A.§ 4 below. Interestingly, this only depends on the ratio $z_i=\frac{m_i}{\sigma_i}$, and with this notation, the uncertainty-aware sign descent direction, $2\rho_i-1$, can be expressed as

Now, a central assumption to uncertainty-aware interpretation of Adam is the following approximation to the erf function.

A comparison plot for these two functions is given in Figure 1 in Appendix A.5 below.

Also, we can multiply the numerator and denominator of the approximation by $\sigma_i$, which changes the expression to

Which starts to resemble the Adam update. If we expand the denominator to see that using the definition of expectation and variance we have

So if we replace the average of the gradient $m_i$ by the EMA of the stochastic gradient, and $\sigma_i^2$ by the EMA of the squared stochastic gradient, we get very close to the Adam update.

The justification for the latter approximation is the topic of the next section.

Derivation of Equation 5

Let

and denote the cumulative distribution function of a standard normal random variable by $\Phi$. Since $\sigma_i>0$, standardizing $g_i$ gives

where the last equality follows from the symmetry of the standard normal distribution. Finally, using

we obtain

which is Equation (5).

Error-function Approximation

Figure 1 compares the error function with its rational approximation. In Balles and Hennig [1], a similar plot is shown but is expressed in terms of $z_i^{-1}$, and is only shown when $z_i>0$.

Adam's EMA as variational inference

Orvieto and Gower [18] give a variational inference interpretation of the EMA used in Adam optimizer. First, they show empirically that setting $\beta_1 = \beta_2$ recovers most of the performance gap between Adam and Signum. Then, they formulate a variational inference problem

and show that the result of plugging in the solution of this optimization problem into the equation (8) gives the same formulation of adam update, namely $\operatorname{EMA}_{\beta}[g_i]/\sqrt{\operatorname{EMA}_{\beta}[g_i^2]}$.

This last formulation has a simple interpretation, the parameterization of the Adam update tries to fit the mean and variance of the gradient distribution to the current batch gradient, while also keeping the distribution close to the previous one. The KL divergence term acts as a smoothness penalty, and the $\lambda$ parameter controls how much weight is given to the previous distribution versus the current batch gradient. This regularized variational inference approach is a design choice that we are going to use.

Note that, in Adam, first the mean and variance of the stochastic batch gradient are estimated using smoothed out variational inference, and then the signal-to-noise ratio $z_i$ and the update direction $2\rho_i-1$ are computed using Gaussian assumption and the approximation to the erf function. This ordering of operations is not necessary, and we can reverse the order of operations, and first compute $z_i$ and $\rho_i$ from the batch statistics, and then use variational inference to smooth out the estimated $\rho_i$ over time. This is the main idea behind our new optimizer, ZEMA. By this change of ordering of operations, we can reduce the number of state variables per parameter from two to one, while maintaining the variance adaptation properties of Adam.

Additional proofs and mathematical details

In this section we collect the proofs and mathematical details specific to ZEMA.

Proof of Equation 2

The problem in Equation (1) is separable across coordinates, so fix a coordinate $i$ and a step $t$, and abbreviate the two fixed probabilities and the optimization variable by

Writing out the two Bernoulli divergences, the objective is

The terms $p\log p + (1-p)\log(1-p)$ and $\lambda\left[q\log q + (1-q)\log(1-q)\right]$ do not involve $r$; collecting them into a constant $C$ and grouping the rest by $\log r$ and $\log(1-r)$ gives

Note that this is where the direction of the two divergences matters: because the optimization variable appears as the second argument of both, it enters only through $\log r$ and $\log(1-r)$, and the two fixed probabilities enter only through the coefficients.

Differentiating Equation (21),

Setting $J'(r)=0$ and multiplying by $r(1-r)>0$,

where the last equality uses that the two bracketed coefficients sum to $1+\lambda$. Hence the unique stationary point is

It is a minimum, and the only one: the second derivative

is strictly positive for every $r\in(0,1)$, since both coefficients are nonnegative and sum to $1+\lambda>0$, so $J$ is strictly convex on $(0,1)$.

Finally, substituting $\beta=\frac{\lambda}{1+\lambda}$, so that $1-\beta=\frac{1}{1+\lambda}$, Equation (25) reads

which applied coordinate-wise is Equation (2). As $\lambda$ ranges over $(0,\infty)$ the smoothing parameter $\beta=\lambda/(1+\lambda)$ ranges over $(0,1)$, so every smoothing level used in Algorithm 1 corresponds to some penalty weight $\lambda$.

The minimizer is interior. Equation (25) is a convex combination of $p$ and $q$, so $r^\star\in(0,1)$ whenever $p,q \in(0,1)$, and the argument above applies as stated. Both conditions hold in ZEMA: the z-score is clipped to $\pm z_{\max}$ before either link function is applied (Algorithm 1, line 6), so $\rho_t^{\mathrm{CLT}}$ and $\rho_t^{\mathrm{Adam}}$ are bounded away from 0 and 1, and the state is initialized at $\hat{u}_0=0$, that is $\hat{\rho}_0=\frac{1}{2}$.

Why the update is an average of probabilities. Had both divergences been taken in the opposite direction, with the fitted distribution as the first argument, the same calculation would have given a weighted average of the log-odds rather than of the probabilities,

with the same $\beta$. The direction chosen in Equation (1) is therefore what makes the smoothing a plain exponential moving average in probability space.

Smoothing $\rho$ and smoothing $u$ agree. Because $u=2\rho-1$ is affine in $\rho$, applying it to Equation (2) gives

the same recursion with the same $\beta$. This justifies computing the batch-level update direction first and smoothing it afterwards, which is what Algorithm 1 does.

Two properties of the confidence-scaled step $u_t$ in ZEMA

Fix a coordinate and drop the index $i$. The per-example gradients $g^{(1)},\ldots,g^{(B)}$ are i.i.d. with mean $\mu$ and finite variance $s^2>0$; $m$ is their mean, $\sigma=s/\sqrt{B}$ its standard error, $z=m/\sigma$, and $u=h(z)$ with $h(z)=\operatorname{erf}(z/\sqrt{2})$. Write $\zeta=\sqrt{B}\mu/s$ for the population evidence. All statements hold for any link $h$ that is continuous, odd, strictly increasing and bounded by 1 with $h(z)\to\pm1$ as $z\to\pm\infty$, which includes the Adam-shaped link $z(1+z^2)^{-1/2}$; the closed forms are for erf.

Proof By the strong law of large numbers $m\to\mu$ a.s., so $z=\sqrt{B}m/s\to\operatorname{sign}(\mu)\cdot\infty$ a.s. Since $h$ is continuous with $h(z)\to\pm1$, $h(z)\to\operatorname{sign}(\mu)$ a.s. With $\hat{\sigma}=\hat{s}/\sqrt{B}$ and $\hat{s}\to s$ a.s., the ratio $m/\hat{\sigma}$ has the same limit.

Hence for every coordinate with a non-zero full-dataset gradient, the ZEMA step converges to the Signum step, and the damping in Proposition 3 disappears. Adam’s step $m/\sqrt{v}$ does not share this limit: $v$ is a temporal average of squared batch gradients, so $v\approx\bar{\mu}^{2}+\operatorname{Var}_t(\mu_t)+s^2/B$ where $\operatorname{Var}_t(\mu_t)$ is the variation of the mean gradient across steps. As $B\to\infty$ the last term vanishes but the middle one does not, so $\lvert m/\sqrt{v}\rvert$ stays strictly below one whenever the gradient moves during training.

1. $|u| < 1$ for every finite $z$, whereas the Signum step satisfies $|u| = 1$;

2. $\mathbb{E}[u] = \operatorname{erf}(\zeta/2)$, which is odd, strictly increasing in $\zeta$, has the sign of $\mu$, and vanishes at $\zeta= 0$;

3. at $\zeta= 0$, $\mathbb{E}[u^2] = \frac{1}{3}$.

Proof. (1) is immediate from $|\operatorname{erf}(\cdot)| < 1$ on $\mathbb{R}$. For (2), write $u = 2\Phi(z) - 1$. For $z \sim\mathcal{N}(\zeta, 1)$ and an independent $W \sim\mathcal{N}(0, 1)$, $\mathbb{E}[\Phi(z)] = \mathbb{P}(W \le z) = \mathbb{P}(W - z \le0) = \Phi(\zeta/\sqrt{2})$, since $W - z \sim\mathcal{N}(-\zeta, 2)$. Hence $\mathbb{E}[u] = 2\Phi(\zeta/\sqrt{2}) - 1 = \operatorname{erf}(\zeta/2)$; the listed properties are those of $\operatorname{erf}$. For (3), at $\zeta= 0$ the variable $\Phi(z)$ is uniform on $(0,1)$, so $\mathbb{E}[(2\Phi(z) - 1)^2] = \mathbb{E}[(2U - 1)^2] = \frac{1}{3}$.

Property (2) says the expected per-coordinate step is $\eta\operatorname{erf}(\zeta/2)$: the learning rate is scaled down exactly where the batch gives weak evidence, and the scaling is a fixed function of the evidence rather than of the gradient magnitude. Property (3) quantifies the damping at a stationary coordinate: ZEMA’s mean squared step there is one third of Signum’s, which is the mechanism we associate with the reduced seed-to-seed variance in Table 1 and Table 2. Monotonicity in (2) holds without the Gaussian assumption for any symmetric noise, since $z$ then belongs to a location family and $h$ is increasing; the closed form does not. When $\sigma$ is estimated from $k$ examples, $z$ is Student-$t$ with $k - 1$ degrees of freedom rather than Gaussian; replacing $\Phi$ by the corresponding $t$ distribution gives the ZEMA-$t$ which we have not showed empirically since it performs approximately the same as ZEMA-CLT.

Experimental details and additional experiments

In this section we report the fine-tuning, instruction fine-tuning and pretraining details, and add a number of additional experiments to further validate ZEMA.

E2E and SAMSum fine-tuning details

Setup. LoRA adapters of rank 16 with $\alpha= 32$ are attached to seven projections per block; the base weights are frozen in bfloat16 and the adapters held in float32. E2E NLG is trained for five epochs at batch 64 with a linear schedule (3,290 steps); SAMSum for two epochs at batch 16 with a cosine schedule and 6% warmup (1,842 steps). The batch is held fixed across all four models, because the ZEMA statistic is a $z$-score of a $B$-sample mean and changing $B$ between cells would change the estimator rather than only the model. This is also why the Qwen3-1.7B numbers here supersede any run of ours at a different batch: at a different $B$ the arm is not the same estimator, so such runs cannot be placed on the same axis. All arms share $\beta= 0.95$ (the ZEMA smoothing and AdamW’s $\beta_1 = \beta_2$), $\sigma_0 = 10^{-8}$, and $z_{\max} = 10$. Runs are on a single NVIDIA L40S.

Selection metric. We select on the minimum validation perplexity over epochs, not the final-epoch value. On E2E every cell turns over between epochs two and four, so the final-epoch value measures where a run happened to sit on its own overfitting curve rather than the quality of the rate; ranking on it bends the sweep curves into a shape that is an artifact of the epoch budget. On SAMSum, which trains for two epochs and never turns over, the minimum and the final value coincide in every cell, so the choice changes nothing there.

Rate selection. Each optimizer is swept independently on each (model, task) cell, at a single seed, with generation metrics disabled. On the six Llama cells the initial grid is $\{2, 5\} \cdot10^{-4}$ and $\{1, 2\} \cdot10^{-3}$; it was then extended downwards—through $10^{-4}$, $5 \cdot10^{-5}$, $2.5 \cdot10^{-5}$, $10^{-5}$ and $5 \cdot10^{-6}$—for whichever cells needed it, until every cell’s argmin had two finished grid points on each side. We adopted this counting rule rather than a significance test on the neighbouring point, because the question a bracket has to answer is whether anything below the current edge is better, not whether the edge beats its neighbour by more than noise. The two Qwen cells were swept in a single pass on the wider grid those extensions had established, $\{5 \cdot10^{-5}, 10^{-4}, 2 \cdot10^{-4}, 5 \cdot10^{-4}, 10^{-3}, 2 \cdot10^{-3}\}$; every argmin is interior, and all but one have two points on each side. The exception is ZEMA-Adam on E2E, whose argmin of $10^{-3}$ has two points below and only one above; we report it as selected because the point above it is worse by 0.058 perplexity, roughly ten times the seed standard deviation of that cell, but it is a weaker bracket than every other cell in this table. The selected rate is then frozen and rerun with seeds 2 and 3; it is not re-derived at report time, so a later extension of the grid cannot silently repoint a seed that has already run.

Figure 2 shows the resulting sweeps and Table 3 the rate each cell selected. The perplexities in the figure are single-seed, from the sweep stage, and are therefore not directly comparable to the three-seed means in Table 1; they establish that every optimum is interior and give the scale of one grid step. Signum was also swept on every cell but is omitted from the tables: on Llama its optimum sits an order of magnitude below every other arm and it loses at every rung, so it is not part of this comparison. We note in passing that on Qwen3-1.7B it does not lose—its paired differences from AdamW are +0.002 ± 0.032 on E2E and +0.003 ± 0.006 on SAMSum, both indistinguishable—at an optimum of $10^{-4}$, five times below AdamW’s rather than ten. We have not investigated this and do not draw a conclusion from a single model.

Model	Task	AdamW	ZEMA-CLT	ZEMA-Adam
Llama-3.2-1B	E2E	$2 \cdot10^{-4}$	$5 \cdot10^{-4}$	$2 \cdot10^{-4}$
Llama-3.2-1B	SAMSum	$2 \cdot10^{-4}$	$5 \cdot10^{-4}$	$5 \cdot10^{-4}$
Llama-3.2-3B	E2E	$2 \cdot10^{-4}$	$2 \cdot10^{-4}$	$2 \cdot10^{-4}$
Llama-3.2-3B	SAMSum	$2 \cdot10^{-4}$	$5 \cdot10^{-4}$	$5 \cdot10^{-4}$
Llama-3.1-8B	E2E	$1 \cdot10^{-4}$	$2 \cdot10^{-4}$	$2 \cdot10^{-4}$
Llama-3.1-8B	SAMSum	$2 \cdot10^{-4}$	$2 \cdot10^{-4}$	$2 \cdot10^{-4}$
Qwen3-1.7B	E2E	$5 \cdot10^{-4}$	$5 \cdot10^{-4}$	$1 \cdot10^{-3}$
Qwen3-1.7B	SAMSum	$5 \cdot10^{-4}$	$5 \cdot10^{-4}$	$5 \cdot10^{-4}$

Two features of Figure 2 are worth noting. The curves are locally flat near their minima, so on most cells one grid step costs between 0.002 and 0.02 perplexity—of the same order as, and usually larger than, the entire gap between AdamW and ZEMA in Table 1. And the two ZEMA variants track each other almost exactly at every rate, which is expected: their links, $\operatorname{erf}(z/\sqrt{2})$ and $z(1+z^2)^{-1/2}$, agree closely over the range of $z$ these runs occupy (Appendix A.5).

SlimPajama details

Setup. A decoder-only transformer with 12 layers, $d_{\mathrm{model}} = 768$, 12 heads, SwiGLU feed-forward of width 2048, RoPE, context 2048 and an untied embedding and output head over the GPT-NeoX 50,280-token vocabulary: 162,183,936 parameters, matching the configuration of Orvieto and Gower [18]. Global batch 256 sequences (8 micro-batches of 8 on each of four ranks), 6,103 steps, 3.2B tokens. Linear warmup over 610 steps then cosine decay to $10^{-5}$; weight decay 0.1 on matrices only, gradient clipping 1.0, and $\beta= 0.95$ for every arm—AdamW in the equal-$\beta_1 = \beta_2$ form that paper identifies as the strong baseline. Four NVIDIA L40S.

As in the fine-tuning campaign the ZEMA arms are hybrid: the estimator covers the body linears (84.9M parameters) while the embedding, the untied head and the norm gains stay on AdamW at the same rate, leaving each arm a single learning-rate knob.

Corpus. The paper’s cerebras/SlimPajama-627B no longer resolves, so we use instead DKYoon/SlimPajama-6B, a published subset of the same source corpus, with its own held-out validation split (9.2M tokens) rather than a resampled tail of the training data.

What was and was not tuned. AdamW at $8 \cdot10^{-3}$ and Signum at $10^{-3}$ are the rates Orvieto and Gower [18] select at this scale; we ran no search of our own. The ZEMA arms have no published rate to copy, and rather than sweep them—which would have given them an advantage the baselines did not get—both were run at AdamW’s rate. The result in Table 2 is therefore a lower bound on what ZEMA does here.

Distance from the published number. Our AdamW reaches 22.79 $\pm$ 0.46 against the 21.86 that paper reports, so this setup runs somewhat harder than theirs—plausibly the corpus substitution above. That offset applies to every arm equally and does not affect the within-table comparison, which is why Table 2 is reported as paired differences against our own AdamW rather than against the published value.

Distributed correctness. All four arms trained on four L40S under DDP. The ZEMA arms do not use DistributedDataParallel: the bound weights carry no .grad, and their evidence statistic is pooled across ranks by a hand-written all-reduce, verified against a single-GPU run to $4 \cdot10^{-8}$ relative error in fp32 before the campaign was launched.

Curves. Figure 4 shows the full trajectories. Signum separates from the other three within the first 0.5B tokens and never rejoins them; AdamW and the two ZEMA arms stay together until roughly 1.5B, after which the ZEMA curves sit marginally below. The seed bands are widest early and narrow as the runs converge—at the end the ZEMA-CLT band is about a fifth the width of AdamW’s, the same ordering of seed-to-seed spread that Table 2 reports.

Paired differences over the run. Table 4 tabulates what Figure 4 plots, as paired differences at matched seed. The run has the same three phases as the TinyStories campaign of Appendix D: AdamW leads by 21.45 perplexity at the first checkpoint, the arms cross at about 1.1B tokens, and from 1.31B onward all three seeds favour ZEMA-CLT at every checkpoint. The difference then settles at −0.6 perplexity and stays there, resolving at the 95% level over the final four checkpoints. Unlike TinyStories, where the gap closes as the cosine anneals, here it persists to the endpoint of Table 2.

SlimPajama learning-rate selection

Every rate in Table 2 is either taken from Orvieto and Gower [18] or is the argmin of a sweep at seed 1 over the grid below. We report the whole grid, not only the winner: a rate quoted without its neighbours cannot be checked, and a selection procedure that stops at the first interior-looking point is exactly how a sweep reports a boundary as an optimum. Every argmin below is interior to its own grid. Table 5 lists the values and Figure 5 plots them.

Arm	validation perplexity by learning rate (seed 1)
Lion, $\lambda= 0.1$	31.8 $5\cdot10^{-5}$	29.6 $10^{-4}$	29.8 $2\cdot10^{-4}$	36.7 $5\cdot10^{-4}$	554 $10^{-3}$	749 $2\cdot10^{-3}$	871 $4\cdot10^{-3}$
Lion, $\lambda= 0.8$	31.4 $5\cdot10^{-5}$	28.8 $10^{-4}$	28.4 $2\cdot10^{-4}$	29.5 $5\cdot10^{-4}$
RMSProp	32.7 $3\cdot10^{-4}$	26.7 $10^{-3}$	27.8 $3\cdot10^{-3}$	531 $10^{-2}$
Adagrad	133 $10^{-3}$	55.9 $3\cdot10^{-3}$	38.7 $10^{-2}$	48.8 $3\cdot10^{-2}$	83.3 $10^{-1}$	279 $3\cdot10^{-1}$

Lion’s grid is an order of magnitude below its published one. Lion’s recommended rate for language-model pretraining is $3$–$10\times$ smaller than AdamW’s, which at AdamW’s $8\cdot10^{-3}$ would put it near $10^{-3}$. Every rate in that band diverged here: $10^{-3}$, $2\cdot10^{-3}$ and $4\cdot10^{-3}$ finish at perplexity 554, 749 and 871. The grid had to be moved a full decade down before Lion trained at all, and its optimum sits at $10^{-4}$, $80\times$ below AdamW’s. We report this because a reader reproducing Lion from its published guidance at this scale will see divergence, not a weak result.

Weight decay: $\lambda= 0.1$ is reported, $\lambda= 0.8$ is a control. Lion is normally run with a decoupled decay several times larger than AdamW’s, because its update is a sign and so has a fixed magnitude that does not shrink as gradients do. Table 2 reports $\lambda= 0.1$, the value every other arm uses, so that the table varies one thing. We also ran Lion’s own $\lambda= 0.8$ with its own sweep: it selects $2\cdot10^{-4}$ and reaches $27.97 \pm1.53$, a paired improvement of $1.13\pm0.24$ over the matched-decay row. So the larger decay does help, and materially—but it does not change the ordering in Table 2, where Lion trails AdamW by 6.32; $\lambda= 0.8$ would still trail by 5.19. Promoting the $\lambda= 0.8$ row into the main table would silently put an $8\times$ regularization difference inside a comparison whose point is that only the optimizer changes.

Lion with $\beta_1 = \beta_2$ is Signum, exactly. This is an algebraic identity, not an approximation, and it is worth stating because it constrains what a Lion baseline can show. Lion forms its update from an interpolation and then updates its momentum,

Setting $\beta_1 = \beta_2 = \beta$ makes the two expressions the same tensor, so $u_t = \operatorname{sign}(m_t)$ with $m_t$ the Signum momentum: the arms coincide for every input, at every step. Running Lion at the $\beta= 0.95$ this paper’s other arms share would therefore not have produced a Lion baseline at all, only a second copy of the Signum row. Lion is run at its own published $(\beta_1,\beta_2) = (0.9,0.99)$ for this reason, and it is the one arm in Table 2 not matched to $\beta= 0.95$.

Adagrad has no decay to match. Adagrad accumulates an undecayed sum of squared gradients; adding a decay to it is what makes it RMSProp. It therefore has no $\beta$ to set to 0.95, and the two arms together span that axis: RMSProp is the same second-moment normalization at $\rho= 0.95$, and is also exactly AdamW at $\beta_1 = \beta_2 = 0.95$ with the first moment removed. That makes the pair a clean ablation of the two AdamW moments taken separately, which is how they should be read—not as competitive optimizers at this scale.

Reading the one-state baselines. The four one-state arms of Table 2 separate widely, and which moment they keep does not predict where they land: keeping a first moment and taking its sign gives 25.25 (Signum) and 29.10 (Lion), while keeping a second and normalizing by it gives 26.71 (RMSProp) and 38.31 (Adagrad). Signum is the informative arm, as the one that takes the sign without scaling it by a confidence: its 2.46 deficit is what dropping that scaling costs in this setting (Appendix A.§ 1). ZEMA’s single state is of a different kind from any of these—the smoothed confidence, formed from the batch mean and its standard error together—which is why it is the only one-state arm in the table that does not pay for the reduction.

Instruction fine-tuning details (MATH and TULU)

In this section we present additional instruction fine-tuning experiments to evaluate non-stationarity. Instruction fine-tuning illustrates this because task mixtures can introduce significant drift (TULU) and rare tokens (MATH). We fine-tune Qwen2.5-0.5B (non-instruction-tuned) on a controlled dialogue/math/code Tulu-3 drift mixture using rank-8 LoRA adapters (4.40M trainable parameters, sequence length 512, 128 validation examples per source), and Llama-3.2-1B on the rare-token MATH split using rank-16 LoRA adapters (11.27M trainable parameters, sequence length 1024, 754 validation examples). Each cell performs exactly 48 optimizer updates at a global batch drawn from $\{8, 32, 128, 512\}$, drawing examples with replacement where the cell’s example budget exceeds the dataset, and is evaluated at eight checkpoints; we report the best validation loss over those checkpoints, exponentiated to a token-level perplexity. Adapters are held in float32 with weight decay 0.01 and $\beta= 0.95$ for both optimizers. Both campaigns use a single selection seed (11), so the tables below carry no confidence intervals and differences should be compared against the scale of seed variation rather than treated as resolved.

In the parameter-efficient setting the two optimizers are indistinguishable. Every difference in Table 6 falls within ±0.3%, and the sign is not consistent: ZEMA is nominally ahead at the two largest Tulu-3 batches and at the smallest MATH batch, behind elsewhere. With a single selection seed these margins are smaller than the seed variation one should expect from a 48-update run, and we therefore read the fine-tuning results as parity rather than as evidence for either method. Two observations survive this caveat. First, the fixed-update budget removes the large-batch degradation seen in the fixed-token setting: ZEMA does not fall behind as the batch grows, and at $B = 128$ and $B = 512$ on Tulu-3 it is marginally ahead, so the degradation reported elsewhere is at least partly attributable to update count rather than to batch size as such. Second, ZEMA’s selected learning rate is consistently two to four times AdamW’s at matched batch, a systematic offset rather than per-cell scatter; a shared grid would have handicapped ZEMA, and we recommend independent selection for any comparison of this kind. We note that the MATH $B = 512$ ZEMA cell remains unbracketed above, so its already small deficit may narrow further.

Batch	AdamW lr	ZEMA-CLT lr	$\Delta$	$\Delta$ %
Tulu-3 drift, Qwen2.5-0.5B, rank-8 LoRA
8	2.4568	$4 \times10^{-4}$	2.4594	$1 \times10^{-3}$	+0.0026	+0.11
32	2.4247	$2 \times10^{-4}$	2.4266	$5 \times10^{-4}$	+0.0019	+0.08
128	2.3968	$5 \times10^{-4}$	2.3915	$1 \times10^{-3}$	−0.0053	−0.22
512	2.3739	$5 \times10^{-4}$	2.3720	$1 \times10^{-3}$	−0.0020	−0.08
MATH rare-token split, Llama-3.2-1B, rank-16 LoRA
8	2.7018	$2 \times10^{-4}$	2.6997	$1 \times10^{-3}$	−0.0021	−0.08
32	2.6364	$5 \times10^{-4}$	2.6373	$2 \times10^{-3}$	+0.0009	+0.03
128	2.5503	$1 \times10^{-3}$	2.5577	$2 \times10^{-3}$	+0.0074	+0.29
512	2.5116	$1 \times10^{-3}$	2.5165†	$2 \times10^{-3}$	+0.0048	+0.19

Sign estimation accuracy in SlimPajama pre-training

To isolate sign estimation from trajectory divergence, a separate probe trains one shared AdamW trajectory. At every checkpoint, shadow AdamW and ZEMA states receive the same training batch at the same parameter vector; neither shadow state changes the trajectory. We instrument the QKV matrix in block zero and the untied output head. The reference direction is a deterministic held-out gradient, disjoint from training and reused at every checkpoint to compute the trget gradient at that checkpoint.

For direction $d$ and held-out gradient $g^{*}$, accuracy is the coordinate mean of $1\{\operatorname{sign}(d_i) = \operatorname{sign}(g_i^{*})\}$. The $\Delta$ columns are ZEMA minus AdamW in absolute percentage points; positive values favour ZEMA.

We run the probe on a SlimPajama model (12 layers, $d_{\mathrm{model}} = 768$, 12 heads, FFN width 2048, vocabulary 50,280, context 2048) trained at global batch 256 for 100M tokens, probing every 48 updates against a fixed 256-sequence oracle, with three seeds. Table 7 reports every probe rather than only the last, since the two layer groups evolve differently over training.

Step	Body QKV AdamW	Body QKV ZEMA-CLT	Body QKV $\Delta$	Output head AdamW	Output head ZEMA-CLT	Output head $\Delta$
1	68.30	68.30	+0.00 ± 0.00	69.84	69.84	+0.00 ± 0.00
48	50.38	55.30	+4.92 ± 7.60	52.34	61.68	+9.35 ± 0.90
96	53.53	54.62	+1.09 ± 3.25	54.09	63.68	+9.59 ± 0.73
144	55.44	51.62	−3.82 ± 5.65	52.20	63.73	+11.53 ± 0.67
190	56.13	53.65	−2.48 ± 6.37	52.69	62.10	+9.41 ± 1.45

The two layer groups separate sharply. The output head shows a large and stable gain, between +9.4 and +11.5 points at every probe after the first, with seed standard deviations below 1.5 points; this is the clearest sign-recovery result we obtain anywhere. Body QKV, by contrast, is not separable from zero at any probe: the seed standard deviation exceeds the mean at all four, and the point estimate changes sign over training, drifting from +4.92 at step 48 to −2.48 at step 190. We therefore report a sign-recovery advantage for this model only on the output head, and explicitly not on the body, and we do not claim a general mechanism.

Task	Tensor	AdamW	ZEMA-CLT	$\Delta$
E2E	L0 q-proj	51.50	51.29	$-0.21 \pm0.43$
E2E	L31 down-proj	51.61	51.54	$-0.07 \pm0.60$
SAMSum	L0 q-proj	48.28	48.82	$+0.54 \pm0.44$
SAMSum	L31 down-proj	52.58	52.05	$-0.54 \pm0.36$
E2E	pooled	—	—	$-0.14 \pm0.37$
SAMSum	pooled	—	—	$+0.00 \pm0.37$

Sign estimation accuracy in fine-tuning

We repeat the probe on the LoRA setting of the main text: Llama-3.1-8B-Instruct SAMSum ($B = 16$, 1,842 updates), probing every 460 updates against a 64-example oracle with three seeds. Two adapter tensors are tracked, an early attention projection and a late MLP projection.

Neither method recovers held-out gradient signs on adapter tensors: absolute accuracies lie between 48.3% and 52.6%, that is, within a few points of chance, and every difference is smaller than its seed standard deviation. The pooled effect is $-0.14 \pm0.37$ points on E2E and $+0.00 \pm0.37$ on SAMSum. We read this as a null result rather than as weak evidence in either direction, and note that where ZEMA governs only a thin adapter atop a frozen backbone, it achieves nothing measurable. Sign recovery on LoRA adapters should not be offered as evidence for the method.

Optimizers from other families: Muon, GaLore, and ZEMA inside GaLore

The arms in this appendix are not coordinate-wise sign methods, so they do not belong in Table 2; but they are the two best-known ways to train a transformer on less than two optimizer states per parameter, and leaving them out would make the memory claim in this paper look less contested than it is. They run in the identical setting: same architecture, corpus, schedule, batch, decay and three seeds.

All three use the same hybrid split ZEMA uses. Only the 84.9M body linear weights (52.4% of the model) go to the arm; the embedding, the untied head and the norms stay on AdamW, because Muon and GaLore both require matrices and neither the reference implementations nor ours applies them to an embedding table. State totals therefore include an AdamW pair on the remaining 77.2M parameters.

Muon. One momentum buffer per matrix, whose direction is replaced by its nearest semiorthogonal matrix through a five-step Newton–Schulz iteration in bfloat16, at the reference coefficients. It is the same one-state budget as ZEMA and it wins this comparison outright, by 2.65 perplexity against AdamW, four times ZEMA’s margin. We report this plainly because it bounds what the rest of the paper claims: a single state per parameter is clearly sufficient to beat Adam at this scale, and ZEMA is not the only way to get there. What ZEMA and Muon do is different in kind—Muon changes the geometry the step is measured in, and is defined only for matrices; ZEMA changes the per-coordinate scale, and applies to any parameter—and nothing in our results says the two cannot be combined. Its rate was swept on $\{3\cdot10^{-3}, 10^{-2}, 3\cdot10^{-2}, 10^{-1}\}$, selecting $3\cdot10^{-2}$ (perplexities 22.2, 20.2, *20.2,

21.5; Figure 6), inside the published 0.02–0.05 band. A ten-step Newton–Schulz variant was also run and finished within 0.10 of the five-step arm at every rate, so the five-step reference configuration is what we report.

Optimizer	lr	body state	total state	val. ppl	$\Delta$ vs. AdamW
AdamW	$8\cdot10^{-3}$	169.9M	324.4M	$22.79 \pm0.46$	—
ZEMA-CLT	$8\cdot10^{-3}$	84.9M	239.4M	$22.18 \pm0.12$	$-0.61 \pm0.55$
Muon	$3\cdot10^{-2}$	84.9M	239.4M	$20.14 \pm0.11$	$-2.65 \pm0.40$
GaLore	$1\cdot10^{-2}$	73.1M	227.6M	$23.58 \pm0.22$	$+0.79 \pm0.28$
ZEMA-CLT/GaLore	$1\cdot10^{-2}$	44.8M	199.3M	$24.10 \pm0.07$	$+1.31 \pm0.52$
ZEMA-Adam/GaLore	$1\cdot10^{-2}$	44.8M	199.3M	$24.14 \pm0.09$	$+1.36 \pm0.54$

GaLore. Adam run inside a rank-$r$ subspace of each gradient, taken from that gradient’s own SVD and refreshed every 200 steps. At $r = 256$ its state is $r(\min+ 2\max)$ per matrix, which for a $768 \times768$ matrix is exactly $d^2$: the published rank for $d_{\mathrm{model}} = 768$ and the rank at which GaLore matches ZEMA’s state coincide, so this is GaLore at its own setting, not squeezed to fit our comparison. It trails AdamW by 0.79 and ZEMA by 1.40, at 0.86× ZEMA’s body state. We swept both weight decays: $\lambda= 0.1$ (matched to every other arm) selects $10^{-2}$ and is the row above; the reference’s own pretraining default $\lambda= 0$ selects a different rate, $3 \cdot10^{-2}$, and is worse by 0.85 at $24.45 \pm0.15$. Running only the published default would have made GaLore the sole unregularized arm here and cost it most of a perplexity point.

ZEMA inside GaLore. The two memory mechanisms compose in principle: GaLore keeps three states of a smaller shape, ZEMA keeps one state of full shape, and running ZEMA as the inner optimizer keeps one state of the smaller shape. This is the cheapest arm we ran, at 0.53× ZEMA’s body state and 0.26× AdamW’s. It is cheap to compute as well as to store: projecting a per-example gradient requires projecting only one of its two operands, since $G_i = \delta_i^{\mathsf{T}} x_i$ gives $G_i P^{\mathsf{T}} = \delta_i^{\mathsf{T}}(x_i P^{\mathsf{T}})$ for a tall matrix and $P^{\mathsf{T}}G_i = (\delta_i P)^{\mathsf{T}}x_i$ for a wide one, so ZEMA’s per-example reduction is unchanged and simply runs at the smaller shape.

It does not compose in practice. Against GaLore-Adam at the same rate, decay and subspace—the comparison in which only the inner optimizer changes—the swap costs $+0.52 \pm0.29$ (CLT) and $+0.57 \pm0.28$ ($t_2$), paired over three seeds. The rate is not the explanation: we swept it over $\{3 \cdot10^{-3}, 10^{-2}, 3 \cdot10^{-2}, 10^{-1}\}$ (Figure 6) and both arms optimize at the same $10^{-2}$, with the swap costing 0.42–0.49 at every rate near the optimum and 1.90 at a rate $3\times$ too small.

Two observations point the same way. ZEMA’s advantage inverts under projection: at full shape ZEMA-CLT beats AdamW by 0.61, inside a rank-256 subspace it loses to Adam by 0.52. And the choice of link, which separates the two ZEMA variants by 0.18 at full shape, stops mattering inside the subspace—the paired difference between them there is $+0.045 \pm0.061$, an interval containing zero. Both are consistent with the projection having already discarded the per-coordinate structure that ZEMA’s standardization exploits: the top-$r$ directions of the gradient are not where the two links disagree.

Training horizon in TinyStories pretraining

The comparison in Section 3.2 is made at a single token budget. This appendix asks how much the outcome depends on that budget, in a second and much smaller pretraining setting: a 9.31M-parameter model on TinyStories, trained to 3.2B tokens at two batch sizes with three seeds per arm. The comparison there has three phases. AdamW leads early, ZEMA leads through the middle of training, and the two finish level, so a run stopped early can report either optimizer ahead. The same configuration stopped at a fixed update count also produces an apparent growth of AdamW’s margin with batch size, which the full horizon shows to be an effect of where training stops rather than of batch size.

Setup

The model is a 9.31M-parameter decoder-only transformer (6 layers, $d_{\mathrm{model}} = 256$, 8 heads, FFN width 768, context 128) trained on TinyStories with a native 8,192-token BPE vocabulary, with exact within-batch dispersion, body-only ZEMA coverage, and a micro-batch of 64 sequences. Each run trains for 3.2B tokens under a cosine schedule with 10% linear warmup over the full horizon, and is evaluated every 0.1B tokens on 4M held-out tokens, giving 32 checkpoints per run. Validation batches are drawn from a generator reseeded identically at every checkpoint, so all checkpoints of all runs are scored on the same held-out sequences and the training stream is unaffected by the evaluation cadence.

We run $B = 256$ and $B = 1024$ with three seeds per optimizer, twelve runs total. Learning rates are the bracketed optima of a six-point grid, searched separately for each optimizer at 2,000 updates at that batch: 0.006 at $B = 256$ and 0.012 at $B = 1024$. At both batches the two optimizers selected the same value, so a single learning rate is simultaneously each one’s own optimum and neither arm is advantaged by the choice. Runs used one NVIDIA H100 each on Trillium.

Results

Table 10 shows the same three-phase structure at both batch sizes. AdamW leads early by a margin that is significant at 0.10B and larger at the larger batch (+0.235 versus +0.090). The curves cross at approximately 0.4B tokens. ZEMA then leads for the remainder of training, by 0.01 to 0.03 perplexity, with all three seeds agreeing at every checkpoint from 0.8B onward and several intervals excluding zero. As the cosine anneals the gap closes, and at 3.2B the two are level: −0.004 [−0.015, +0.008] at $B = 256$ and −0.011 [−0.030, +0.008] at $B = 1024$, neither separable from zero.

Tokens	AdamW	ZEMA-CLT	$\Delta$	95% CI	seeds
$B = 256$, $lr$ 0.006
0.10B	7.2566	7.3466	+0.0900	$[+0.036, +0.144]$	0/3
0.20B	6.6933	6.7274	+0.0341	$[-0.002, +0.070]$	0/3
0.40B	6.4488	6.4458	−0.0030	$[-0.044, +0.038]$	2/3
0.80B	6.2208	6.1947	−0.0261	$[-0.083, +0.031]$	3/3
1.20B	6.0241	5.9959	−0.0282	$[-0.055, -0.002]$	3/3
1.60B	5.7759	5.7559	−0.0200	$[-0.035, -0.005]$	3/3
2.00B	5.5002	5.4793	−0.0209	$[-0.050, +0.008]$	3/3
2.40B	5.1610	5.1492	−0.0117	$[-0.038, +0.015]$	2/3
2.80B	4.7832	4.7736	−0.0095	$[-0.021, +0.002]$	3/3
3.20B	4.5885	4.5850	−0.0035	$[-0.015, +0.008]$	2/3
$B = 1024$, $lr$ 0.012
0.10B	7.3837	7.6183	+0.2346	$[+0.152, +0.317]$	0/3
0.20B	6.5267	6.6138	+0.0871	$[+0.029, +0.146]$	0/3
0.40B	6.0843	6.0821	−0.0022	$[-0.032, +0.027]$	1/3
0.80B	5.7894	5.7766	−0.0127	$[-0.056, +0.030]$	3/3
1.20B	5.6291	5.6011	−0.0280	$[-0.064, +0.008]$	3/3
1.60B	5.4257	5.3933	−0.0324	$[-0.046, -0.019]$	3/3
2.00B	5.1954	5.1694	−0.0260	$[-0.041, -0.011]$	3/3
2.40B	4.9196	4.9021	−0.0175	$[-0.031, -0.005]$	3/3
2.80B	4.6519	4.6357	−0.0162	$[-0.046, +0.014]$	3/3
3.20B	4.5319	4.5208	−0.0112	$[-0.030, +0.008]$	3/3

Two points follow for short-horizon comparisons. First, measured at 3.2B tokens the difference is −0.1% at $B = 256$ and −0.2% at $B = 1024$, whereas the same configuration stopped at 2,000 optimizer updates—0.065B tokens at $B = 256$ and 0.26B at $B = 1024—puts AdamW ahead by 2.09% and 3.56%. Second, a fixed update budget also generates an apparent batch-size trend: AdamW’s early advantage is itself larger at the larger batch, and every cell is sampled inside that early window, so a horizon effect is read off as a batch effect. At the full horizon the trend does not survive, because there is no advantage to either method.

The crossover is determined by tokens rather than by updates. It occurs near 0.4B tokens at both batch sizes, which is 12,207 optimizer steps at $B = 256$ but only 3,052 at $B = 1024$. Any explanation resting on optimizer-state warm-up would predict a fixed step count; the data instead tracks data seen, consistent with ZEMA’s evidence weighting becoming useful once gradient noise grows relative to the gradient mean.

Finally, note that $B = 1024$ reaches a lower perplexity than $B = 256$ at matched tokens for both optimizers (4.53 and 4.52 versus 4.59 for both), so the larger batch is not harmful in this setting once training is long enough.

Smoothing rule ablation

The belief EMA in Algorithm 1 applies an update of much smaller magnitude than the per-step belief it smooths: at $B = 256$ the recorded mean absolute belief is 0.50–0.83 while the mean absolute applied update is 0.07–0.085, a factor of 6.5–9.8, and the ratio is largest early in training. This invites the hypothesis that the shrinkage is a defect and that correcting it would remove ZEMA’s early disadvantage. It does not. Table 11 compares the four alternative smoothing rules implemented alongside the paper’s, each with its own learning rate selected from a common six-point grid spanning $2.5 \times10^{-4}$ to $8 \times10^{-3}$ at $B = 256$ and 2,000 updates.

Every alternative is worse, and the ordering is monotone in how aggressively the variant removes the shrinkage: the selected learning rate falls from 0.006 to 0.0005 across the table as the effective update grows, and perplexity degrades in step, with the most thorough corrector worst by 1.6 perplexity. The damping is therefore doing useful work. Averaging bounded beliefs cancels coordinates whose sign disagrees between steps, and the variants that undo it trade that evidence weighting for larger, noisier updates. We keep the post-link EMA of Algorithm 1 on this evidence.

Table 11: Smoothing rule ablation, $B = 256$, 2,000 updates, seed 7. Each variant’s learning rate is selected independently; all optima are interior to the grid searched.

Variant	Rule	Val. PPL	lr
AdamW	—	6.2607	0.006
ZEMA-CLT	EMA of the belief (Algorithm 1)	6.3917	0.006
ZEMA-CLT-z	EMA of $z$, then link	7.1721	0.002
ZEMA-pool	pooled $z$, bias corrected by $\sqrt{1-\beta^{2t}}$	7.3559	0.001
ZEMA-logodds	accumulate log-odds, squash with tanh	7.5795	0.001
ZEMA-kalman	posterior mean, bias corrected by $1-\beta^t$	8.0148	0.0005

Caveats

Intermediate checkpoints are mid-schedule: the cosine is tied to the full 3.2B horizon, so the value at 0.4B in these runs is not what a run terminating at 0.4B with a fully annealed learning rate would produce. The comparison between arms at each checkpoint is unaffected, since both share the schedule, but these intermediate numbers are not comparable to annealed endpoints. The learning rates were selected at 2,000 updates and reused at 97,656; both arms carry the same value at each batch, so the comparison remains fair, but absolute perplexities may be off-optimum. The 3.2B horizon is roughly 7.7 epochs over the 416M-token training split; validation perplexity is still decreasing at the final checkpoint, so neither optimizer has been trained to saturation. Finally, perplexity is specific to the 8,192-token vocabulary used here and is not comparable to published TinyStories perplexities obtained with different tokenizers.

Efficient implementation of ZEMA

§ 3.2 reports ZEMA at a per-step cost of 1.046× AdamW. A direct implementation of Algorithm 1 costs 1.59×. This appendix specifies the three changes that close that gap and the end-to-end runs that validate them. Two of the three are exact rearrangements of the same computation; the third replaces one estimator of $\boldsymbol{\sigma}_t$ with another estimator of the same quantity, at a cost of three degrees of freedom out of 255. Section E.7 then treats subsampling, which is an independent and larger reduction available on top of these.

Where the overhead comes from

Fix one bound linear weight of shape $d_{\mathrm{out}} \times d_{\mathrm{in}}$, and let $L$ be the context length, $B$ the global batch of sequences, and $W$ the number of data-parallel ranks. For example $n$ write $\boldsymbol{x}^{(n)} \in\mathbb{R}^{L \times d_{\mathrm{in}}}$ for the layer input, $\boldsymbol{\delta}^{(n)} \in\mathbb{R}^{L \times d_{\mathrm{out}}}$ for the gradient of the loss with respect to the layer output, and

for that example’s own weight gradient. Algorithm 1 needs the mean $\boldsymbol{m}_t$ and, through line 5, the dispersion

The asymmetry. $\boldsymbol{m}_t$ is the ordinary weight gradient that any backward pass already produces, and costs nothing extra. $M_2$ is a function of the spread across examples, which the usual batched backward destroys the instant it sums over the batch. Every cost below follows from that one fact, and it appears twice. In arithmetic, AdamW’s weight-gradient GEMM contracts over all $B \cdot L$ tokens at once, whereas any per-example statistic must stop at every example boundary. In communication, AdamW’s second moment $v_t = \beta_2v_{t-1} + (1-\beta_2)g_t^2$ is a pure function of the already reduced gradient, which every rank holds after the gradient all-reduce, so it costs no extra bytes; $M_2$ is a function of the pre-reduction spread, which that same reduction has destroyed.

The direct implementation. Written literally, Equation (31) materializes the $[B,d_{\mathrm{out}},d_{\mathrm{in}}]$ per-example gradient tensor, writes it once and reads it at least twice; and because $M_2$ is defined against the global mean, it forces three serialized collectives per step—reduce the mean, agree a common scale, then reduce the rescaled dispersion—of which the second cannot begin until the first has finished. Table 12 gives the resulting ladder.

Configuration	Train time	Tokens/s	vs. AdamW
AdamW	3.572 h	248,815	1.000×
Direct implementation of Algorithm 1	5.643 h	157,496	1.580×
+ bfloat16 statistic accumulation	4.828 h	184,103	1.352×
+ rank-local pooling (§E.2)	4.322 h	205,644	1.210×
+ overlapped reduction (§E.4) and fused per-example kernel (§E.3)†	3.862 h	230,129	1.081×

Configuration	Val. loss	$\Delta$	Perplexity	Verdict
Exact baseline	3.096886	—	22.129	—
+ bfloat16 statistic	3.099870	+0.002984	22.195	pass
+ rank-local pooling	3.100192	+0.003306	22.202	pass
+ fused kernel	3.096976	+0.000090	22.131	pass

Rank-local pooling of the dispersion

The change. Pool the dispersion within each rank only. Partition the batch into the $W$ rank-sized blocks it is already sliced into, write $\bar{\boldsymbol{g}}_w$ for the mean over the $B/W$ examples resident on rank $w$, and decompose the exact sum of squares:

If examples are i.i.d. across ranks, the three terms have expectations $(B-1)\varsigma^2$, $(B-W)\varsigma^2$ and $(W-1)\varsigma^2$ for a per-example variance $\varsigma^2$. Discarding the between-rank term and dividing by $B-W$ rather than $B-1$ therefore leaves an unbiased estimator of the same quantity:

As with subsampling (Equation (35)), the estimand is unchanged—it is still the standard error of the $B$-example mean—and $m_t$ itself is untouched, still the exact global mean formed by an all-reduce over every example. Only the sample used for the dispersion changes, and with it the relative error of $\sigma_t$, which for a $\chi^2$ estimator on $\nu$ degrees of freedom is $1/\sqrt{2\nu}$:

Estimator of $\sigma_t$	$\nu$	rel. error
Exact pooling (Algorithm 1)	255	4.43%
Rank-local pooling, $W=4$	252	4.45%
Subsampling at $k=32$ (§E.7)	31	12.70%

Why it is cheaper. With no between-rank term, nothing in the dispersion depends on the pooled mean, so the three-deep chain collapses to two independent sums:

where $\|$ denotes two collectives issued back to back with no dependency between them. Because each rank’s window holds only $B/W$ examples, the partials can further be carried as raw sums $S_1 = \sum_{n \in\mathcal{W}} g^{(n)}$ and $S_2 = \sum_{n \in\mathcal{W}} (g^{(n)})^2$, converted locally by $M_2 = S_2 - S_1^2 W/B$ with no collective at all, which removes the MAX round that the factored form needs for underflow protection. The cancellation this invites is bounded by the identity $S_2 = M_2 (1 + z^2/(B/W - 1))$: at $z_{\max} = 10$ and $B/W = 64$ the worst attainable ratio is 2.6, under one bit of float32, and at the $\lvert z\rvert\approx0.51$ these runs actually exhibit it is 1.004.

Assumptions, and when they fail. Two conditions are load-bearing. First, unbiasedness in Equation (33) requires the between-rank term to be pure noise. Here the global permutation is drawn once per step and then sliced across ranks, so rank membership carries no information and the condition holds by construction; an implementation that instead gave each rank its own shard, or sorted examples by length, would discard real between-rank spread, report $\sigma_t$ low, $\lvert z_t\rvert$ high, and be overconfident. Second, the raw-sum form assumes no underflow. It is safe at pretraining scale—the fraction of coordinates reporting zero variance is exactly 0 on every run reported here—but not universally: LoRA initializes one factor at zero, producing gradients of order $10^{-22}$ whose squares underflow float32 to exactly zero and read as unanimous agreement among examples. The fine-tuning experiments of § 3.1 therefore retain the factored form and its MAX round.

A fused per-example kernel

Why the ordinary GEMM structure does not apply. A matrix-product kernel tiles the output $[d_{\mathrm{out}}, d_{\mathrm{in}}]$ and streams the whole $B \cdot L$ contraction into a register accumulator that is written once. The per-example square in Equation (31) requires the sum over the $L$ token positions to complete at every example boundary, so the accumulator must be finalized $B$ times rather than once—which is why the direct path falls back to a batched product and materializes $[B, d_{\mathrm{out}}, d_{\mathrm{in}}]$.

The change. Reorder the contraction so that the example index is the outer loop and finalize in registers, never writing the per-example tensor (Algorithm 2). Each program owns a disjoint tile of the output, so the final accumulation is a plain read–modify–write rather than an atomic. Operands are loaded in their own precision with a float32 accumulator—the same arithmetic the ordinary weight gradient performs—so tensor cores are retained; upcasting both operands before the product instead forces the float32 CUDA-core path and is 4.7–6.5× slower than the plain weight gradient.

It replaces work rather than only adding it. $S_1$ is the weight gradient, so the kernel subsumes the backward pass’s own weight-gradient product for that layer, which must then be suppressed. On the setting of § 3.2 the kernel costs 181 ms of device time per step while removing a 148 ms product, a marginal cost of about 33 ms; an accounting that prices it as pure addition overstates it roughly fivefold. The extra arithmetic proper is $B\,d_{\mathrm{out}}d_{\mathrm{in}}$ squares against $B L\,d_{\mathrm{out}}d_{\mathrm{in}}$ multiply–accumulates, smaller by a factor of $L$. The binding constraint is registers: three float32 tiles are live at once, which caps tile size well below what a plain GEMM would choose.

When it pays. The saving is the traffic of the $[B, d_{\mathrm{out}}, d_{\mathrm{in}}]$ tensor, weighed against operands that are read either way, so the governing ratio is

Algorithm 2: Fused per-example first and second sums, for one bound weight and one rank. Replaces line 3 of Algorithm 1; $S_1$ is the rank’s weight gradient and $S_2$ feeds Equation (33). The $[B,d_{\mathrm{out}},d_{\mathrm{in}}]$ tensor is never formed.

Input: activations $x$, output sensitivities $\delta$; tile sizes $(b_{\mathrm{out}}, b_{\mathrm{in}}, b_L)$

Output: $S_1,S_2$

1. foreach tile $T$ of size $b_{\mathrm{out}} \times b_{\mathrm{in}}$ of the output, in parallel do

   2. $s_1 \leftarrow0$; $s_2 \leftarrow0$ // float32, in registers

   3. for $n \leftarrow1$ to $B$ do

      4. $p \leftarrow0$ // this example’s own gradient tile

      5. for each block of $b_L$ token positions do

         6. $p \leftarrow p+\delta_{\mathrm{block}}^{(n)\mathsf{T}}x_{\mathrm{block}}^{(n)}$ // tensor-core product

      6. end

      7. $s_1 \leftarrow s_1+p$; $s_2 \leftarrow s_2+p^2$ // finalize at the example boundary

   4. end

   5. write $s_1,s_2$ into $T$ of $S_1,S_2$

2. end

Here $r=0.19$–$0.27$, but $r>1$ whenever the context is short and the layer wide—1.94 for a [8192, 256] head at $L=128$. Short context and wide layers favour the kernel; long context and narrow layers do not.

Overlapping the reduction with the backward pass

With gradient accumulation, a weight’s $S_1$ and $S_2$ are complete only once the last micro-batch’s backward has passed it. Reducing everything after the accumulation window closes leaves the collective fully exposed at the end of the step. But weights become final in reverse layer order, so the reduction for the last layer can be launched while the earlier layers are still computing. Grouping weights into byte-sized buckets and launching each bucket asynchronously as it fills—the scheme DistributedDataParallel uses for gradients—recovers 59 ms per step here. It changes only when bytes move, not what is computed.

Two conditions are required. Bucket boundaries must be identical on every rank, or ranks issue mismatched collectives and deadlock; this holds because weights complete in the same deterministic order everywhere and buckets fill by the same byte rule. And the scheme requires rank-local pooling: under the exact three-round chain of §E.2, round two depends on round one and round three on round two, so only the first could be launched early. Removing the between-rank term is what makes both payloads independent, and therefore launchable as soon as they exist.

A variant we discarded. One can go further and use the previous step’s dispersion with the current step’s mean, leaving the current dispersion in flight across the step boundary. We do not recommend it. Once the kernel and the bucketing above had shrunk the surrounding work, bucketed overlap alone hid the dispersion collective essentially completely, and the stale variant measured −1 ms—within noise—while costing a persistent $[d_{\mathrm{out}},d_{\mathrm{in}}]$ float32 buffer, +0.32 GiB, which is precisely the footprint of AdamW’s $v$. ZEMA’s claim is one state against Adam’s two; that variant buys a second full-size state for no speed.

What remains, and why it cannot reach zero

After all three changes, ZEMA moves two weight-shaped payloads per step on the bound parameters where AdamW moves one:

Per step, bound parameters	Bytes	vs. AdamW
AdamW (gradient only)	648,735,744	1.00×
ZEMA, rank-local (mean + dispersion + auxiliary)	988,474,368	1.52×

One of the two is not a premium: the mean is the gradient, the same bytes data-parallel training moves anyway. The second has no AdamW counterpart and cannot acquire one, for the reason given in §E.1: $v_t$ is recoverable from the reduced gradient and $M_2$ is not. 1.52× of AdamW’s collective traffic is therefore the floor for this statistic as specified, not a target to optimize past. Reducing it further means computing a different statistic—which is what subsampling (§E.7) does, and it is priced in degrees of freedom accordingly.

Validation

Of the four changes in Table 12, two alter the numbers the optimizer sees—bfloat16 accumulation of the statistic, through rounding, and rank-local pooling, through three degrees of freedom—and two do not. The fused kernel computes Equation (30) exactly in exact arithmetic and is pinned against a batched-product reference in unit tests; bucketed overlap changes only scheduling.

Both numerical changes were carried to a complete run—6,103 steps and 3.2B tokens, ending in a full validation pass over the 9.2M held-out tokens—rather than to a short smoke, and compared against the exact baseline with a tolerance of ±2$\varsigma$, where $\varsigma$ = 0.002189 nats is that baseline’s own three-seed standard deviation.

Short runs are not a substitute. A rounding change and a three-degree-of-freedom change both produce drift that is invisible early and only interpretable against a same-schedule multi-seed band, which exists only at the end of training. A 20-step smoke of this configuration reported a negative-variance fraction of 0.142 where the completed runs give $3.9 \cdot10^{-9}$, because the smoke’s compressed warmup pooled 8 examples where the real schedule pools 256.

Measured throughput. Table 14 gives per-step training times for the two endpoints. AdamW’s own step time varies by about 7% across nodes of this cluster, so each block is one job on one node and ratios are taken within a block, never across. Peak memory is below AdamW’s in every ZEMA row, because ZEMA carries one persistent state tensor on the bound parameters where AdamW carries two.

Reconciling the two ratios. Table 12 ends at 1.081× and Table 14 at 1.046×. Both are training-time measurements and the difference is configuration, not metric. The complete runs of Table 12 exist to gate the two numerical changes, and were launched before the implementation settled: that final row still carries the stale-dispersion variant of §E.4 and predates a set of host-side cleanups—chiefly the removal of 672 per-step device-to-host synchronizations—together worth 90 ms per step, taking 2,278 ms down to the 2,188 ms of Table 14. The configuration we recommend has been measured per step but not carried to a complete run. We therefore report 1.046× as its per-step cost and 1.081× as the only ratio backed by a full 3.2B-token run, rather than quoting one number for both.

Diagnostics for a reimplementation. Three quantities are worth monitoring. The fraction of coordinates reporting zero variance must be 0; a nonzero value means the raw-sum cancellation bound of §E.2 has failed for that regime and the factored form is required. The mean $|z|$, about 0.51 here, sets the realized cancellation ratio and should be checked before trusting the raw-sum path. And the assignment of examples to ranks must be verified to be exchangeable, or Equation (33) is biased rather than merely wider. Kernel tile sizes are register-bound and shape-dependent, so they should be retuned per architecture—after computing Equation (34), which decides whether the kernel is worth using at all.

Subsampling the per-batch variance

The standard error $\sigma_t$ is the only statistic in ZEMA that needs per-example gradients, and it does not need all of them: the mean can be taken over the whole batch while the variance is estimated from a fraction of it. Figure 7 sweeps that fraction on the setting of § 3.2.

Down to $k = 8$ of $B = 256$ examples both arms stay within seed noise of their whole-batch result, and at every $k$ we tried they remain below AdamW. Details of the estimator and the Algorithm 3 can be found below.

The estimator. The mean $m_t$ and the standard error $\sigma_t$ need not be estimated from the same examples. We keep the mean over all $B$ examples and estimate the dispersion from a subset $S$ of size $k$:

The divisor is $B(k - 1)$, not $k(k - 1)$. The sum over $S$ divided by $k - 1$ estimates the per-example variance from $k$ draws; dividing by $B$ instead rescales that to the standard error of the $B$-example mean, which is the quantity $z_t$ is defined against. The estimand is therefore unchanged—only the sample used to estimate it shrinks, and with it the relative error of $\sigma_t$, which grows as roughly $1/\sqrt{2(k - 1)}$.

This is not a smaller batch. Every example still enters $m_t$, so the direction the optimizer follows is formed from the full batch, and the $n$ behind the standard error remains $B$. A batch of $k$ would change both.

Why it is cheaper. Algorithm 3 gives the replacement. The saving is in traffic, not arithmetic. The full path materializes a $[B,\mathrm{out},\mathrm{in}]$ per-example gradient tensor and makes several elementwise passes over it; here the mean comes from a single matrix product over all $B \cdot T$ tokens—the same contraction, and the same floating-point operations, as the ordinary weight gradient—and only $k$ examples ever enter the batched product that forms per-example gradients. Arithmetic rises by a factor $(B+k)/B$; the intermediate tensor shrinks by $B/k$.

Algorithm 3: Subsampled variance. Replaces lines 3–5 of Algorithm 1; everything from $z_t$ onwards is unchanged.

Input: batch size $B$; variance sample size $k \le B$; activations $\mathbf{x}$ and output sensitivities $\boldsymbol{\delta}$ for the batch

1. sample a batch of size $B$

2. $m_t \leftarrow\frac{1}{B}\boldsymbol{\delta}^{\mathsf{T}}\mathbf{x}$ // one product over all $B \cdot T$ tokens; no per-example tensor

3. $S \leftarrow$ the first $k$ examples of the batch // the batch is shuffled, so $S$ is a random draw from the data

4. form per-example gradients $\mathbf{g}_t^{(j)}$ for $j \in S$ only // a $[k,\mathrm{out},\mathrm{in}]$ tensor, not $[B,\mathrm{out},\mathrm{in}]$

5. $\bar{\mathbf{g}}_S \leftarrow\frac{1}{k}\sum_{j\in S}\mathbf{g}_t^{(j)}$ // centre on the subsample’s own mean

6. $\boldsymbol{\sigma}_t \leftarrow\left[\frac{1}{B(k-1)}\sum_{j\in S}\left(\mathbf{g}_t^{(j)}-\bar{\mathbf{g}}_S\right)^2\right]^{1/2}$ // still the standard error of the $B$-example mean

Results. Figure 7 sweeps $k$ from 256 down to 4, one seed per $k$. From $k = 128$ down to $k = 8$—a $32\times$ reduction in the per-example tensor—ZEMA-CLT moves between 22.22 and 22.28 against its whole-batch 22.18, and ZEMA-Adam between 22.30 and 22.45 against 22.36; in both cases the excursion is at or below the seed-to-seed spread at that point, which is 0.05 and 0.09 in standard deviation over the three whole-batch seeds. Only at $k = 4$ does a penalty appear that is clearly bigger than noise, and it remains small: ZEMA-CLT rises to 22.49 and ZEMA-Adam to 22.53. At every $k$, including $k = 4$, both arms sit below AdamW’s whole-batch mean of 22.79, though at $k = 4$ they are inside AdamW’s 95% interval and we do not claim a separation there.