Research map / Architectures and efficiency / Attention and transformers
Learning and Transferring Sparse Contextual Bigrams with Linear Transformers
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.
Proposes the Sparse Contextual Bigram (SCB) task and analyzes training a one-layer linear transformer with a preconditioned, l1-regularized projected proximal gradient method, showing a two-stage sample complexity and that pretraining with nontrivial correlation lets finetuning skip the sample-intensive stage.
Problem
Theoretical understanding of how transformers learn language modeling tasks that require both in-context (contextual) retrieval and global transition knowledge, including training dynamics, sample complexity, and transfer learning.
Approach
Extends the bigram model so the next token depends on a sparse, last-token-dependent set of earlier positions, shows it is representable by a one-layer linear transformer, and uses preconditioned l1-proximal gradient descent with thresholding stages to prove recovery with polynomial sample complexity, with a faster finetuning path when a pretrained model is correlated with the downstream task.
Key results
- Preconditioned l1-proximal GD recovers the ground-truth transition and attention matrices, while vanilla SGD either diverges with large step or converges slowly with small step at batch size 64.
- Training from scratch splits into a poly(T)-sample signal-boosting stage followed by poly(N,Q,1/eps,log T)-sample convergence.
- Finetuning from a pretrained model with nontrivial correlation skips the poly(T) stage, needing only poly(N,Q,1/eps,log T) samples.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ℓ1-regularized proximal gradient descent after Stage 1 instead of vanilla SGD | One-layer linear transformer · Synthetic SCB (N=3, Q=2, T=5000) | distance to ground truth | Helped |
| Use vanilla SGD with the larger learning rate after Stage 1 (ablation) | One-layer linear transformer · Synthetic SCB (N=3, Q=2, T=5000) | convergence | Hurt |
| Use vanilla SGD with the smaller learning rate after Stage 1 (ablation) | One-layer linear transformer · Synthetic SCB (N=3, Q=2, T=5000) | convergence rate | Hurt |
| Replace linear attention with softmax attention | Transformer trained on SCB · SCB (T=50, N=10, Q=2) | loss | No clear effect |
Surprising findings
- Vanilla SGD fails to converge at the same small batch size and large learning rate due to gradient variance scaling with T.
Limitations
- Restricted to one-layer single-head linear transformers with merged/fixed weight blocks, far from practical transformers.
- Non-standard multi-stage training algorithm with manual thresholding steps.
- Data model generates a single next token with independent context tokens, not natural sequential language.
Experimental setup
Datasets: synthetic Sparse Contextual Bigram data
Benchmarks: synthetic SCB task
Metrics: Frobenius/mu-norm distance to ground truth (||V-P||_mu, ||A-
Baselines: vanilla SGD
Builds on: bigram model linear transformer induction head analysis
Closest papers by meaning
- Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.88)
- Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers (NeurIPS 2025 · similarity 0.87)
- Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.87)
- Pruning Large Language Models with Semi-Structural Adaptive Sparse Training (AAAI 2025 · similarity 0.87)
- In-Context Learning with Representations: Contextual Generalization of Trained Transformers (NeurIPS 2024 · similarity 0.87)
- Adaptive Layer Sparsity for Large Language Models via Activation Correlation Assessment (NeurIPS 2024 · 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.