Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks

2025/06/03 by Luca Arnaboldi, Arnaboldi, Luca, Bruno Loureiro +7 · 3 citations
Computer Science · #Class (philosophy) #Convergence (economics) #Disordered Systems and Neural Networks (cond-mat.dis-nn) #Encoding (memory) #FOS: Computer and information sciences #FOS: Physical sciences #Generative Adversarial Networks and Image Synthesis #Initialization #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Neural Networks and Applications #Population #Sequence (biology) #Sequence learning #Space (punctuation) #Stochastic Gradient Optimization Techniques

paper · pdf · doi:10.48550/arxiv.2506.02651

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2025/06/03 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/05

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.

Citations

• When Do Transformers Outperform Feedforward and Recurrent Networks? A Statistical Perspective
• Dynamical Decoupling of Generalization and Overfitting in Large Two-Layer Networks
• Understanding Generalization in Transformers: Error Bounds and Training Dynamics Under Benign and Harmful Overfitting
• Optimal Spectral Transitions in High-Dimensional Multi-Index Models
• Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds
• Unraveling the Gradient Descent Dynamics of Transformers
• On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent
• Attention layers provably solve single-location regression
• High-dimensional learning of narrow neural networks
• Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot
• Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions
• Computational-Statistical Gaps in Gaussian Single-Index Models
• A phase transition between positional and semantic learning in a solvable model of dot-product attention
• On the Convergence of Encoder-only Shallow Transformers
• On Learning Gaussian Multi-index Models with Gradient Flow
• A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks
• Symmetric Single Index Learning
• Gradient-Based Feature Learning under Structured Data
• Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models
• Representational Strengths and Limitations of Transformers
• How Two-Layer Neural Networks Learn, One (Giant) Step at a Time
• Escaping mediocrity: how two-layer networks learn hard generalized linear models with SGD
• Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models
• Leveraging the two timescale regime to demonstrate convergence of neural networks
• SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics
• From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks
• Neural Networks can Learn Representations with Gradient Descent
• The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks
• Language Models are Few-Shot Learners
• Mean Field Analysis of Neural Networks: A Central Limit Theorem
• On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport
• Attention Is All You Need
• Orthogonal Decomposition of Symmetric Tensors
• A Stochastic Approximation Method
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

Cited by

• Inductive Bias and Spectral Properties of Single-Head Attention in High Dimensions
• Statistical Advantage of Softmax Attention: Insights from Single-Location Regression
• A statistical physics framework for optimal learning

Related

• Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds
• SGD on Neural Networks Learns Functions of Increasing Complexity
• Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models
• Escaping mediocrity: how two-layer networks learn hard generalized linear models with SGD
• SGD with Large Step Sizes Learns Sparse Features