Transformer Learns Optimal Variable Selection in Group-Sparse Classification
transformer;self-attention;variable selection;group-sparse classification
Yes
No ethics review needed.
the input x_j is concatenated with position encoding, does Theorem 3.2 still hold if using addition instead of concatenation?
The paper is well written and easy to follow. 

The theoretical analysis is thorough and the results align with the group sparsity assumption, although I did not check all the math and proof in detail.

In a sense, it provides new insights into sparsity analysis on the workings of attention based models.
This paper investigates one-layer transformers trained on specific dataset, where the input variables are generated from multiple groups, while the true label of this input is determined by variables from a single group. Based on these simplifications and assumptions, it theoretically demonstrates that the one-layer transformers can almost attend to the variables from the label-relevant group. Moreover, it provides a tight lower and upper bound for the population cross-entropy loss of a one-layer transformer trained by gradient descent. It also shows that the well pre-trained one-layer transformers can be efficiently transferred to a downstream task sharing a similar “group-sparse” structure and further provides an improved generalization error bound for one-layer transformers fine-tuned by SGD, which surpasses that of linear logistic regression applied to vectorized features. The numerical experiment observations support the theoretical findings.
(1) too much data assumption and model simplification. According to Definition 3.1, each patch x_j is i.i.d from Gaussian, and its label is determined from a given v (which can be learned from samples later). These data assumptions make group attention trivial.

(2) And the model is one-layer transformers, which may cannot capture attention-based models to a sufficiently satisfactory extent. With these simplifications, the contribution of this paper is limited, especially considering the pre-work from Jelassi et al. (2022).
Yes
No ethics review needed.
See weakness
The author presented a novel theoretical analysis on how the transformer can learn to select variables for the group sparse classification, which is interesting and novel.
This paper analyzes the variable selection of the Transformer theoretically.
1. Given your previous definition of the notations, pls define $\Theta$ in advance （otherwise people would get confused about this common notation.What is the 'variable' in the paper? Are they some features or some attributes of the data that define the label? What is the $v_1$ in the Thereom 3.2-2？ The weight corresponding to the label 1 or the value vector (output from value matrix in Transformer)？\\
2. The proof stretch should be placed in the Supplementary instead of the main paper. \\
3. The experiments are conducted on synthetic data, however, in the higher dimensional images (e.g. 3*224*224), things would change. I have concerns on these too simple experiments. \\
4. The Lemma 5.1, I do not see why the $W^{(T^*)}_{1,2}$ holds and $W$ is diagonal. In NLP, it seems that the feature is correlated to its position and would correspond to the output, is that a too strong assumption?
5. Can you show a more detailed demonstration on why the inequality of $\alpha^(T^*)$ holds (or is just a definition?) and why Lemma5.2 can be incorporated into Lemma 5.3? Also, I do not think Lemma 5.2. holds as the underlying assumption is the Transformer is naturally with low error. Lemma 5.2 is more like a strong assumption and definition to me. \\
6. How you define the group sparse? What is this different from the standard classification?


I would like to raise my scores if more details are provided.
Yes
No ethics review needed.
Please refer to the weaknesses section.
1. This paper investigates the theoretical understanding of transformers, an important and interesting problem that may greatly benefit the deep learning community.
2. This paper is well-written and has a good motivation.
3. The conclusion that a one-layer transformer can use attention mechanisms to select relevant variables and ignore irrelevant ones for classification makes sense to me.
This paper investigates the theoretical understanding of transformers, which have shown success in various applications but lack a clear theoretical foundation. The study focuses on how transformers can learn a statistical model with group sparsity, where only a subset of input variables (groups) influences the label. The research shows that a one-layer transformer can use attention mechanisms to select relevant variables and ignore irrelevant ones for classification. Additionally, it demonstrates that a well-pretrained one-layer transformer can adapt to new tasks with limited samples and achieve good accuracy. This study provides insights into how transformers learn structured data effectively.
1. Though the problem is interesting, the conclusion, a one-layer transformer can use attention mechanisms to select relevant variables and ignore irrelevant ones for classification, seems to be well-known in the deep learning community and has already been verified by the original paper (Attention Is All You Need).
2. The theoretical understanding of transformers given by this paper mainly focuses on the well-pretrained one-layer transformer. However, it will be much more interesting to study the multi-layer transformer since the scaling law shows that the loss of transformer scales as a power-law with model size, dataset size, and the amount of compute used for training, with some trends spanning more than seven orders of magnitude. 
3. The study focuses on how transformers can learn a statistical model with group sparsity, where only a subset of input variables (groups) influences the label. However, the concept of group sparse is not properly defined, making me confused to fully understand it. Further, what is a similar “group-sparse” structure as mentioned in the Introduction section? What kind of “group-sparse” inputs are similar?
@inproceedings{
anonymous2024transformer,
title={Transformer Learns Optimal Variable Selection in Group-Sparse Classification},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=fuoM5YDBX4},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
transformer self-attention variable selection group-sparse classification
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/384c7992f72314ae73c49cfbefbae3127bd43f78.pdf
learning theory
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Transformer Learns Optimal Variable Selection in Group-Sparse Classification
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission