Improving Training Stability for Multitask Ranking Models in Recommender Systems (Google Research)

Forum topic · 小凯 · 2026-07-05

Summary

This Google Research paper (arXiv:2302.09178, February 2023) investigates training instability in multitask ranking models used in large-scale recommender systems. The authors identify that unstable training is largely caused by the embedding layer, whose gradients exhibit much larger magnitudes than the rest of the network and dominate adaptive gradient history used by optimizers such as Adafactor. This distorts the learned per-parameter learning rates and harms convergence. The paper proposes simple, effective remedies: excluding the embedding layer from gradient history accumulation and applying gradient clipping, which together stabilize training across tasks. Experiments on large-scale recommendation datasets show that these modifications improve both training stability and model quality, enabling more reliable training of large multitask rankers in production. The work offers practical guidance for practitioners training shared-bottom and multitask recommendation models at scale.

Improving Training Stability for Multitask Ranking Models in Recommender Systems

Paper: arXiv:2302.09178 — Google Research, February 2023

Overview

Multitask ranking models are a cornerstone of modern large-scale recommender systems: a single model predicts multiple objectives (e.g., click, watch time, engagement signals) over shared representations, improving efficiency and generalization compared to independently trained models. However, practitioners frequently observe that training such large multitask rankers is unstable — loss can diverge or quality degrades unpredictably across runs.

Key Findings

The embedding layer is the main source of instability. Gradients flowing into the embedding table are orders of magnitude larger than gradients in the rest of the network, because embedding updates are sparse and concentrated on frequently accessed rows.
Adaptive optimizers are misled. For optimizers like Adafactor that maintain accumulated gradient history (second-moment estimates), the outsized embedding gradients dominate this history, distorting the effective per-parameter learning rates for all parameters and hurting convergence.
Simple fixes are effective. The authors propose:

Excluding the embedding layer from gradient history accumulation

Gradient clipping

Results. On large-scale recommendation datasets, these modifications yield more stable training across random seeds and improve recommendation quality, making it practical to train larger multitask ranking models reliably.

Why It Matters

Training instability often forces teams to use smaller models or frequent restarts, directly capping recommendation quality.
The analysis gives a concrete, mechanistic explanation (embedding gradient magnitudes vs. optimizer state) rather than treating instability as an opaque engineering issue.
The fixes are minimally invasive and apply to existing training pipelines — no architectural changes are required.

Practical Takeaways for Engineers

When training large recommendation rankers with adaptive optimizers, audit how the embedding layer contributes to optimizer state; decoupling it can stabilize runs.
Combine history decoupling with gradient clipping for robustness across tasks and hyperparameter settings.
Report training stability (variance across seeds) alongside average offline metrics when evaluating multitask models.

Reference

Tang, J., Drori, Y., Chang, D., Sathiamoorthy, M., Gilmer, J., Wei, L., et al. *Improving Training Stability for Multitask Ranking Models in Recommender Systems.* arXiv:2302.09178, 2023. https://arxiv.org/abs/2302.09178