Research map / Optimization / Implicit bias and generalization
Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes high-dimensional one-pass spherical SGD for a tied single-layer attention model learning Gaussian sequence single-index targets, deriving population-loss sufficient statistics, a sequence information exponent, and a sequence-length speed-up.
Problem
Understanding the sample complexity and dynamics of SGD for attention-based (single-layer tied attention) models learning sequence single-index targets in high dimensions, including the effects of sequence length and positional encoding.
Approach
Introduces the sequence information exponent (SIE), a generalization of the information exponent to sequence single-index models, and reduces the population risk of tied attention to two sufficient statistics (semantic overlap m and positional overlap e), giving an ODE description of SGD and a gain over untied networks that scales with sequence length.
Key results
- SIE determines sample complexity: O(d) for SIE=1, O(d log^2 d) for SIE=2, O(d^(SIE-1)) for SIE>=3 (informal).
- Tied network speed-up over untied network scales as L^2 for the SIE=2 target sum_i He_2(z_i) (numerical fit with d=1000, ReLU).
- Semantic-to-positional transition in the phase diagram occurs at omega_trans ≈ 0.64 for a=1, matching simulations with d=1000 (sharpening with d).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add positional encoding to the attention model (ablation) | Tied single-layer attention SGD · SSI target with SIE = 4 | sequence information exponent: 4 -> 2 | Helped |
| Use tied rather than untied weights | Sequence learning with SIE = 1 · SSI, SIE = 1 | weak-recovery gain: gain proportional to L | Helped |
| Use tied rather than untied weights | Sequence learning with g(z*) = sum_i He_2(z*_{i}) · SSI, SIE = 2 | weak-recovery gain: gain asymp L^2 | Helped |
| Add positional encoding (ablation) | Tied single-layer attention SGD · SSI models | sequence information exponent | Helped |
| Add positional encoding (ablation) | Tied single-layer attention SGD · SSI models where the loss landscape changes but sample complexity does not | sample complexity | No clear effect |
| Use tied weights rather than untied weights | Sequence learning with general target function g · SSI models | weak-recovery speed | Mixed |
| Vary omega and a, changing the target between semantic and positional behavior (ablation) | Attention with positional encoding · Semantic/positional target | convergence to global minimum: omega_trans = 0.64 | Mixed |
Surprising findings
- Positional encoding can lower the SIE (e.g., from 4 to 2), but it is not always beneficial; some cases change the landscape without changing sample complexity.
- SGD can converge to a local positional minimum rather than the global semantic minimum, so it does not always reach the global optimum.
Limitations
- Analysis restricted to Gaussian i.i.d. tokens and tied single-head attention with d_head=1 and identity value matrix
- Results rely on gradient-flow (small learning rate) approximation and high-dimensional limit d large
- Target class restricted to even SIE in main text; odd cases require injected positional encoding
Experimental setup
Datasets: synthetic Gaussian sequence single-index data
Metrics: overlap with target (weak recovery time) gain in gradient steps (tied vs untied) empirical probability of convergence to semantic minima
Baselines: untied network fully connected network
Builds on: single-index model SGD analysis (information exponent) sequence multi-index models Hermite expansion
Closest papers by meaning
- Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds (ICML 2025 · similarity 0.89)
- Bayes optimal learning of attention-indexed models (NeurIPS 2025 · similarity 0.87)
- Understanding the Differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks (NeurIPS 2024 · similarity 0.87)
- A Phase Transition between Positional and Semantic Learning in a Solvable Model of Dot-Product Attention (NeurIPS 2024 · similarity 0.87)
- High-Dimensional Learning Dynamics of Attention-Indexed Models (NeurIPS 2026 · similarity 0.86)
- Attention layers provably solve single-location regression (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.