Mpemba Effect in Large-Language Model Training Dynamics: A Minimal Analysis of the Valley-River model
This provides a principled justification for tuning learning rates in LLMs, potentially reducing hyperparameter sweeps, but is incremental as it builds on existing thermodynamic analogies and focuses on specific loss landscapes.
The paper tackles the lack of mechanistic explanation for learning rate schedules in large language model training by analyzing a 'valley-river' loss landscape using the Mpemba effect, showing that a high plateau learning rate can accelerate convergence during decay by eliminating the slowest mode.
Learning rate (LR) schedules in large language model (LLM) training often follow empirical templates: warm-up, constant plateau/stable phase, and decay (WSD). However, the mechanistic explanation for this strategy remains underexplored, and the choice of plateau height and decay schedule is largely heuristic. In this paper, we connect training dynamics to a thermodynamic analogy via the Mpemba effect - a phenomenon in which a hotter system cools faster than a colder one when quenched into the same bath. We analyze a class of "valley-river" loss landscapes, where sharp (valley) directions equilibrate quickly, while flatter (river) directions govern global descent. The Mpemba effect provides an explanation for the necessity of the warm-up phase and motivates a high plateau - rather than a low one - for accelerating loss decrease during decay. We show that for certain loss landscapes, there exists an optimal plateau learning rate - the "strong Mpemba point" - at which the slowest mode vanishes, resulting in faster convergence during the decay phase. We derive analytical conditions for its existence and estimate decay dynamics required to preserve the Mpemba advantage. Our minimal model and analysis offer a principled justification for plateau-based schedulers and provide guidance for tuning LR in LLMs with minimal hyperparameter sweep.