Transformer Learns Optimal Variable Selection in Group-Sparse Classification

Authors: Chenyang Zhang, Xuran Meng, Yuan Cao
Primary area: learning theory
Keywords: transformer, self-attention, variable selection, group-sparse classification

Abstract:
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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730368020405

thread_type:
official_review

<thread_object>
creator: reviewer_teud

<content>
summary:
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.

strengths:
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.

weaknesses:
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?

questions:
Please refer to the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596768419

thread_type:
official_review

<thread_object>
creator: reviewer_bj8k

<content>
summary:
This paper analyzes the variable selection of the Transformer theoretically.

strengths:
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.

weaknesses:
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.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701293962

thread_type:
official_review

<thread_object>
creator: reviewer_czbc

<content>
summary:
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.

strengths:
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.

weaknesses:
(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).

questions:
the input x_j is concatenated with position encoding, does Theorem 3.2 still hold if using addition instead of concatenation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732443181818

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Overview

comment:
Dear Reviewers,

We appreciate your detailed and constructive comments on our paper. We are glad to hear that all reviewers agree our paper provides interesting and novel insights for the theoretical understanding of transformers. We have addressed all your questions in detail in our individual responses and have updated the manuscript accordingly. We would greatly appreciate it if you could review our responses and the revised manuscript to see whether you are satisfied, and let us know if you have any further comments or suggestions. Below, we give an overview of the major changes we have made in the revised manuscript. 

In particular, we appreciate the valuable comments from Reviewer BJ8K and Reviewer tEUD regarding the absence of a formal definition of “group sparsity” in our original manuscript. We have included a formal definition of "group sparsity" from lines 105 to 120 in the revised manuscript, and we would like to clarify that "group sparsity" is a popular concept in statistics and the goal of our work is to study the capability of simple one-layer transformers in solving the classic statistical problem of “linear classification with group sparsity”. We have also provided further explanations to clarify the connections between “linear classification with group sparsity” and our data model in Definition 2.1 (Definition 3.1 in the original manuscript). We believe that the formal definition and the additional explanations can address your concerns and clarify the motivation and significance of our study.

Additionally, we have conducted further experiments on synthetic data as well as real data, to demonstrate our theoretical conclusions on high-dimensional data,  as suggested by Reviewer BJ8K. The results on synthetic and real data are given in Appendices F and G respectively. We would like to remark that all the experiment results align with our theoretical conclusion, offering more robust empirical evidence to support our findings.

We look forward to your further feedback and comments, and we are happy to address any remaining concerns.

$\ $

Best regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work provides a theoretical analysis of transformers, demonstrating their ability to learn structured data with group sparsity. It shows that a one-layer transformer trained via gradient descent can use attention mechanisms to focus on relevant variables while ignoring irrelevant ones. Additionally, a well-pretrained transformer can adapt efficiently to downstream tasks with limited samples, offering insights into how transformers excel in structured data learning.

This paper is borderline at this stage. Some of the reviewers have highlighted a few issues with the paper, and the authors seem to have addressed many of them in the rebuttal. I would encourage the authors to go through the reviews carefully and improve their submission for the next version of this paper.