Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention
Abstract
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.
Related
- Unbiased kinetic Langevin Monte Carlo with inexact gradients
- Unbiased Low-Rank Approximation for Repeated Matrix-Vector Multiplication: Stability and Optimal Sampling
- Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training
- Mesh-Uniform Power Stability of Two-Relaxation-Time Vector Lattice Boltzmann Schemes with Reversible Boundaries
- Wasserstein Stability and Free Boundaries in Measure-Parameterized Bilevel Obstacle Problems
BibTeX
@misc{balasubramanian2026unbiased,
title = {{Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention}},
author = {Krishnakumar Balasubramanian},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2605.21292},
eprint = {2605.21292},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.21292},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention
AU - Krishnakumar Balasubramanian
PY - 2026
DA - 2026/09/21
PB - arXiv
DO - 10.48550/arxiv.2605.21292
UR - https://arxiv.org/abs/2605.21292
AB - 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.
ER -