Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks
Abstract
We study the dynamics of stochastic gradient descent (SGD) for a class of sequence models termed Sequence Single-Index (SSI) models, where the target depends on a single direction in input space applied to a sequence of tokens. This setting generalizes classical single-index models to the sequential domain, encompassing simplified one-layer attention architectures. We derive a closed-form expression for the population loss in terms of a pair of sufficient statistics capturing semantic and positional alignment, and characterize the induced high-dimensional SGD dynamics for these coordinates. Our analysis reveals two distinct training phases: escape from uninformative initialization and alignment with the target subspace, and demonstrates how the sequence length and positional encoding influence convergence speed and learning trajectories. These results provide a rigorous and interpretable foundation for understanding how sequential structure in data can be beneficial for learning with attention-based models.
Cite
Text
Arnaboldi et al. "Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks." Advances in Neural Information Processing Systems, 2025.Markdown
[Arnaboldi et al. "Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks." Advances in Neural Information Processing Systems, 2025.](https://mlanthology.org/neurips/2025/arnaboldi2025neurips-asymptotics/)BibTeX
@inproceedings{arnaboldi2025neurips-asymptotics,
title = {{Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks}},
author = {Arnaboldi, Luca and Loureiro, Bruno and Stephan, Ludovic and Krzakala, Florent and Zdeborova, Lenka},
booktitle = {Advances in Neural Information Processing Systems},
year = {2025},
url = {https://mlanthology.org/neurips/2025/arnaboldi2025neurips-asymptotics/}
}