Backpropagated Output Momentum: Relocating Optimizer History from Parameters to Task Space

Yuchen Li; Zongqi Fan; Nguyen H. Tran; Ken-Tye Yong

cs.LGSep 29, 2026

Backpropagated Output Momentum: Relocating Optimizer History from Parameters to Task Space

Authors: Yuchen Li, Zongqi Fan, Nguyen H. Tran, Ken-Tye Yong

Organizations: University of Sydney

Abstract

Optimizer momentum is usually stored as a parameter-sized moving average of past gradients, which makes history costly and fixes each past signal in the coordinates in which it was computed. We introduce Backpropagated Output Momentum (BOM), which instead stores a compact moving average of prediction errors at the model output and reprojects that history through the current network at every step. A batch-level analysis characterizes the information retained and omitted by this relocation, while the implementation preserves the current supervised gradient and can replace the first-moment component of several adaptive optimizers. As a plug-in for momentum-based optimizers, including ones that already compress their state, BOM reduces parameter-shaped optimizer state by 49.7-99.8% in three compositions and, averaged over three language backbones, paired step time by 4.0%. It also improves mean validation performance across language and vision fine-tuning, by 1.42 points in the primary five-task comparison. Language and vision pretraining studies, together with matched mechanism controls, further test the construction across output spaces and model scales.

Figures & tables

Panel of Figure 1 — Figure 1: Where first-order history lives. AdamW projects each $s_{\tau}$ through its own step’s model and sampled examples before averaging, storing historical projections inside $m_{t}$ ; BOM averages the task signals first and applies one projection through the current model and current examples, refreshing the same weighted history at every step (covariance terms omitted; Section 3.3 is exact).

Backbone	AdamW	BOM	$\Delta$ Score	$\Delta$ Time
RoBERTa-base	0.7962	0.8104	+0.0142	$-0.7\%$
DeBERTa-v3-base	0.8329	0.8467	+0.0138	$-4.8\%$
Qwen3-1.7B	0.8169	0.8304	+0.0135	$-6.4\%$
Coverage	3/3 backbones: higher mean score, lower mean step time

Table 1: Primary epoch-5 NLP results. Scores average five tasks and five evaluation seeds; taskwise intervals are in Appendix A.1 .

\Delta

Time is the mean paired bf16 step-time change of BOM relative to AdamW; with FP32 master weights, RoBERTa-base changes by

-4.3\%

Base optimizer	Base	With BOM	$\Delta$	Base	With BOM	$\Delta$
	Five-task macro			Optimizer state (MiB)		Step time
AdamW	0.7962	0.8104	$+0.0142$	475.5	237.7	$-0.7\%$
Adam-mini	0.7659	0.7782	$+0.0123$	238.2	0.42	$-2.0\%$
GaLore	0.7739	0.7854	$+0.0114$	155.7	78.3	$-3.1\%$

Table 2: Pairwise BOM composition on RoBERTa-base: five-seed epoch-5 macros, parameter-shaped optimizer state (excluding the

O(d_{\mathrm{out}})

vector), and step-time changes paired within each base’s implementation stack. Each base keeps a parameter-local first-moment EMA, which BOM replaces; Adam-mini shares

v

within blocks, and GaLore keeps its moments in a low-rank subspace. Macro-margin 95% CIs: AdamW

[+0.0027,+0.0258]

, Adam-mini

[-0.0203,+0.0449]

, GaLore

[-0.0078,+0.0307]

. All entries are measured in bf16; with FP32 master weights, the AdamW step-time change is

-4.3\%

. Details: Appendices A.3 , B.3.1 and B.8 .

Control	Macro	$\Delta$	Control	Macro	$\Delta$
CE, no local $m$	$0.7841$	$-2.63$	Own-step projection	$0.7807$	$-2.97$
Mixed + local $m$	$0.7806$	$-2.98$	BOM numerator, CE $v$	$0.7892$	$-2.12$
AdamW, $\beta_{1}=0.45$	$0.8008$	$-0.96$	Classifier-only history	$0.7910$	$-1.94$
Full-gradient kernel	$0.7977$	$-1.28$	$\beta_{1}=0$ (RTE acc.)	$0.5235$	$4/5$ deg.
APM anchor	$0.7934$	$-1.70$	Drift $\cos$ head/enc.	$0.81$ / $0.52$	–

Table 3: RoBERTa-base mechanism controls (epoch-5 macro, five seeds;

\Delta

from BOM ’s

0.8104

). Details: Table 13 .

Fine-tuning (vision), 40 epochs: validation accuracy $\uparrow$
Setting	AdamW	BOM	$\Delta$
STL10, ConvNeXt-Tiny	0.9647	$\mathbf{0.9769}$	$+0.0122$
STL10, ViT-Tiny	0.9567	$\mathbf{0.9669}$	$+0.0102$
Pretraining (vision), 90 epochs: epoch-90 validation top-1 $\uparrow$
ImageNet-1k, ResNet-50	$0.7612\pm 0.0005$	$\mathbf{0.7646\pm 0.0008}$	$+0.0035$
Pretraining (language): final evaluation loss $\downarrow$

Table 4: Matched AdamW– BOM cross-domain results; each group states its metric and

\Delta

is the paired BOM

-

AdamW mean. Five consecutive-seed reporting runs except C4 440M and 1.1B (

n=3

each); pretraining entries are mean

\pm

sample sd. Configurations: Appendices B.4 and B.5 ; paired intervals: Tables 11 and 12 .

Panel of Figure 2 — Figure 2: Pretraining resource efficiency across scales. (a) C4 end-to-end time ratios (paired mean $\pm$ SD), H800 single-step ratios, and the calibrated cost model (Appendix B.8.1 ). (b) Peak allocated-memory savings with the FP32 state reference $4(P-V)$ bytes. (c) Checkpoint sizes with FP32 predictions $12P$ (AdamW) and $8P$ ( BOM ) bytes. Details: Table 32 , Figure 9 , Appendix B.8 .

Appendix figures & tables49 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Optimizer-state placement in AdamW and BOM. AdamW retains dense first- and second-moment tensors for every trainable layer. BOM deletes every dense local first moment (crossed out), retains the local second moment used for adaptive scaling, and stores one compact task-space history $q_{t}$ . The current backward graph re-expresses that history as updates for all trainable tensors, rather than preserving the parameter coordinates induced by earlier Jacobians. Parameter, gradient, and activation memory are omitted because the diagram isolates optimizer state.

Backbone	Dataset	Metric	Method	Train score	Val score	Train loss	Val loss
RoBERTa-base	CoLA	MCC	AdamW	0.8566	0.5743	0.1851	0.4777
RoBERTa-base	CoLA	MCC	BOM	0.9641	0.5970	0.0495	0.5794
RoBERTa-base	MRPC	F1	AdamW	0.9609	0.9077	0.1324	0.3240
RoBERTa-base	MRPC	F1	BOM	0.9862	0.9191	0.0616	0.3305
RoBERTa-base	QQP	F1	AdamW	0.8501	0.8463	0.2620	0.2683
RoBERTa-base	QQP	F1	BOM	0.9004	0.8626	0.1960	0.2481

Appendix

Table 5: Per-dataset results for AdamW and BOM on the three pretrained backbones, at the epoch-5 endpoint over the five evaluation seeds. Score is the metric named in each row; loss is cross-entropy.

(a) Task-level paired margins
Dataset	Metric	$\Delta$	95% CI
CoLA	MCC	$+0.0226$	$[-0.0226,\,+0.0679]$
MRPC	F1	$+0.0114$	$[-0.0079,\,+0.0306]$
QQP	F1	$+0.0163$	$[+0.0059,\,+0.0267]$
RTE	Acc	$+0.0195$	$[+0.0074,\,+0.0316]$
SST-2	Acc	$+0.0014$	$[-0.0116,\,+0.0144]$

Appendix

Table 6: Paired BOM

-

AdamW validation-score margins on RoBERTa-base at the epoch-5 endpoint over the five evaluation seeds. Panel (a) reports the mean paired difference in each task’s native metric with its descriptive paired 95% interval; panel (b) reports the five-task macro comparison of Section 4.3 . Intervals use

t_{0.975,4}=2.776

on the five paired differences. Seeds are matched across the two arms, so the paired spread is the relevant one: on RTE each arm’s own standard deviation (

0.021

and

0.029

) exceeds the margin, while the paired standard deviation is

0.0098

. Per-seed values are in Table 7 .

Task	Method	A	B	C	D	E	Mean
RoBERTa-base
CoLA (MCC)	AdamW	0.5435	0.5778	0.6131	0.5834	0.5539	0.5743
	BOM	0.5885	0.5536	0.6107	0.6111	0.6209	0.5970
$\Delta$ (95% CI): $+0.0226$ $[-0.0226,\,+0.0679]$
MRPC (F1)	AdamW	0.9004	0.9127	0.9171	0.9084	0.8998	0.9077
	BOM	0.9255	0.9308	0.9020	0.9206	0.9164	0.9191

Appendix

Table 7: Per-seed validation scores for the primary AdamW– BOM comparison on all three backbones, at the epoch-5 endpoint. Score is the task-appropriate metric named with each block; the five evaluation seeds are disjoint from the two HPO seeds. Labels A–E identify the matched evaluation runs, and the Mean column is the value reported in Table 5 .

\Delta

is the mean paired difference with its descriptive 95% interval (

t_{0.975,4}=2.776

on the five paired differences); intervals use the variation of matched seed differences, which need not be smaller than either arm’s marginal variation.

Panel of Figure 4 — Figure 4: RoBERTa-base control effects relative to AdamW at the epoch-5 endpoint (Table 8 ). Each point is the mean method-minus-AdamW difference; bars are descriptive, unadjusted 95% paired $t$ confidence intervals over five evaluation seeds ( $df=4$ ). For each seed, we first average the task-appropriate score or cross-entropy across the five tasks, then pair the resulting macro with AdamW at the same seed. Score differences are multiplied by 100; positive values favor the method. Negative cross-entropy differences favor the method. The dashed line marks no difference. These intervals describe seed variation conditional on the selected configurations and this fixed task set.

Method	Train score	Val score	Train loss	Val loss
BOM -AdamW	0.9447	0.8104	0.1324	0.3826
AdamW	0.8832	0.7962	0.2235	0.3714
Lion	0.9302	0.7965	0.1340	0.4340
Muon	0.8833	0.7785	0.2029	0.5628
GaLore	0.8134	0.7739	0.3128	0.3873
Adam-mini	0.7967	0.7659	0.3406	0.3826

Appendix

Table 8: RoBERTa-base five-task macro endpoints for the seven methods in the fixed-backbone comparison, at the epoch-5 endpoint over the five evaluation seeds. Parameter-shaped optimizer state for the substitution pairs is reported with the composition comparison (Table 2 ).

Dataset	Metric	Method	Train score	Val score	Train loss	Val loss
CoLA	MCC	AdamW	0.8566	0.5743	0.1851	0.4777
CoLA	MCC	BOM -AdamW	0.9641	0.5970	0.0495	0.5794
CoLA	MCC	Lion	0.9173	0.5780	0.0945	0.5178
CoLA	MCC	Adam-mini	0.6403	0.5348	0.3614	0.4639
CoLA	MCC	BOM -Adam-mini	0.7098	0.5478	0.3197	0.4568
CoLA	MCC	GaLore	0.6668	0.5495	0.3521	0.4386

Appendix

Table 9: RoBERTa-base per-dataset results for every optimizer control and every BOM composition, at the epoch-5 endpoint over the five evaluation seeds. Score is the metric named in each row; loss is cross-entropy. Train quantities are measured on a fixed probe subset of the training set.

Panel of Figure 5 — Figure 5: RoBERTa-base training and validation trajectories for the six baseline optimizers and BOM -AdamW. Each curve averages the five evaluation seeds within each task and then gives the five tasks equal weight. Training scores and losses use the training probe; validation scores use the task-appropriate metric and losses are cross-entropy. The logged intermediate evaluations describe training progress; the reported comparison uses the fixed epoch-5 endpoints in Table 8 .

Panel of Figure 6 — Figure 6: Pairwise effects of substituting BOM for the first-moment component of AdamW, Adam-mini and GaLore on RoBERTa-base. Each arrow starts at the base optimizer (open circle) and ends at its BOM composition (filled square); movement up and left denotes a higher five-task validation macro and less parameter-shaped optimizer state. The annotations give each method’s mean paired bf16 step-time change over its main comparison: the three NLP backbones for AdamW, and RoBERTa-base for Adam-mini and GaLore. All three compositions move up and left and have lower mean step time, so each composition Pareto-dominates its own base on the three reported mean estimates. Averaged over the three methods’ bf16 main comparisons, adding BOM raises the macro by $1.25$ points and lowers paired step time by $3.0\%$ . State bytes use the common accounting convention. Scores are epoch-5 means over five evaluation seeds. Full estimates and paired intervals are in Tables 2 and 30 .

Method	Epoch at reference training score
AdamW	5.00
BOM	2.94

Appendix

Table 10: Training-probe score progress on RoBERTa-base, macro-averaged over five tasks and five evaluation seeds. The reference is AdamW’s epoch-5 training-probe macro of

0.8832

. The first BOM crossing is linearly interpolated between adjacent half-epoch measurements. This diagnostic uses only training-probe scores; the validation evaluations logged every half epoch are shown descriptively in Figure 5 and do not enter it, and the reported validation comparison uses the fixed epoch-5 endpoint.

Backbone	Method	Metric	Train score	Val score	Train loss	Val loss
ConvNeXt-Tiny	AdamW	Acc	0.9914	0.9647	0.0298	0.1487
ConvNeXt-Tiny	BOM	Acc	1.0000	0.9769	0.0007	0.0855
ViT-Tiny	AdamW	Acc	0.9930	0.9567	0.0255	0.1711
ViT-Tiny	BOM	Acc	0.9992	0.9669	0.0015	0.1262

Appendix

Table 11: STL10 transfer results for AdamW and BOM on two pretrained vision backbones, at the 40-epoch endpoint over the five evaluation seeds. Both arms share the same 40-epoch budget (Table 40 ) and are read under the same endpoint convention within each backbone pair. Score is accuracy; loss is cross-entropy. The paired BOM

-

AdamW mean accuracy differences and descriptive 95%

t

intervals are

+0.0122\;[-0.0052,+0.0297]

for ConvNeXt-Tiny and

+0.0102\;[-0.0093,+0.0298]

for ViT-Tiny. Intervals use the sample standard deviation of five matched seed differences and

t_{0.975,4}=2.776

, without multiplicity correction.

Setting	Metric	AdamW	BOM	$\Delta$
ImageNet-1k, ResNet-50	Epoch-90 val top-1 $\uparrow$	$0.7612\pm 0.0005$	$\mathbf{0.7646\pm 0.0008}$	$+0.0035$
	Val loss at that epoch $\downarrow$	$0.9795\pm 0.0027$	$\mathbf{0.9530\pm 0.0011}$	$-0.0265$
Python code, Qwen3-55M	Final eval loss $\downarrow$	$1.7353\pm 0.0093$	$\mathbf{1.7337\pm 0.0106}$	$-0.0016$
	Perplexity $\downarrow$	$5.6706\pm 0.0529$	$\mathbf{5.6619\pm 0.0603}$	$-0.0087$
FineWeb-Edu, Qwen3-55M	Final eval loss $\downarrow$	$\mathbf{3.2931\pm 0.0092}$	$3.2948\pm 0.0045$	$+0.0017$
	Perplexity $\downarrow$	$\mathbf{26.9262\pm 0.2501}$	$26.9723\pm 0.1221$	$+0.0460$

Appendix

Table 12: Aggregate from-scratch pretraining endpoints, reported as mean

\pm

sample sd over five consecutive-seed reporting runs except the C4 440M and 1.1B studies (

n=3

each).

\Delta

is the paired mean difference ( BOM

-

AdamW). ImageNet-1k reports epoch-90 top-1 and validation loss. For the language corpora, perplexity is computed as

\exp(\mathrm{loss})

for each run before averaging. All language rows use the final evaluation. Descriptive paired 95% intervals for the C4 loss differences are

[-0.0226,+0.0065]

at 55M,

[-0.0766,-0.0329]

at 110M, and

[-0.0199,+0.0111]

at 440M. Budgets and schedules are in Table 42 .

Panel of Figure 7 — Figure 7: Qwen3-55M validation-loss trajectories on Python code and FineWeb-Edu. Each curve is the mean over five runs and each band is $\pm 1$ sample standard deviation; validation cross-entropy is evaluated on held-out batches. C4 validation-loss trajectories across four scales are reported separately in Figure 8 .

Panel of Figure 8 — Figure 8: C4 validation-loss trajectories at four model and token-budget scales. Curves are means and bands are $\pm 1$ sample standard deviation over five consecutive-seed reporting runs at 55M/110M and three at 440M/1.1B. The 1.1B setting uses $10^{-3}$ in both arms. Configurations are in Table 42 ; summary losses and perplexities are in Table 12 .

Question	Control	Reported quantity	Comparison / diagnostic	Reading
Is it buffer removal alone?	CE + no local $m$	$0.7841$	$-2.63$	no degenerate seeds
Is it the objective alone?	Mixed + local $m$	$0.7806$	$-2.98$	below AdamW
Is it just the current/history ratio?	AdamW $\beta_{1}=0.45$ ( $0.55/0.45$ )	$0.8008$	$-0.96$	vs. $\beta_{1}=0.9$ : CI includes zero
Is historical projection drift real, and where?	$\cos(\hat{M}_{t},h_{t})$ , head / encoder	$0.81$ / $0.52$	–	rel. drift $0.46$ / $0.79$
Does a local scalar anchor help?	APM : per-tensor anchor, $+O(L)$ scalars	$0.7934$	$-1.70$	lower mean score
What changes under current reprojection?	same history, own-step projection	$0.7807$	$-2.97$	matched scaling rule; observed drop

Appendix

Table 13: RoBERTa-base mechanism controls. Most rows report the five-task epoch-5 macro over five seeds and their difference from BOM (

0.8104

); rows 4 and 7 report labeled diagnostics. The classifier-only row uses its protocol-selected configurations; rescue searches are reported in Appendix A.6.6 . APM is the anchored per-tensor variant.

Construction	First-order history	Update numerator	Second-moment source
AdamW	Parameter-space full gradients	$\widehat{\operatorname{EMA}}(g)_{t}$	$g_{t}$
CE + no local $m$	None	$g_{t}$	$g_{t}$
Mixed + local $m$	Task residuals and parameter-space mixed updates	$\widehat{\operatorname{EMA}}(u)_{t}$	$u_{t}$
BOM	Task-space residual EMA $\hat{q}_{t}$	$u_{t}$	$u_{t}$
Own-step	Parameter-space mean-residual projections $\hat{M}_{t}$	$(1-\lambda)g_{t}+\lambda\hat{M}_{t}$	Own numerator
Full temporal kernel	Parameter-space full gradients $\hat{m}_{t}$	$(1-\lambda)g_{t}+\lambda\hat{m}_{t}$	Own numerator

Appendix

Table 14: Update constructions behind the mechanism comparisons.

g_{t}

is the full CE gradient,

u_{t}=(1-\lambda)g_{t}+\lambda\bar{J}_{t}^{\top}\hat{q}_{t}

is the BOM numerator, and

\widehat{\operatorname{EMA}}

includes bias correction. The final column is the vector whose elementwise square feeds the second moment. The own-step buffer averages

\bar{J}_{\tau}^{\top}s_{\tau}

and omits historical covariance; the full-kernel buffer averages

g_{\tau}

and retains it. Each independently trained arm develops its own trajectory and adaptive statistics.

Dataset	Metric	AdamW	Mixed + local $m$	CE + no local $m$	BOM
CoLA	MCC	0.5743	0.5864	0.5635	0.5970
MRPC	F1	0.9077	0.9016	0.9011	0.9191
QQP	F1	0.8463	0.8452	0.8547	0.8626
RTE	Acc	0.7155	0.6318	0.7069	0.7350
SST-2	Acc	0.9372	0.9378	0.8945	0.9385
Macro mean	–	0.7962	0.7806	0.7841	0.8104

Appendix

Table 15: RoBERTa-base objective–momentum control. Entries are mean task-appropriate validation scores at epoch 5 over all five evaluation seeds. Both controls receive the same HPO budget as the primary pair; selected rates are in Table 43 . Every condition retains parameter-local adaptive scaling and decoupled weight decay.

Varied coefficient	Value	Fixed coefficient	MRPC F1	$\Delta$	SST-2 Acc.	$\Delta$
$\lambda$	0.25	$\beta_{1}=0.9$	$0.9187\pm 0.0075$	$-0.0004$	$0.9353\pm 0.0019$	$-0.0032$
$\lambda$	0.75	$\beta_{1}=0.9$	$0.9156\pm 0.0076$	$-0.0035$	$0.9385\pm 0.0039$	$0.0000$

Appendix

Table 16: Post-selection local sensitivity of BOM ’s fixed mixing coefficient

\lambda

on RoBERTa-base. Entries are epoch-5 validation mean

\pm

population standard deviation over the five evaluation seeds;

\Delta

is relative to the paper configuration

\lambda=0.5,\ \beta_{1}=0.9

, whose scores are

0.9191

on MRPC and

0.9385

on SST-2. Only non-default settings are shown.

$\beta_{1}$	score	deg.	score	score	deg.	score	score
	RTE (Acc.)		MRPC (F1)	SST-2 (Acc.)		CoLA (MCC)	QQP (F1)
$0.00$	$0.5235$	4/5	$0.9057$	$0.8505$	1/5	$0.5848$	$0.8590$
$0.80$	$0.6383$	2/5	$\mathbf{0.9210}$	$0.9378$	0/5	$0.5890$	$0.8581$
$0.90$	$0.7350$	0/5	$0.9191$	$\mathbf{0.9385}$	0/5	$\mathbf{0.5970}$	$\mathbf{0.8626}$
$0.99$	$\mathbf{0.7437}$	0/5	$0.9149$	$0.9358$	0/5	$0.5897$	$0.8529$

Appendix

Table 17: Task-space EMA decay

\beta_{1}

across the five RoBERTa-base tasks, with the

0.5/0.5

current/history mixture, the selected learning rates and all remaining coefficients held fixed. Entries are epoch-5 validation means over the five evaluation seeds; deg. counts seeds whose endpoint sits at the majority-class solution (

0.5271

on RTE,

0.5092

on SST-2) with validation cross-entropy at chance, and no seed of any other task is degenerate at any setting. The

\beta_{1}=0.9

row is the sweep anchor; the best entry in each column is bold.

Validation score
Dataset	Metric	$\beta_{1}=0.9$	$\beta_{1}=0.45$	$\Delta$
CoLA	MCC	$0.5743$	$0.5827$	$+0.0083$
MRPC	F1	$0.9077$	$0.9039$	$-0.0038$
QQP	F1	$0.8463$	$0.8577$	$+0.0115$
RTE	Acc.	$0.7155$	$0.7191$	$+0.0036$
SST-2	Acc.	$0.9372$	$0.9408$	$+0.0037$

Appendix

Table 18: AdamW-only current/history-ratio control on RoBERTa-base. The main-experiment AdamW configuration is unchanged except for

\beta_{1}

, which moves from

0.9

0.45

and changes the asymptotic current/past total weights from

0.10/0.90

0.55/0.45

. This matches the mass split of BOM ’s batch-mean component, while the past tail decays at

0.45

rather than

0.9

. Per-dataset entries are epoch-5 means over the five evaluation seeds; the macro row averages the five task means. Score is the task-appropriate validation metric, loss is validation cross-entropy, and

\Delta

\beta_{1}=0.45

minus

\beta_{1}=0.9

Task	Metric	AdamW	Full kernel	BOM	$\Delta$	95% CI
CoLA	MCC	0.5743	0.5625	0.5970	$+0.0345$	$[+0.0053,+0.0636]$
MRPC	F1	0.9077	0.9209	0.9191	$-0.0018$	$[-0.0159,+0.0122]$
QQP	F1	0.8463	0.8584	0.8626	$+0.0042$	$[-0.0030,+0.0114]$
RTE	Acc.	0.7155	0.7235	0.7350	$+0.0116$	$[-0.0376,+0.0607]$
SST-2	Acc.	0.9372	0.9232	0.9385	$+0.0154$	$[-0.0115,+0.0423]$
Macro	–	0.7962	0.7977	0.8104	$+0.0128$	$[+0.0066,+0.0190]$

Appendix

Table 19: Full temporal-kernel control on RoBERTa-base under the unified HPO policy (Appendix B.6 ). Scores are five-seed means from the one-shot final test after epoch 5.

\Delta

is BOM minus the full-kernel control; intervals are descriptive paired 95%

t

intervals (

df=4

), without multiplicity correction. The macro interval averages tasks within each seed before forming paired differences.

Dataset	Metric	AdamW	Own-step control	BOM	$\Delta$	Paired 95% CI
CoLA	MCC	0.5743	0.5632	0.5970	$+0.0338$	$[-0.1003,+0.1678]$
MRPC	F1	0.9077	0.9162	0.9191	$+0.0029$	$[-0.0113,+0.0171]$
QQP	F1	0.8463	0.8561	0.8626	$+0.0065$	$[-0.0044,+0.0173]$
RTE	Acc	0.7155	0.6440	0.7350	$+0.0910$	$[+0.0161,+0.1658]$
SST-2	Acc	0.9372	0.9241	0.9385	$+0.0144$	$[+0.0037,+0.0252]$
Macro mean	–	0.7962	0.7807	0.8104	$+0.0297$	$[+0.0014,+0.0580]$

Appendix

Table 20: Historical-projection control on RoBERTa-base, evaluated once after five epochs on each of the five evaluation seeds under the final-test convention of Appendix B . The control keeps BOM ’s task-space history but transports it through each step’s own model and examples rather than the current ones. On a common trajectory, its numerator differs from BOM ’s by

\lambda

times the bias-corrected historical-projection drift term of Eq. ( 7 ). Both arms share the coefficients, temporal weighting, covariance treatment, AdamW-style adaptive-scaling rule and weight decay; each updates its scale statistic from its own numerator. Its independently selected per-task learning rates are listed in Table 36 .

\Delta

is the mean paired BOM -minus-control score difference. Intervals are descriptive, unadjusted 95% paired-

t

intervals over five matched evaluation runs (

df=4

), using the sample standard deviation of the paired differences. Macro differences first average the five tasks equally within each matched run before estimating uncertainty.

	$\cos(\hat{M}_{t},h_{t})$		$\lVert\hat{M}_{t}-h_{t}\rVert/\lVert\hat{M}_{t}\rVert$		Cancellation
Task	head	encoder	head	encoder	ratio
CoLA	$0.84$	$0.46$	$0.45$	$0.86$	$0.27$
MRPC	$0.63$	$0.26$	$0.71$	$0.96$	$0.28$
QQP	$0.81$	$0.58$	$0.47$	$0.75$	$0.25$
RTE	$0.98$	$0.84$	$0.20$	$0.50$	$0.14$
SST-2	$0.79$	$0.44$	$0.49$	$0.88$	$0.33$

Appendix

Table 21: Historical-projection drift measured inside the own-step control on RoBERTa-base, over the five evaluation seeds and all telemetry steps from step 100 to the epoch-5 endpoint.

\hat{M}_{t}

accumulates the task-space signals through each step’s own model and examples;

h_{t}=\bar{J}_{t}^{\top}\hat{q}_{t}

reprojects the same history through the current model and examples. head is the classification head, encoder every other trainable tensor. Cancellation is

\lVert s_{t}\rVert/B^{-1}\sum_{i}\lVert r_{t,i}\rVert

. Tasks are weighted equally in the mean.

Task	Batch-only	Param. $\Delta=20$	Param. $\Delta=100$	Param. $\Delta=300$	Mixed $\Delta=20$	Mixed $\Delta=300$
RTE	$0.686$	$0.740$	$0.588$	$0.349$	$0.577$	$0.294$
MRPC	$0.655$	$0.702$	$0.579$	$0.493$	$0.545$	$0.432$
CoLA	$0.619$	$0.830$	$0.704$	$0.622$	$0.559$	$0.480$
SST-2	$0.548$	$0.850$	$0.736$	$0.656$	$0.507$	$0.432$
QQP	$0.475$	$0.848$	$0.750$	$0.672$	$0.456$	$0.402$

Appendix

Table 22: Decomposing historical-projection drift on BOM trajectories across the five RoBERTa-base tasks (the five evaluation seeds, probes every 20 steps from step 100 to the epoch-5 endpoint). Fixed probe batches

A

and

B

and a fixed unit task-space direction

s^{\ast}

define

v_{X}(\theta)=\bar{J}_{X}(\theta)^{\top}s^{\ast}

. Batch-only compares

v_{A}

and

v_{B}

at the same parameters; parameter-only compares

v_{A}

at parameters

\Delta

steps apart; mixed changes both. Entries are encoder cosines averaged over seeds and probe steps; lower values indicate larger departure. The EMA at

\beta_{1}=0.9

has a mean lag of nine steps.

Task	LR	A	B	C	D	E	Mean score	Mean loss	Status
CoLA	$2.5\times 10^{-5}$	0.6407	0.4624	0.6547	0.5417	0.6007	0.5800	0.4775	Primary
MRPC	$2.0\times 10^{-5}$	0.9242	0.9078	0.9158	0.8538	0.9154	0.9034	0.3237	Primary
MRPC	$3.0\times 10^{-5}$	0.7722	0.9154	0.9007	0.9061	0.9144	0.8818	0.4517	Rescue
QQP	$3.0\times 10^{-5}$	0.8633	0.8561	0.8631	0.8627	0.8292	0.8549	0.2624	Primary
QQP	$2.0\times 10^{-5}$	0.8537	0.8440	0.8588	0.8396	0.8556	0.8503	0.2694	Rescue
RTE	$2.0\times 10^{-5}$	0.5632	0.7365	0.7401	0.7978	0.5740	0.6823	0.5960	Primary

Appendix

Table 23: Classifier-only past-history control on RoBERTa-base. Each row retains all five evaluation runs (A–E) at epoch 5. Score is Matthews correlation for CoLA, F1 for MRPC/QQP and accuracy for RTE/SST-2; loss is cross-entropy. Primary marks the protocol-selected configurations used in the primary macro and Rescue the configurations selected by the three-HPO-seed rescue searches. Selection rules are in Appendix B.6 .

			Cell B: BOM numerator CE second moment		Cell C
Dataset	AdamW	BOM	at BOM ’s rate	at $3\times 10^{-5}$	AdamW numerator BOM $v$
CoLA	0.5743	0.5970	0.5501	0.5565	0.1146
MRPC	0.9077	0.9191	0.9134	0.9134	0.8551
QQP	0.8463	0.8626	0.8427	0.8427	0.5175
RTE	0.7155	0.7350	0.6578	0.6968	0.5747
SST-2	0.9372	0.9385	0.9365	0.9365	0.6716

Appendix

Table 24: Numerator/second-moment cross-pairings on RoBERTa-base. Entries are epoch-5 means over all five evaluation seeds, including degenerate runs. Cell B combines BOM ’s numerator with the CE-gradient second moment; cell C combines AdamW’s numerator with the BOM -numerator second moment. Search rules and final rates are in Appendix B.6 and Table 43 .

Task	FullV Train	OutRMS Train	FullV Val	OutRMS Val	$\Delta$ Val	OutRMS parameter- shaped state (MiB)
RTE	95.46	49.44	75.45 Acc	49.28 Acc	-26.17 pt	0
MRPC	98.28	50.00	92.15 F1	40.61 F1	-51.54 pt	0

Appendix

Table 25: Strict output-only second-moment diagnostic on RoBERTa-base at the fixed epoch-5 endpoint. Means use two paired seeds under the primary fine-tuning schedule. BOM-FullV retains parameter-space

v_{\ell,t-1}

; BOM-OutRMS replaces it with a task-space second-moment EMA and backpropagates the resulting RMS-normalized task signal. The output-only variant removes all parameter-shaped optimizer state while retaining task-space EMA state, but stays near a one-class/chance solution.

Dataset	Metric	Val score	Val loss
CoLA	MCC	0.5508	0.5313
MRPC	F1	0.9140	0.3737
QQP	F1	0.8547	0.2607
RTE	Acc	0.7105	0.6167
SST-2	Acc	0.9369	0.1982
Macro mean	–	0.7934	0.3961

Appendix

Table 26: RoBERTa-base anchored APM variant by dataset. The corresponding BOM aggregate is reported in Table 13 . Entries are epoch-5 endpoints over the five evaluation seeds. Score denotes Matthews correlation for CoLA, F1 for MRPC and QQP, and accuracy for RTE and SST-2; loss is validation cross-entropy.

Validation score
Dataset	Metric	AdamW	BOM	$\Delta$
CoLA	MCC	$0.5789\pm 0.0337$	$\mathbf{0.6150}\pm 0.0176$	$+0.0361$
MRPC	F1	$0.9048\pm 0.0058$	$\mathbf{0.9186}\pm 0.0075$	$+0.0138$
QQP	F1	$\mathbf{0.8809}\pm 0.0074$	$0.8777\pm 0.0070$	$-0.0032$
RTE	Acc	$0.7350\pm 0.0161$	$\mathbf{0.7625}\pm 0.0065$	$+0.0274$
SST-2	Acc	$0.9278\pm 0.0047$	$\mathbf{0.9372}\pm 0.0073$	$+0.0094$

Appendix

Table 27: FP32-state replication of the RoBERTa-base AdamW– BOM comparison. Both arms follow the shared fine-tuning protocol and are evaluated once after epoch 5 on the official validation split used as the final test (Appendix B ) over the five evaluation seeds. Scores use the task metric shown; losses are validation cross-entropy. Entries are mean

\pm

sample standard deviation over five seeds, and

\Delta

is the paired BOM

-

AdamW mean. The macro score margin has paired 95% CI

[+0.0073,+0.0261]

(

p=0.0078

). Task-level paired intervals in Table 28 use the same endpoint. All 50 endpoints are nondegenerate. BOM has the higher mean score on four of five tasks and a five-task macro of

0.8222

versus AdamW’s

0.8055

; its validation cross-entropy is lower on four tasks and on the macro, while AdamW is lower on QQP.

Dataset	$\Delta$ score	Paired 95% CI	Two-sided $p$
CoLA	$+0.0361$	$[-0.0176,+0.0898]$	0.1354
MRPC	$+0.0138$	$[+0.0011,+0.0264]$	0.0392
QQP	$-0.0032$	$[-0.0197,+0.0134]$	0.6250
RTE	$+0.0274$	$[+0.0010,+0.0539]$	0.0450
SST-2	$+0.0094$	$[+0.00032,+0.01849]$	0.0453
Five-task macro	$+0.0167$	$[+0.0073,+0.0261]$	0.0078

Appendix

Table 28: FP32-state replication: paired validation-score margins from the one-shot final test after epoch 5, matching Table 27 and the protocol in Appendix B . Intervals are descriptive, unadjusted 95% paired-

t

intervals over five matched evaluation seeds (

df=4

). Macro uncertainty is computed after averaging the five tasks within each seed. The macro margin excludes zero; QQP’s interval includes zero.

Backbone	Projection	A	B	C	D	E	Mean
ConvNeXt-Tiny	current	0.9784	0.9688	0.9789	0.9771	0.9814	0.9769
ConvNeXt-Tiny	own-step	0.9541	0.9773	0.9774	0.9808	0.9555	0.9690
ViT-Tiny	current	0.9675	0.9703	0.9626	0.9634	0.9709	0.9669
ViT-Tiny	own-step	0.9559	0.9655	0.9628	0.9611	0.9578	0.9606

Appendix

Table 29: Independently tuned multiclass own-step projection control on STL10. Entries are accuracy at the 40-epoch endpoint; A–E identify the five matched evaluation runs. Current-reprojection rows reuse Table 11 . Both arms use the same selection budget; training and HPO settings are in Appendix B.7 .

	Optimizer state	Peak allocated (MiB)			Step time
Method	(MiB)	Mean	Paired ratio	Paired $\Delta$	Mean (ms)	Paired ratio
AdamW (FP32)	951.0	3194.0	–	–	49.35	–
BOM -AdamW (FP32)	475.5	2865.1	0.905	$-328.9$	47.15	0.957
AdamW	475.5	2276.9	–	–	39.84	–
BOM -AdamW	237.7	2045.2	0.894	$-231.7$	39.50	0.993
Adam-mini	238.2	1724.5	–	–	35.08	–

Appendix

Table 30: RoBERTa-base memory and step time for three base optimizers and their BOM compositions over five tasks and five evaluation seeds. FP32 rows keep FP32 master weights and optimizer state with bf16 compute; the other rows use bf16 storage. State counts parameter-shaped optimizer buffers in MiB, excluding the task-space vector; peak denotes allocated GPU memory. Ratios and differences are paired within task and seed. Measurement settings and aggregation are specified in Appendix B.8 .

Model / tokens	$B\!\times\!T$	$n$	AdamW	BOM	$\Delta$ (%)	AdamW	BOM	time ratio
			Peak allocated (GiB)			E2E time (h)		Mean paired
55M / 3.0B	$128\times 256$	5	27.116	26.901	$-0.215$ ( $-0.8\%$ )	8.175	8.714	1.066
110M / 5.0B	$24\times 256$	5	7.557	7.151	$-0.406$ ( $-5.4\%$ )	19.355	20.063	1.037
440M / 8.5B	$160\times 256$	3	87.239	83.118	$-4.121$ ( $-4.7\%$ )	31.959	32.975	1.032
1.1B / 15.0B	$80\times 256$	3	83.503	78.159	$-5.344$ ( $-6.4\%$ )	178.193	177.561	0.996

Appendix

Table 31: C4 resource measurements used in Figure 9 . The study uses five consecutive-seed reporting runs at 55M/110M and three at 440M/1.1B. Times are arithmetic means of the selected end-to-end intervals, including data I/O, evaluation, checkpoint writes and other process-wall overhead. Peak denotes allocated GPU memory, measured as defined in Appendix B.8 . Timing intervals and aggregation are defined in Appendix B.8 .

Panel of Figure 9 — Figure 9: C4 resource comparison at four scales. (a) End-to-end throughput, $\sum_{i}N_{i}/\sum_{i}T_{i}$ , includes in-window data I/O, evaluation, checkpoint writes and other process-wall overhead. (b) Peak allocated-memory difference. (c) Checkpoint sizes; 55M used no intermediate checkpoint. (d) Mean paired duration ratio with sample-standard-deviation bars over $n_{t}=5,5,3,3$ pairs; the dotted line marks parity. The 440M pairs use the common interval, steps 96,436–207,519. At 1.1B, one pair uses steps 388,184–732,421 and two pairs use complete runs; each pair has identical token endpoints. See Table 31 and Appendix B.8 for timing intervals and aggregation.

Panel of Figure 10 — Figure 10: Paired BOM /AdamW step-time ratio against parameter count on the 250M–4B ladder (single H800, 32k vocabulary, matched FP32 protocol, median of 200 timed steps). The ratio falls monotonically, crosses parity between the 1B and 2B rungs, and reaches $0.938$ at 4B. Per-device tokens per step halve from the 1B rung onward (Table 32 ), so $r$ grows faster than $P$ along the ladder. The observed crossing describes these implementations and workloads.

Rung	$P$	$V$	mb	Tokens /step	$r$	AdamW (ms)	BOM (ms)	$\Delta$ (ms)	Ratio
250M	$2.60\times 10^{8}$	32,768	128	32,768	0.24	504.1	544.9	$+40.9$	1.081
500M	$4.94\times 10^{8}$	32,768	128	32,768	0.46	1041.8	1081.5	$+39.7$	1.038
1B	$9.84\times 10^{8}$	32,768	64	16,384	1.83	889.0	897.9	$+9.0$	1.010
2B	$2.00\times 10^{9}$	32,768	32	8,192	7.4	864.0	845.9	$-18.1$	0.979
4B	$3.95\times 10^{9}$	32,768	16	4,096	29.5	914.5	858.2	$-56.3$	0.938

Appendix

Table 32: H800 (Hopper, 80 GB) single-GPU ladder. Median step time over 200 timed steps;

\Delta

is BOM

-

AdamW. Micro-batch (mb) is selected by the stated out-of-memory halving procedure on the BOM arm; both arms share it.

Study	Selection policy	Reporting and qualifications
NLP primary pairs; RoBERTa optimizer controls	Unified LR envelope $[10^{-6},3\times 10^{-4}]$ ; 15 configurations, two-epoch screening on the internal 10% holdout	Fresh five-epoch training on full official train; fixed epoch-5 official-validation endpoint; bf16 compute and stated storage. Muon additionally searches five update scales after LR selection.
FP32-state replication	Repeat the same train-only LR policy independently for each arm and task	FP32 parameters/state, bf16 compute; same five evaluation seeds and one final evaluation on the official validation split.
NLP own-step; full temporal kernel	Unified fine-tuning LR policy	Per-task rates in Table 36 ; constructions in Appendix A.6 .
Classifier-only past history	Unified 15-configuration, two-HPO-seed policy; primary score uses the protocol-selected configurations	Rescue searches on MRPC (one), QQP (one) and RTE (three) are all reported, and none replaces the primary configurations (Appendix B.6 ).
Objective–momentum controls	Unified fine-tuning LR policy; same 15-configuration, two-HPO-seed budget	Independently selected taskwise rates; five evaluation seeds at the epoch-5 endpoint (Table 43 ).
APM; coefficient sensitivity	Task-specific diagnostic configurations	Table 36 ; coefficient sweeps are post-selection diagnostics.

Appendix

Table 33: Protocol overview. LR denotes peak learning rate. HPO uses deterministic 90/10 partitions of the applicable training pool and does not access the designated validation split. After HPO selection, or after the peak rate is specified for the no-HPO 1.1B study, final runs train afresh on the complete training pool. Supervised fine-tuning reports the fixed official-validation endpoint, with logged RoBERTa intermediate evaluations shown descriptively; pretraining uses its designated validation split and reports it as validation. Each candidate uses two HPO seeds unless stated otherwise; reported quality uses five consecutive evaluation seeds disjoint from the HPO seeds, except the three-run C4-440M and C4-1.1B groups. Timing is excluded from learning-rate selection.

Coefficient	AdamW	BOM
First-order decay $\beta_{1}$ (parameter EMA / task-space EMA)	0.9	0.9
Second-moment decay $\beta_{2}$	0.999	0.999
Denominator constant $\epsilon$	$10^{-8}$	$10^{-8}$
Decoupled weight decay $\omega$	0.01	0.01
Current/history mixture $(1-\lambda)/\lambda$	–	$0.5/0.5$

Appendix

Table 34: Fixed optimizer coefficients for the primary supervised AdamW– BOM comparisons. The BOM objective uses a fixed

0.5/0.5

current/history mixture, and its task-space EMA decay is set to the same value as AdamW’s first-moment decay

\beta_{1}

, so the two stored histories share the EMA weights

w_{t,\tau}

. After EMA bias correction becomes negligible, BOM ’s batch-mean projected residual has current/past total weights

0.55/0.45

, with the current covariance term weighted

0.5

(Eq. 5 ); AdamW’s first moment weights current/past full gradients

0.10/0.90

Item	Value
NLP backbones	roberta-base ; microsoft/deberta-v3-base ; Qwen/Qwen3-1.7B
Classification head	each backbone’s AutoModelForSequenceClassification head, num_labels $=2$
Tokenization	each backbone’s fast AutoTokenizer ; the decoder-only backbone pads with its end-of-sequence (EOS) token
NLP data	GLUE ( nyu-mll/glue via Hugging Face datasets ); CoLA, MRPC, QQP, RTE, SST-2; HPO uses a deterministic 10% holdout from official train, and the official validation split supplies the fixed epoch-5 reporting scores; logged RoBERTa intermediate evaluations are descriptive
Software	PyTorch $\geq$ 2.2, Transformers $\geq$ 4.53, datasets $\geq$ 2.18, torchvision $\geq$ 0.17, timm 1.0.15 (pinned), adam-mini 1.1.1 and galore-torch 1.0 (pinned)
Control implementations	Adam-mini and GaLore use the authors’ released packages, adam-mini 1.1.1 and galore-torch 1.0, using the parameter grouping and projection scope described in Appendix B.3.1 ; all other controls use the sources cited with them

Appendix

Table 35: Environment for the reported runs. Model identifiers are Hugging Face model ids; every backbone is loaded through its standard sequence-classification interface, so no custom head or pooling is introduced. Precision, batch sizes, schedules, and the seed protocol are specified in the surrounding text of this appendix.

Method	CoLA	MRPC	QQP	RTE	SST-2
AdamW	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2.5\times 10^{-5}$	$3\times 10^{-5}$
BOM -AdamW	$2.5\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$
APM	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$
Lion	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$	$1\times 10^{-5}$
Adam-mini	$7\times 10^{-6}$	$1\times 10^{-5}$	$3\times 10^{-6}$	$3\times 10^{-6}$	$1\times 10^{-5}$
BOM -Adam-mini	$5\times 10^{-6}$	$5\times 10^{-6}$	$1\times 10^{-5}$	$3\times 10^{-6}$	$5\times 10^{-6}$

Appendix

Table 36: Task-specific learning rates for the RoBERTa-base optimizer comparisons and the APM diagnostic, by dataset. Values are peak learning rates. All cells use max length 128, 10% linear warmup, and subsequent linear decay to zero.

Dataset	AdamW LR	BOM LR	Epochs	Max length
CoLA	$3\times 10^{-5}$	$3\times 10^{-5}$	5	128
MRPC	$3\times 10^{-5}$	$3\times 10^{-5}$	5	128
QQP	$3\times 10^{-5}$	$3\times 10^{-5}$	5	128
RTE	$3\times 10^{-5}$	$3\times 10^{-5}$	5	128
SST-2	$3\times 10^{-5}$	$3\times 10^{-5}$	5	128

Appendix

Table 37: DeBERTa-v3-base reproducibility parameters by dataset, for the methods reported on this backbone (AdamW and BOM ). Each arm independently selects its learning rate using the common 15-configuration, two-HPO-seed protocol; the selected rates coincide. Values are peak learning rates; all cells use 10% linear warmup followed by linear decay to zero and train for five epochs.

Dataset	AdamW LR	BOM LR	Epochs	Max length
CoLA	$2\times 10^{-5}$	$2\times 10^{-5}$	5	256
MRPC	$2\times 10^{-5}$	$2\times 10^{-5}$	5	256
QQP	$2\times 10^{-5}$	$2\times 10^{-5}$	5	256
RTE	$1\times 10^{-5}$	$1\times 10^{-5}$	5	256
SST-2	$2\times 10^{-5}$	$2\times 10^{-5}$	5	256

Appendix

Table 38: Qwen3-1.7B reproducibility parameters by dataset, for the methods reported on this backbone (AdamW and BOM ). Each arm independently selects its learning rate using the common 15-configuration, two-HPO-seed protocol; the selected rates coincide. Values are peak learning rates; all cells use 10% linear warmup followed by linear decay to zero and train for five epochs.

Setting	Value
GaLore, BOM -GaLore
Implementation	galore-torch 1.0 (authors’ release)
Projected modules	the authors’ GLUE target list: attention and feed-forward Linear weights
Projection rank	8
Projection scale	1.0
Projection refresh interval	200 steps

Appendix

Table 39: Settings specific to individual optimizers, beyond the learning rates in Table 36 . Every value is fixed before formal HPO and is not varied during it, except Muon’s update scale, fixed at

0.2

during learning-rate screening and then given an additional five-point search after the common learning-rate search, with the selected learning rate held fixed (Appendix B ). GaLore’s projected modules and rank follow the authors’ GLUE fine-tuning script, and its projection scale and refresh interval are the galore-torch 1.0 defaults; all are fixed for every GaLore and BOM -GaLore cell. Parameters outside the projected module list—embeddings, LayerNorms, biases and the classification head—remain unprojected. The GaLore baseline applies the package’s AdamW update to these parameters. In BOM -GaLore, both branches omit the local first moment and retain their second-moment rules and decoupled weight decay; the second moments are updated from the corresponding BOM numerator. SCALE applies its AdamW branch to one-dimensional parameters at the selected matrix rate, so only that rate is searched. SCALE uses the authors’ official implementation. Its backbone matrices carry no first moment; the

0.9

decay is the output-layer momentum. The backbone-matrix update normalizes its gradient before applying the learning rate.

Backbone	Dataset	LR	Epochs	Image size
ConvNeXt-Tiny	STL10	$3\times 10^{-5}$	40	224
ViT-Tiny	STL10	$3\times 10^{-5}$	40	224

Appendix

Table 40: STL10 transfer parameters for AdamW and BOM . Both backbones fine-tune from ImageNet-pretrained weights with batch size 32, bf16 compute,

\beta_{1}=0.9

\beta_{2}=0.999

\epsilon=10^{-8}

, weight decay

0.01

, and

\lambda=0.5

for BOM . Both arms use linear warmup over the first

10\%

of updates followed by linear decay to zero. Independently selected learning rates coincide. Both arms report the final checkpoint under the shared five-seed convention. Own-step settings are in Appendix B.7 .

Micro-batch	Minimum throughput (tokens/s)
24 (selected)	64,698
32	64,598
128	59,121

Appendix

Table 41: Selected C4-110M batch-qualification measurements. Throughput is the minimum over the four pre-run configurations described above, rounded to the nearest token/s. These measurements select a common micro-batch; they do not compare validation quality.

Setting	Model	Budget	Schedule (warmup)	AdamW LR	BOM LR
ImageNet-1k	ResNet-50 (random init)	90 epochs	linear (5 epochs)	$1\times 10^{-3}$	$1.5\times 10^{-3}$
Python code	Qwen3-55M, 32k vocab	3.0B tokens	linear (10%)	$5\times 10^{-3}$	$1.5\times 10^{-3}$
FineWeb-Edu	Qwen3-55M, 32k vocab	3.0B tokens	linear (10%)	$4\times 10^{-3}$	$1\times 10^{-3}$
C4	Qwen3-55M, 32k vocab	3.0B tokens	linear (10%)	$4\times 10^{-3}$	$2\times 10^{-3}$
C4	Qwen3-110M, 32k vocab	5.0B tokens	linear (10%)	$3\times 10^{-3}$	$1.5\times 10^{-3}$
C4	Qwen3-440M, 32k vocab	8.5B tokens	linear (10%)	$4\times 10^{-3}$	$4\times 10^{-3}$

Appendix

Table 42: From-scratch pretraining parameters. Reporting uses five consecutive-seed runs except for the C4 440M and 1.1B studies, which each use three. Budgets refer to final schedules. Each ImageNet HPO evaluation screens its learning rate for 15 epochs over two HPO seeds. Language HPO uses scale-specific screening endpoints: 500M tokens for the 55M/3.0B settings, 1.0B for C4-110M/5.0B, and 1.5B for C4-440M/8.5B. The 55M and 110M screens average two HPO seeds. C4-440M retains the same coarse grid, admissible range, 10-evaluation allocation, loss criterion, and selection rule but uses one HPO seed; both separate screens select

4\times 10^{-3}

. C4-1.1B performs no HPO and fixes

10^{-3}

for both arms a priori. Both arms use

\beta_{1}=0.9

and

\beta_{2}=0.999

: AdamW applies

\beta_{1}

to its parameter-space first moment and BOM applies it to its task-space output EMA. The BOM arm has no parameter-space first-moment tensor and uses

\lambda=0.5

Diagnostic	CoLA	MRPC	QQP	RTE	SST-2
Mixed + local $m$	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$
CE + no local $m$	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$
Cell B, at BOM ’s rate	$2.5\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$
Cell B, at $3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$
Cell C, at AdamW’s rate	$3\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2.5\times 10^{-5}$	$3\times 10^{-5}$
Fixed-batch drift probe	$2.5\times 10^{-5}$	$3\times 10^{-5}$	$3\times 10^{-5}$	$2\times 10^{-5}$	$3\times 10^{-5}$

Appendix

Table 43: Learning rates used for the RoBERTa-base mechanism diagnostics. The objective–momentum rows report independently selected learning rates under the same 15-configuration, two-HPO-seed HPO budget as AdamW and BOM (Table 15 ); the cell B/C rows correspond to Table 24 ; the fixed-batch drift probe corresponds to Table 22 . Cell B’s first row uses the primary BOM rates; its second row uses the rates selected by targeted CoLA/RTE searches and retains the anchor rates on the other tasks. Cell C uses the primary AdamW rates, with a separate CoLA stability screen described in Appendix B.6 . The two cell B rows reuse the same MRPC, QQP and SST-2 runs, where their rates coincide.

Rung	$P$	$L$	$d$	$d_{\mathrm{ff}}$	$H_{q}/H_{kv}$	$d_{h}$	Vocab.	Pos.	GC
250M	260,089,344	18	1024	3072	8/4	128	32,768	4096	No
500M	494,207,488	23	1280	3840	10/5	128	32,768	4096	Yes
1B	983,658,240	24	1792	5376	14/7	128	32,768	4096	Yes
2B	1,995,527,680	31	2304	6912	18/6	128	32,768	4096	Yes
4B	3,954,407,168	35	3072	9216	24/8	128	32,768	4096	Yes

Appendix

Table 44: Architectures of the single-GPU systems benchmark in Table 32 , taken from its model builder;

P

is the exact parameter count, including query/key normalization weights. All use Qwen3ForCausalLM , tied input/output embeddings and disabled KV caching; the timed context length is 256.

L

d

, and

d_{\mathrm{ff}}

denote decoder layers, hidden width and intermediate width;

H_{q}/H_{kv}

are query/key–value head counts, and

d_{h}

is head width. Pos. is the configured maximum position count, and GC denotes activation checkpointing. Remaining fields in constructed configurations use the benchmark’s Qwen3Config defaults.

Setting	Output coordinates	$P$	$N_{\mathrm{out}}$	$r$
NLP fine-tuning, RoBERTa-base ( $B{=}32$ , $C{=}2$ )	$B\cdot C$	$1.25\times 10^{8}$	$64$	${\sim}2\times 10^{6}$
STL10 fine-tuning, ViT-Tiny / ConvNeXt-Tiny ( $B{=}32$ , $C{=}10$ )	$B\cdot C$	$5.7\times 10^{6}$ – $2.8\times 10^{7}$	$3.2\times 10^{2}$	${\sim}10^{4}$ – $10^{5}$
ImageNet-1k pretraining, ResNet-50 ( $B{=}256$ , $C{=}1000$ )	$B\cdot C$	$2.56\times 10^{7}$	$2.6\times 10^{5}$	${\sim}10^{2}$
LM pretraining, Qwen3-55M ( $B\,T{=}32{,}768$ , $V{=}32{,}768$ )	$B\cdot T\cdot V$	$5.45\times 10^{7}$	$1.1\times 10^{9}$	${\sim}0.05$

Appendix

Table 45: The dimensionless element-count ratio

r=P/N_{\mathrm{out}}

of Eq. ( 19 ) across the training configurations of the main suite. These are structural element counts. Timing predictions additionally require the implementation-specific coefficients in Eq. ( 19 ).

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 206 papers

Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule
Aug 19, 2026Euijin Hong, Guannan QuDeep Learning Optimization Momentum Methods
Electrical and Computer Engineering, Carnegie Mellon University
A Step Towards Forgetting: Optimiser History and the Loss of Answer Mass
Oct 2, 2026Valeria Ruscio, Seth Nabarro, Keiran ThompsonLLM Fine-Tuning Catastrophic Forgetting
Intuition Machines Inc
SOLO: Pretraining Billion-Parameter Language Models with Shared-Output Local Learning
Sep 28, 2026Bojian Yin, Shurong Wang, Yuqi Pan +1Language Model Pretraining Efficient Language Model Training
Institute of Automation, Chinese Academy of Sciences

Aug 19, 2026J/K move · Enter open · S save

Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule

Euijin Hong, Guannan Qu

Electrical and Computer Engineering, Carnegie Mellon University

Deep Learning Optimization Momentum Methods Neural Network Optimization Accelerated Gradient Methods

Open paper AskPDF arXiv

Most modern optimizers form their momentum as an exponential moving average (EMA) of past gradients, forgetting every direction at one fixed rate. However, the inputs a deep network sees during training can be highly anisotropic, with a few directions queried frequently while most are seen rarely. Preconditioning methods address this anisotropy by wrapping extra processing around this buffer and leave the momentum update itself unchanged. We propose Activation-Keyed Momentum (AK-Momentum), which builds direction-awareness into the momentum update rule. The gradient of a linear layer splits into an input activation that acts as a key and an output-side error that acts as a value. Keying on that activation, AK-Momentum updates the momentum buffer by the canonical delta rule, so each direction is forgotten at a rate set by how often it appears. We prove that it is a valid momentum, that it applies the input-side curvature correction without matrix inversion, and that it clears stale directions faster than EMA under both a fixed and a drifting optimum. It is a drop-in replacement for the momentum buffer of any optimizer, its coefficient transfers across widths under

μ

P, and its extra compute stays between

22.2\%

and

25.0\%

of a gated-MLP block's linear cost with no persistent memory. In FineWeb-Edu pretraining, AdamW with AK-Momentum (AK-AdamW) reaches AdamW's validation loss in up to

46.39 \pm 4.32\%

fewer steps at 67M and

22.12 \pm 0.80\%

at 370M over three seeds, and the gain persists at 1B on a Chinchilla-optimal budget. A Muon baseline tuned under the same protocol sits above AK-AdamW at both language-model scales, and the gain holds for SGD, ResNet-18, and ViT-Tiny on CIFAR-10. Training-time diagnostics confirm the predicted mechanism, better gradient tracking and healthier input directions.

Abstract

Deep Learning Optimization

Explore similar work

Activation-Keyed Momentum: An Anisotropic Momentum Update via the Delta Rule

A Step Towards Forgetting: Optimiser History and the Loss of Answer Mass

SOLO: Pretraining Billion-Parameter Language Models with Shared-Output Local Learning