On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures

The technical contribution beyond [Zhang et al., 2023a] is not clear to me. You consider the classification problem with a different loss function. But why is it challenging enough compared with [Zhang et al., 2023a]. It is better to include related discussions in the paper.

Compared to [1], we use the cross-entropy loss for the classification problem. Thus, a critical distinction to [1] is that our model does not have a closed-form expression of the global minimizer, while [1] does. This distinction adds a significant challenge and leads to a different technical approach compared to [1]. By analyzing the Taylor expansion at the stationary point, we establish the convergence properties of $W^*$. Moreover, [1] considered optimizing the transformer using gradient flow. In contrast, our work proves the convergence of optimizing the transformer with the more practical gradient descent method. We will add more related discussions in our revised paper.

The practical insight of this work is not clear. For example, how can the theoretical analysis be used to improve ICL in practice? Are the theoretical conclusions aligned with any existing empirical finding? Or can the theory be used to explain any theoretical finding?

Yes, our theoretical conclusions align with existing empirical findings. For example, in Figure 1, we conducted experiments of single-layer and multi-layer transformers for in-context classification of Gaussian mixtures. We can notice that both transformer models' ICL inference errors decrease as training prompt length ($N$) and test prompt length ($M$) increase, and increase as the number of Gaussian mixtures ($c$) increases. These behaviors are consistent with our theoretical claims. Moreover, the results also indicate that some of our insights obtained from studying this simplified model may hold for transformers with more complex structures, and studying this simplified model can help us have a better understanding of the ICL abilities of transformers. We hope this paper can provide valuable insights into the theoretical understanding of the ICL mechanisms of transformers. These insights may be helpful for potential architectural design and building safe and explainable AI systems.

The writing needs some improvement. It is better to include some remark after each Theorem to provide an intuitive explanation. This can help readers understand the key point of each result.

Thanks for your suggestion. We will add more marks and intuitive explanations in the revised paper.

References:

[1] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023a.

[8] Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.

[9] Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402.15607, 2024.

[10] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint arXiv:2310.08391, 2023.

I suggest the author to condense Section 4. Remarks E.1 and G.2 should be included as part of the main body.

Thanks for your suggestions. We will modify them in the revised paper.

It is okay to assume homo-scedasticity (same $\Lambda$ for both classes), but the assumption on the means having the same $\Lambda^{-1}$-weighted norms is hard to justify. I guess this assumption is made for the convenience of the proof. If so, the authors should be upfront about this and explain why the problem is intractable without this assumption.

Yes, this assumption is made for the convenience of the proof. Actually, in Remark E.1 and G.2, we explained why we need this assumption and showed that if this assumption is not satisfied, the single-layer transformer cannot correctly perform the in-context classification tasks. We also verified this in experiments (Section 5.1, Figure 2) that the single-layer transformer cannot perform well for varying norms.

It seems that the derivation is about to connect the pre-trained attention model to linear/quadratic discriminant analysis (LDA/QDA) but then abruptly stops. Actually, I think it is very helpful to describe what kind of learning algorithm is implemented by the pre-trained transformer. Setting $W = 2\Lambda^{-1}$ implements the LDA decision criterion?

Yes, we are connecting the pre-trained attention model to LDA and setting $W = 2\Lambda^{-1}$ implements the LDA decision criterion. We will add more discussions about the connection between the pre-trained attention model and LDA in the revised paper. Thanks for your suggestion.

What happens when you stack multiple layers?

In our experiments (Section 5.1, Figure 2), we showed that multi-layer transformers have better robustness for varying covariances and norms. We believe developing a better understanding of multi-layer and more complex transformers is an intriguing direction for future research.

The well-conditioned property (9)

es, we prove the that there exist $\alpha>0$ and $l<\infty$ that satisfy $(9)$. In Lemma D.2, we prove that in a compact domain, the strong convexity parameter $\alpha$ is larger than 0. In Lemma D.5, we show that $l\leq \frac{1}{4}\sum_{i\in[d^2]}E[(p_{t_1(i)}q_{t_2(i)})^2]$, where $p,q$ are some combinations and rotations of Gaussian random variables. Since Gaussian random variables have finite second moments. Thus, $l<\infty$. We will clarify this in the revised paper.

I noticed that in Section D.1, the description of D.3 does not match the actual lemma

What we exactly did in Lemma D.3 matches what we described in Section D.1. We are analyzing the Taylor expansion of $L(W)$ in Lemma D.3. In line 892-894, we show that as $N\to\infty$, our loss function $L(W)$ point-wisely converges to $\widetilde{L}(W)$. In line 926, we also show that as $N\to\infty$, the global minimizer $W^*$ converges to $2\Lambda^{-1}$. The logic of the proof here is that we first show the loss function $L(W)$ converges point-wisely to $\widetilde{L}(W)$, which implies that the global minimizer $W^*$ converges to $2\Lambda^{-1}$. Then, given the property that the global minimizer $W^*$ converges to $2\Lambda^{-1}$, in Lemma D.4, we can derive a tighter convergence rate for $W^*$. We will clarify this in the revised paper.

Figure 3: How many classes are there?

We mentioned in the paper that we compare the classification of three Gaussian mixtures.

For the "inference error" are you referring to the TV distance in equation 3? I don't think this notion is directly applicable to k-nearest neighbor or SVM.

We used KNeighborsClassifier and SVC from sklearn, and they provided the prediction probability for given test data. We then use the prediction probability to calculate the TV distance. See the API of sklearn for details.

I find this comparison to be unfair since the transformer models have been pre-trained on a lot (what is the exact number?) of data

The exact number of training data can be found in H.2. Experiment Details. It is somewhat unfair to compare Transformer models with classical models. However, the goal of our experiments is not to make a fair comparison between Transformers and classical methods, nor to prove that Transformers are superior. Instead, our main purpose is to demonstrate that trained Transformers have in-context learning capabilities and perform no worse than classical methods.

Why is LDA not part of the classical baseline? You should include your source code

We upload our code and additional results in the following link. https://anonymous.4open.science/r/In-Context-Classification-of-Gaussian-Mixtures-2374 We added LDA in LDA_Comp.png

Adding a few references on the classical methods for Gaussian mixtures, e.g. linear discriminant analysis, would be very helpful.

Thanks for this comment. We will cite and discuss a few references about Gaussian mixtures and linear discriminant analysis in the revised paper.

By assuming an over-simplified attention structure...

Setting some parameters to fixed values and considering spare form parameters is commonly used in ICL theory papers [7,8,9,10], and we adopt a similar parameterization as in [7,8,10]. Even for this simplified structure, our analysis for the convergence of $W^*$ is nontrivial. For example, in order to get a tighter bound for $W^*$, in Lemma D3, we first show the loss function $L(W)$ converges point-wisely to $\widetilde{L}(W)$, which implies that the global minimizer $W^*$ converges to $2\Lambda^{-1}$. Then, given the property we prove in Lemma D3, we can derive a tighter convergence rate for $W^*$ in Lemma D4 than the bound we got in Lemma D3.

total variation cannot be used to evaluate test performance.

We thank the reviewer for this comment, but we believe that total variation distance can be used to evaluate test performance in our setting. We will illustrate this point using a binary case example: Suppose for data $x_q$, the conditional distribution of the ground truth label $y$ is $P(y=1|x_q)=p$, $P(y=-1|x_q)=1-p$. Suppose the prediction of the model is $z$, which depends on $P=(x_1, y_1, ...,x_M, y_M, x_q)$. Suppose the conditional distribution of $z$ is $P(z=1|P)=q$, $P(z=-1|P)=1-q$. Since $(x_i, y_i)$ are i.i.d. sampled, for given $P=(x_1, y_1, x_2, y_2, ...,x_M, y_M, x_q)$, $z|P$ and $y|x_q$ are independent. Thus, the example given by the reviewer is not valid in our setting. If we consider the accuracy, we can calculate it as $Acc(y, z)=pq+(1-p)(1-q)=1+2pq-p-q$. We can notice that, if $p>1/2$, the best $q$ to maximize the accuracy should be $q=1$ and if $p<1/2$, the best $q$ to maximize the accuracy should be $q=0$. Thus, if the output of the model $\widehat y_{out} | P$ is the same as $P(y=0|x_q)$, we can output $z$ based on the value of $\widehat y_{out} | P$ to maximize the possible accuracy. The smaller the difference between $\widehat y_{out} | P$ and $P(y=0|x_q)$, the better the model performs. Actually, the total variation distance in our case measures this difference.

What concerns me most is their decision to drop all terms and factors that are independent of $N$ during the training process...

We'd like to clarify that we did not drop all terms and factors that are independent of $N$ during the training process. For example, in Theorem 3.3 Eqn (10), we retained the coefficient of the $1/N$ term and showed that this coefficient is affected by the distribution of distance between two means $\mu=\mu_{\tau, 1}-\mu_{\tau, 0}$, the query, and the covariance matrix $\Lambda$. Intuitively, the classification problem becomes easier when the distance between the two class means $\mu$ increases and the variance $\Lambda$ decreases; conversely, it becomes harder when $\mu$ is small and $\Lambda$ is large. Eqn (10) reflects this intuition. For instance, when $\mu$ is large and $\Lambda$ is small -- corresponding to an easier classification problem -- the parameter $a$ becomes larger, and the derivative $\sigma'(a)$ decays exponentially with increasing $a$, thereby reducing the overall training error. On the other hand, when $\mu$ is small and $\Lambda$ is large -- corresponding to a harder classification problem -- the training error increases accordingly. This theoretical prediction aligns well with our intuition. We did not explicitly present results in terms of the dimensionality $d$, because intuitively, $d$ is not the primary factor determining the classification difficulty. Instead, it is the relationship between $\mu$ and $\Lambda$ that plays the central role. Our theory captures this key insight effectively.

I also do not understand why the query token $q$ can appear on the RHS of the formula of Theorem 3.6 when you are taking the expectation over the query token pair.

We did not take the expectation over the query token $q$. You can check Theorem 3.6. Our expectation is take over $x_i, y_i; i=1,.., M$, which does include $q$.

I find that Assumptions 3.2 and 3.5, which state that the mean vectors $\mu_1$ and $\mu_2$ have the same $\Lambda^{-1}$ norm, are somewhat strong.

See our second reply to reviewer aULa.

There are several theoretical studies on the optimization of one-layer transformers that consider training $W^v$ and $W^{KQ}$ simultaneously, even in the more challenging softmax attention setting [2, 3, 4, 5]. There are also results on general convergence guarantees of transformers such as [6]. In addition, [7] considers almost the same question as this paper. I suggest that the authors should compare their results with these works.

Thank you very much for providing those related works. We will cite and discuss them in the revised paper.

Is there a typo in line 719? Should it be $\frac{1}{4N^2}$ instead of $\frac{1}{4N}$?

No. This is not a typo. You can verify that the result is $\frac{1}{4N}$.

References can be found in our reply to reviewer M6vU.

We conducted additional experiments and uploaded code in https://anonymous.4open.science/r/In-Context-Classification-of-Gaussian-Mixtures-2374

Fundamentally, I felt that the lack of serious comparison to just the linear-regression approach is a big weakness. Treating classification as linear regression isn't optimal, but it can work decently well especially in the kinds of settings here. So, what parts of the analysis are picking up extra nuances of what single-layer linear-attention transformers can do beyond linear regression and what parts are just casting this problem into its shadow linear-regression form in disguise? I can't tell from the discussion. But this is what I really want to know.

In SNR_LR_Comp.png, we conducted experiments comparing the transformer with linear regression for classification tasks. We can notice that trained transformers, in general, have better performance than linear regression, especially when the prompt length is small. This is not that surprising since the trained transformers somewhat learned $\Lambda$ during training.

Given the work in Zhang, Frei, and Bartlett (2023a), this feels like the critical question.

Compared to [1], we use the cross-entropy loss for the classification problem. Thus, a critical distinction to [1] is that our model does not have a closed-form expression of the global minimizer, while [1] does. This distinction adds a significant challenge and leads to a different technical approach compared to [1]. By analyzing the Taylor expansion at the stationary point, we establish the convergence properties of $W^*$. Moreover, [1] considered optimizing the transformer using gradient flow. In contrast, our work proves the convergence of optimizing the transformer with the more practical gradient descent method. We will add more related discussions in our revised paper.

First straight-up least-squares knowing nothing. Second, one that has learned the covariance $\Lambda$ by magic. Third, one that has learned the covariance $\Lambda$ using the same training data provided to the transformers?

In SNR_LR_Comp.png, we compared the trained transformer with straight-up linear regression, knowing nothing. Actually, as the reviewer aULa pointed out, our pre-trained attention transformer approximately implements LDA. If a LDA model learned the covariance $\Lambda$ by magic or learned the covariance $\Lambda$ using the same training data provided to the transformers, this LDA model should have better or comparable performance as the trained transformer.

If I think about how linear regression for classification works in a high-enough-dimensional space, when the means are large, it works quite decently. Anyone who studies mixture-models as toys for classification knows that the signal-to-noise ratio is important. Given that your plots have inference errors dropping to by $5\%$ or lower, it suggests that the SNR is decent. So how do the results change if we vary the SNR?

In SNR_LR_Comp.png, by changing the magnitude of the covariance matrix of testing data ($\Lambda, 4\Lambda, 16\Lambda$), we compared the accuracy of trained transformers and linear regression with different rates of SNR. Our results show that smaller SNR leads to worse accuracy for both the transformer model and the linear regression model.

References can be found in our reply to reviewer M6vU.