# Training trajectories, mini-batch losses and the curious role of the learning rate

arXiv: https://arxiv.org/abs/2301.02312 · HTML: https://mlsanity.com/paper/2301.02312/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 33 references that are in the map.
- First posted on arXiv 5 Jan 2023.

## Builds on (most cited first, 20 of 20)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Linear Mode Connectivity and the Lottery Ticket Hypothesis](https://mlsanity.com/paper/1912.05671/) (2019)
- [Cyclical Learning Rates for Training Neural Networks](https://mlsanity.com/paper/1506.01186/) (2015)
- [Git Re-Basin: Merging Models modulo Permutation Symmetries](https://mlsanity.com/paper/2209.04836/) (2022)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Stochastic Gradient Descent as Approximate Bayesian Inference](https://mlsanity.com/paper/1704.04289/) (2017)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Stochastic modified equations and adaptive stochastic gradient algorithms](https://mlsanity.com/paper/1511.06251/) (2015)
- [Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks](https://mlsanity.com/paper/1710.11029/) (2017)
- [Asymmetric Valleys: Beyond Sharp and Flat Local Minima](https://mlsanity.com/paper/1902.00744/) (2019)
- [Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging](https://mlsanity.com/paper/2209.14981/) (2022)
- [Beyond the Quadratic Approximation: the Multiscale Structure of Neural Network Loss Landscapes](https://mlsanity.com/paper/2204.11326/) (2022)
- [Strength of Minibatch Noise in SGD](https://mlsanity.com/paper/2102.05375/) (2021)
- [Random initialisations performing above chance and how to find them](https://mlsanity.com/paper/2209.07509/) (2022)
- [Amortized Proximal Optimization](https://mlsanity.com/paper/2203.00089/) (2022)

## Built on by (most cited first, 21 of 21)

- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [The Road Less Scheduled](https://mlsanity.com/paper/2405.15682/) (2024)
- [Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits](https://mlsanity.com/paper/2411.18704/) (2024)
- [No Train No Gain: Revisiting Efficient Training Algorithms For Transformer-based Language Models](https://mlsanity.com/paper/2307.06440/) (2023)
- [Early Weight Averaging meets High Learning Rates for LLM Pre-training](https://mlsanity.com/paper/2306.03241/) (2023)
- [Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs](https://mlsanity.com/paper/2502.15938/) (2025)
- [The Surprising Agreement Between Convex Optimization Theory and Learning-Rate Scheduling for Large Model Training](https://mlsanity.com/paper/2501.18965/) (2025)
- [WSM: Decay-Free Learning Rate Schedule via Checkpoint Merging for LLM Pre-training](https://mlsanity.com/paper/2507.17634/) (2025)
- [Training Dynamics Impact Post-Training Quantization Robustness](https://mlsanity.com/paper/2510.06213/) (2025)
- [Butterfly Effects of SGD Noise: Error Amplification in Behavior Cloning and Autoregression](https://mlsanity.com/paper/2310.11428/) (2023)
- [ZipNN: Lossless Compression for AI Models](https://mlsanity.com/paper/2411.05239/) (2024)
- [Through the River: Understanding the Benefit of Schedule-Free Methods for Language Model Training](https://mlsanity.com/paper/2507.09846/) (2025)
- [When, Where and Why to Average Weights?](https://mlsanity.com/paper/2502.06761/) (2025)
- [ScheduleFree+: Scaling Learning-Rate-Free & Schedule-Free Learning to Large Language Models](https://mlsanity.com/paper/2605.19095/) (2026)
- [Lossless and Near-Lossless Compression for Foundation Models](https://mlsanity.com/paper/2404.15198/) (2024)
- [Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs](https://mlsanity.com/paper/2512.17131/) (2025)
- [EMA Without the Lag: Bias-Corrected Iterate Averaging Schemes](https://mlsanity.com/paper/2508.00180/) (2025)
- [Understanding SGD with Exponential Moving Average: A Case Study in Linear Regression](https://mlsanity.com/paper/2502.14123/) (2025)
- [PADAM: Parallel averaged Adam reduces the error for stochastic optimization in scientific machine learning](https://mlsanity.com/paper/2505.22085/) (2025)
- [QAM: Quadratic-Accurate Checkpoint Merging via Sequential Consistency](https://mlsanity.com/paper/2609.35168/) (2026)
- [Understanding the Weight Averaging Mechanism in LLM Training for Post-Training Quantization](https://mlsanity.com/paper/2610.05329/) (2026)

## Abstract (arXiv)

Stochastic gradient descent plays a fundamental role in nearly all applications of deep learning. However its ability to converge to a global minimum remains shrouded in mystery. In this paper we propose to study the behavior of the loss function on fixed mini-batches along SGD trajectories. We show that the loss function on a fixed batch appears to be remarkably convex-like. In particular for ResNet the loss for any fixed mini-batch can be accurately modeled by a quadratic function and a very low loss value can be reached in just one step of gradient descent with sufficiently large learning rate. We propose a simple model that allows to analyze the relationship between the gradients of stochastic mini-batches and the full batch. Our analysis allows us to discover the equivalency between iterate aggregates and specific learning rate schedules. In particular, for Exponential Moving Average (EMA) and Stochastic Weight Averaging we show that our proposed model matches the observed training trajectories on ImageNet. Our theoretical model predicts that an even simpler averaging technique, averaging just two points a many steps apart, significantly improves accuracy compared to the baseline. We validated our findings on ImageNet and other datasets using ResNet architecture.