# Transformer Learns Optimal Variable Selection in Group-Sparse Classification

arXiv: https://arxiv.org/abs/2504.08638 · HTML: https://mlsanity.com/paper/2504.08638/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 11 Apr 2025.

## Built on by (most cited first, 2 of 2)

- [Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression](https://mlsanity.com/paper/2508.07571/) (2025)
- [Towards Understanding Transformers in Learning Random Walks](https://mlsanity.com/paper/2511.23239/) (2025)

## Abstract (arXiv)

Transformers have demonstrated remarkable success across various applications. However, the success of transformers have not been understood in theory. In this work, we give a case study of how transformers can be trained to learn a classic statistical model with "group sparsity", where the input variables form multiple groups, and the label only depends on the variables from one of the groups. We theoretically demonstrate that, a one-layer transformer trained by gradient descent can correctly leverage the attention mechanism to select variables, disregarding irrelevant ones and focusing on those beneficial for classification. We also demonstrate that a well-pretrained one-layer transformer can be adapted to new downstream tasks to achieve good prediction accuracy with a limited number of samples. Our study sheds light on how transformers effectively learn structured data.