mlsanityData to 14 Oct 2026

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

arXiv 2605.21292 ↗First posted 20 May 2026stat.ML cs.AI cs.LGcited by 1

Not on the map. Why?

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.

Key facts. Cited by 1 paper 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

26 of 31 references.

References · 26

  1. 2022Transformers learn in-context by gradient descentcited by 732
  2. 2021Gradient Descent on Neural Networks Typically Occurs at the Edge of Stabilitycited by 427
  3. 2023Trained Transformers Learn Linear Models In-Contextcited by 337
  4. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 279
  5. 2020The large learning rate phase of deep learning: the catapult mechanismcited by 216
  6. 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 183
  7. 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
  8. 2022Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stabilitycited by 155
  9. 2022Understanding Gradient Descent on Edge of Stability in Deep Learningcited by 128
  10. 2023How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?cited by 96
  11. 2022SGD with Large Step Sizes Learns Sparse Featurescited by 68
  12. 2022Understanding Edge-of-Stability Training Dynamics with a Minimalist Examplecited by 56
14 more
  1. 2021Large Learning Rate Tames Homogeneity: Convergence and Balancing Effectcited by 50
  2. 2021Strength of Minibatch Noise in SGDcited by 39
  3. 2024In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initializationcited by 33
  4. 2020Stochasticity of Deterministic Gradient Descent: Large Learning Rate for Multiscale Objective Functioncited by 26
  5. 2023Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theorycited by 26
  6. 2024High dimensional analysis reveals conservative sharpening and a stochastic edge of stabilitycited by 18
  7. 2020Multiplicative noise and heavy tails in stochastic optimizationcited by 17
  8. 2023From Stability to Chaos: Analyzing Gradient Descent Dynamics in Quadratic Regressioncited by 17
  9. 2021Neural Network Weights Do Not Converge to Stationary Points: An Invariant Measure Perspectivecited by 15
  10. 2026Rod Flow: A Continuous-Time Model for Gradient Descent at the Edge of Stabilitycited by 9
  11. 2025Gradient Descent with Large Step Sizes: Chaos and Fractal Convergence Regioncited by 7
  12. 2026Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioningcited by 5
  13. 2025Softmax as Linear Attention in the Large-Prompt Regime: a Measure-based Perspectivecited by 3
  14. 2026State-Dependent Lyapunov Analysis of Rank-1 Matrix Factorizationcited by 2

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 1.

1 paper.