Beyond One Epoch: Uncertainty-Weighted Sensitivity Regularization for Recommendation Models

Richard Lettich; Shagun Gupta

cs.LGSep 28, 2026

Beyond One Epoch: Uncertainty-Weighted Sensitivity Regularization for Recommendation Models

Authors: Richard Lettich, Shagun Gupta

Organizations: Meta Platforms, Inc. Menlo Park, CA, USA

Abstract

Recommendation models with sparse embeddings and a shared consumer often exhibit the one-epoch phenomenon: a second epoch lowers training loss while sharply degrading generalization. We present a view based on the violation of the prequential principle. On the first epoch, an example's label has not affected the embedding rows used to score it. On later epochs, those rows contain a displacement induced by the labels earlier update. This creates an incentive for the shared consumer to exploit this displacement in subsequent epochs, which fails to generalize. We call this self-influence asymmetry. Using an exact scalar model and local influence analysis, we connect this mismatch to the uncertainty in the embeddings and the consumers incentive to exploit it in subsequent epochs. We verify this hypothesis using an embedding-consumer-update interventions in deep recommendation models and propose uncertainty-weighted sensitivity regularization (UWSR) which counteracts this mismatch by augmenting the loss function to penalize the consumer for relying on uncertain embeddings. Unlike existing remedies, UWSR preserves the learned embeddings and across three benchmarks, four-epoch UWSR reduces test cross-entropy by 1.38%-6.78% and improves AUC by 0.0058-0.0231 relative to one-epoch training.

Figures & tables

	Embedding $\rightarrow$ Consumer			Consumer $\rightarrow$ Embedding
Dataset	$\Delta\mathrm{BCE}_{\mathrm{train}}$	$\Delta\mathrm{BCE}_{\mathrm{val}}$	$\Delta S$	$\Delta\mathrm{BCE}_{\mathrm{train}}$	$\Delta\mathrm{BCE}_{\mathrm{val}}$	$\Delta S$
MovieLens-20M	$+1.43\%$	$-2.90\%$	$+54.0\%$	$+0.19\%$	$-0.21\%$	$-0.1\%$
Amazon Electronics	$+10.93\%$	$-13.14\%$	$+474.3\%$	$+0.22\%$	$-0.44\%$	$-4.8\%$
Avazu	$+1.73\%$	$-1.86\%$	$+56.4\%$	$+0.15\%$	$+0.31\%$	$+4.6\%$

Table 1: One-epoch embedding-consumer-update-order intervention, averaged over two seeds.

\Delta\mathrm{BCE}_{\mathrm{train}}

is the reduction in same-minibatch training loss after the first block update.

\Delta\mathrm{BCE}_{\mathrm{val}}

is the validation-loss reduction relative to ordinary training; positive values indicate improvement.

\Delta S

is the relative change in primary-row logit sensitivity relative to joint training.

Dataset	Method	E1	E2	E3	E4	E2	E3	E4
		Test AUC				$\Delta$ BCE (%)
Amazon	Naive Replay	0.8523	0.8594	0.8403	0.8233	$+0.02$	$-28.51$	$-95.43$
	MEDA	0.8523	0.8608	0.8644	0.8649	$+2.54$	$+3.62$	$+3.92$
	Freeze until final epoch	0.8523	0.8469	0.8483	0.8460	$-1.61$	$-1.10$	$-1.74$
	Active-row AdamW	0.8523	0.8594	0.7852	0.7811	$+0.02$	$-20.06$	$-23.54$
	Embedding dropout	0.8523	0.8575	0.8453	0.8417	$+0.72$	$-4.03$	$-7.16$

Table 2: Test-set performance after each epoch, averaged over five seeds. Each endpoint uses the coefficient selected on validation at that endpoint. The right block gives the paired relative BCE reduction against the same seed’s one-epoch checkpoint; positive values are better.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Method	MovieLens	Amazon	Avazu
Naive replay	$-24.66$	$-101.03$	$-15.71$
MEDA	$-0.93$	$+3.37$	$-0.06$
Freeze embeddings	$-1.89$	$-1.45$	$+0.06$
AdamAR	$-0.65$	$+2.19$	$+1.96$
UWSR Uniform sensitivity (unweighted ablation)	$+0.40$	$-1.05$	$+2.50$
UWSR, rowwise	$+1.47$	$+6.77$	$+3.97$

Appendix

Table 3: Validation

\Delta

BCE (%) at epoch four for algorithms compared in Table 2 . Positive values denote improvement over the corresponding one-pass model.

(a) Relative $\Delta$ BCE (%) with 95% CI
Method	MovieLens	Amazon	Avazu
Naive replay	$-33.66$ $[-45.12,-22.20]$	$-95.43$ $[-102.76,-88.10]$	$-13.85$ $[-15.57,-12.14]$
MEDA	$+0.21$ $[-2.37,+2.79]$	$+3.92$ $[+3.23,+4.62]$	$-0.11$ $[-0.58,+0.37]$
Freeze until final pass	$-0.90$ $[-3.60,+1.81]$	$-1.74$ $[-3.07,-0.42]$	$-0.20$ $[-0.81,+0.40]$
Active-row AdamW	$+0.10$ $[-2.45,+2.65]$	$-23.54$ $[-25.58,-21.51]$	$-5.76$ $[-7.19,-4.34]$
Embedding dropout	$-2.63$ $[-4.36,-0.90]$	$-7.16$ $[-8.82,-5.50]$	$-7.77$ $[-10.72,-4.82]$

Appendix

Table 4: Paired test-set changes at epoch four, averaged over five seeds. Brackets give two-sided 95%

t

-intervals across paired seed-level differences from the corresponding one-pass Adam checkpoint; positive values are improvements.

Accumulator	Scalars per row	$\lambda$	Validation BCE $\downarrow$	AUC $\uparrow$
Row-wise scalar	$1$	$0.03$	$\mathbf{0.442510}$	$\mathbf{0.876256}$
Per-coordinate diagonal	$32$	$0.03$	$0.442851$	$0.875963$

Appendix

Table 5: Accumulator-granularity ablation on Amazon Electronics after four epochs (seed 7). Both variants selected

\lambda=0.03

on validation. Auxiliary state excludes model parameters and optimizer state.

Architecture	BCE	Naive	MEDA	Freeze	AdamAR	Isotropic	UWSR	$\Delta$ AUC
	One pass	$\Delta$ BCE						Fisher
WuKong	0.474634	$-101.03$	$+3.37$	$-1.45$	$+2.19$	$-1.05$	$\mathbf{+6.95}$	$+0.0222$
DCNv2	0.507843	$-45.23$	$+6.71$	$+3.08$	$+5.90$	$+2.26$	$\mathbf{+9.02}$	$+0.0362$
AutoInt	0.491371	$-57.86$	$+5.29$	$-3.35$	$+5.56$	$-0.97$	$\mathbf{+8.58}$	$+0.0303$

Appendix

Table 6: Amazon validation results across WuKong, DCNv2 ( Wang et al., 2021 ) , and AutoInt ( Song et al., 2019 ) . Entries are

\Delta

BCE relative to each architecture’s one-pass baseline.

Method	Epochs	Later-epoch step time	End-to-end overhead
Naive replay / MEDA / AdamAR	4	$1.00\times$	$1.00\times$
UWSR	4	$1.18\times$	$1.13\times$

Appendix

Table 7: Amazon computational overhead relative to naive replay four-pass training, mean 2 seeds. Later-epoch step time is relative to one naive replay step; end-to-end overhead includes all epochs and method-specific setup.

Explore similar work

Memory Layer: Train the In-Model Cache for Recommendation Models
Jul 27, 2026Liangyuan Na, Gufan Yin, Yixin Bao +19Text Embeddings Cache
Meta
RecRec: Recursive Refinement for Sequential Recommendation
Jul 12, 2026Pervez Shaik, Prosenjit Biswas, Abhinav Thorat +2Multimodal Recommendation Model Sequential Recommendation
Sony Research India India
Policy Regret for Embedding Model Routing: Contextual Bandits with Low-Rank Experts
Jun 12, 2026Yan Dai, Negin Golrezaei, Patrick JailletBehavioral Embeddings Low-Rank Structure
Operations Research Center, MIT. · Sloan School of Management, MIT. · Department of EECS, MIT.

Jul 27, 2026, cs.IRJ/K move · Enter open · S save

Memory Layer: Train the In-Model Cache for Recommendation Models

Liangyuan Na, Gufan Yin, Yixin Bao, Xianjie Chen+18

Meta

Text Embeddings Cache Keys And Value

Open paper AskPDF arXiv

Early ranking stages in recommendation systems precompute item embeddings and cache them in-model for scoring within strict latency constraints. Because this cache exists only at serving time, outside the training loop, training and serving use different item representations, a structural discrepancy that limits quality and adds operational fragility. We show that co-designing the training and serving paths removes this representation discrepancy at its source. We introduce the memory layer, an in-model key-value embedding cache co-trained with the model: the item tower writes embeddings during training and the model reads them at serving, one source of truth for item representations by construction. Always-on embeddings cover items not yet cached, so every item receives a prediction, and the design consolidates three separate trainer-to-predictor update paths into a single self-contained pipeline. Deployed in production on Instagram Reels, the memory layer raises prediction coverage from 96% to 100%, improves embedding freshness from

O(5\text{ min})

O(20\text{ s})

, and narrows the training-serving Normalized Entropy (NE) gap by up to 86%, yielding over

2\times

recall for the freshest content and a 5-6% cold start engagement lift. Because embeddings are produced during training, the system needs no separate bulk-evaluation or publish-time recomputation, cutting training-and-publish computational cost by 30% at neutral serving computational cost.

Abstract

Explore similar work

Memory Layer: Train the In-Model Cache for Recommendation Models

RecRec: Recursive Refinement for Sequential Recommendation

Policy Regret for Embedding Model Routing: Contextual Bandits with Low-Rank Experts