Research map / Architectures and efficiency / Attention and transformers
Transformer Learns Optimal Variable Selection in Group-Sparse Classification
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.
Shows theoretically that a one-layer transformer trained by joint gradient descent from zero initialization learns to attend to the label-relevant variable group in group-sparse classification, with a tight loss bound, and that its pretrained weights transfer to downstream tasks with improved sample complexity.
Problem
Understanding how transformers learn structured data: specifically, whether gradient-trained one-layer attention can perform optimal variable selection on group-sparse binary classification data, and whether the pretrained model transfers with few samples.
Approach
Analyzes the full optimization trajectory of all parameters (value vector and combined query-key matrix) of a simplified one-layer softmax-attention transformer with positional encodings under joint gradient descent from zero initialization, yielding matching upper and lower loss bounds and a sample complexity improvement over vectorized logistic regression for transfer.
Key results
- Theorem gives matching loss bounds C_1(eps ∧ 1/D^2) ≤ L ≤ eps ∧ 1/D^2 after T* = Θ(D^3 ∨ 1/(D^3 eps^3)) GD iterations.
- Downstream online SGD yields generalization error Õ((d+D)/n), versus Ω(dD/eps) sample complexity for vectorized logistic regression.
- Empirically, test accuracy converges to 1 on downstream synthetic tasks and attention concentrates on the label-relevant group row.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a one-layer self-attention transformer by gradient descent on group-sparse classification data | One-layer self-attention · Synthetic group-sparse classification | training loss | Helped |
| Fine-tune a pre-trained one-layer transformer on downstream tasks sharing the group sparsity pattern | Pre-trained one-layer self-attention · Downstream synthetic classification | test accuracy | Helped |
| Use self-attention to select the label-relevant group | One-layer self-attention · Synthetic group-sparse classification | attention score on the label-relevant group | Helped |
Limitations
- Analysis restricted to a one-layer self-attention model with reparameterized query-key and value vector.
- Uses population loss, assuming infinite training data.
- Requires number of groups D to scale with log^2(1/epsilon) and Gaussian/sub-Gaussian data assumptions.
Experimental setup
Datasets: synthetic group-sparse Gaussian data
Benchmarks: synthetic group-sparse data
Metrics: population cross-entropy loss attention score sparsity cosine similarity test accuracy generalization error bound
Baselines: linear logistic regression on vectorized features
Builds on: one-layer transformer optimization analysis (Tarzanagh et al.) sparse linear regression with transformers
Closest papers by meaning
- Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.89)
- One-Layer Transformer Provably Learns One-Nearest Neighbor In Context (NeurIPS 2024 · similarity 0.88)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.86)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.86)
- Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context (ICLR 2025 · similarity 0.86)
- On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures (ICML 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.