# Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 31 references that are in the map.
- First posted on arXiv 20 May 2026.

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

- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Small-scale proxies for large-scale Transformer training instabilities](https://mlsanity.com/paper/2309.14322/) (2023)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2209.15594/) (2022)
- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [SGD with Large Step Sizes Learns Sparse Features](https://mlsanity.com/paper/2210.05337/) (2022)
- [Understanding Edge-of-Stability Training Dynamics with a Minimalist Example](https://mlsanity.com/paper/2210.03294/) (2022)
- [Large Learning Rate Tames Homogeneity: Convergence and Balancing Effect](https://mlsanity.com/paper/2110.03677/) (2021)
- [Strength of Minibatch Noise in SGD](https://mlsanity.com/paper/2102.05375/) (2021)
- [In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization](https://mlsanity.com/paper/2402.14951/) (2024)
- [Stochasticity of Deterministic Gradient Descent: Large Learning Rate for Multiscale Objective Function](https://mlsanity.com/paper/2002.06189/) (2020)
- [Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory](https://mlsanity.com/paper/2307.04204/) (2023)
- [High dimensional analysis reveals conservative sharpening and a stochastic edge of stability](https://mlsanity.com/paper/2404.19261/) (2024)
- [Multiplicative noise and heavy tails in stochastic optimization](https://mlsanity.com/paper/2006.06293/) (2020)
- [From Stability to Chaos: Analyzing Gradient Descent Dynamics in Quadratic Regression](https://mlsanity.com/paper/2310.01687/) (2023)
- [Neural Network Weights Do Not Converge to Stationary Points: An Invariant Measure Perspective](https://mlsanity.com/paper/2110.06256/) (2021)
- [Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stability](https://mlsanity.com/paper/2602.01480/) (2026)
- [Gradient Descent with Large Step Sizes: Chaos and Fractal Convergence Region](https://mlsanity.com/paper/2509.25351/) (2025)
- [Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning](https://mlsanity.com/paper/2603.01514/) (2026)
- [Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspective](https://mlsanity.com/paper/2512.11784/) (2025)
- [State-Dependent Lyapunov Analysis of Rank-1 Matrix Factorization](https://mlsanity.com/paper/2604.26993/) (2026)

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

- [Learning Decision-Stump Thresholds in Context: Dynamics of Softmax Attention](https://mlsanity.com/paper/2610.07074/) (2026)

## Abstract (arXiv)

Unbiased stochastic gradients can match the full gradient in expectation while changing finite-step stability. We study this effect in one-layer linear self-attention for in-context linear regression, where shared-mode mini-batch training reduces exactly to a random two-factor map. The full-batch map preserves an elliptic region, whereas sampled curvature and target correlation create batch-dependent stability boundaries. Their combined fluctuation is centered at the parameter-update level but induces a strictly outward drift of the full-batch boundary coordinate. We derive an exact one-step crossing criterion, adaptive concentration bounds, and conditional drift identities, and prove almost-sure escape in a balanced switching model. For independent long prompts, mode leakage decays at the inverse square root of prompt length. For finite-token softmax attention, we derive an explicit boundary-crossing law governed by an effective context length that shrinks as attention scores concentrate. Experiments in ambient attention objectives verify the reduced predictions, and four trained causal softmax Transformers show larger mean curvature shifts for smaller first-step batches, consistent with the predicted nonlinear outward bias.