In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization

Summary

This paper studies the in-context learning (ICL) in the linear regression setting with a Gaussian prior with a non-zero mean. They first prove that linear Transformer block (LTB) will enjoy a smaller approximation error than linear self-attention (LSA) in this setting, where the non-zero mean of Gaussian emerges. Second, they show that the core benefit brought by linear MLP is that LTB can implement the one-step gradient descent with learnable initialization (GD-$\beta$). Besides, all global minimizers in the LTB class are equivalent to the unique global minimizer in the GD-$\beta$ class. Finally, they investigate the non-convex dynamics of the GD-$\beta$ class and prove that it will converge to the unique global minimizer. Simulations on the GPT-2 support their theory.

Strengths

1. Clarity: The paper is presented very clearly and is easy to follow. 2. Significance: The paper characterizes the benefit brought by the linear MLP in the ICL for linear regression setting, which is undoubtedly an important topic in the ICL theory community. 3. Quality: The paper investigates the additional effect of linear MLP from approximation and optimization perspectives in detail, and the derivation is solid.

Weaknesses

1. In this paper, all non-linear parts in the transformer block (softmax and ReLU) are dropped. However, similar works on the non-linear transformer block have emerged in the ICL theory. [a] considers the ReLU MLP followed by a linear attention layer in the linear regression setting. [b] considers the standard transformer block in the in-context classification regime. Authors are suggested to discuss the potential impacts of these studies about non-linear transformers on the setting in this paper. 2. Though the LTB class enjoys the same global minimizer as the GD-$\beta$ class and the training dynamics of the GD-$\beta$ class can converge to the global minima, whether the training dynamics of the LTB class can converge to the global minimizer is not discussed by the authors. The authors should clarify this point. [a] Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape, ICML, 2024 [b] Training Nonlinear Transformers for Efficient In-Context Learning: A Theoretical Learning and Generalization Analysis, ICML, 2024

Questions

1. In the LTB class, the paper does not merge $W_K^TW_Q$ and $W_P^TW_V$ into one matrix parameter, respectively. Are there any benefits? Does it will improve the approximation ability of the LTB? 2. In line 272, I think the $\Gamma$ is a precondition but not equivalent to the Newton step. Is it a typo? 3. The pretraining goal of this paper is different from the autoregression training in practice. Are there any potential approaches to study the autoregression pretraining and few-shot regression inference?

Rating

7

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have adequately addressed the limitations.

Summary

his submission extends the earlier theoretical analyses of in-context learning with linear transformers by taking into account the data covariance and a non-zero mean on the task parameters. It is shown that a linear transformer with a linear MLP block (LTB) can implement the algorithm of one-step gradient descent with non-zero initialization. This additional flexibility of learned initialization allows LTB to match the Bayes optimal risk for linear regression, whereas the linear transformer studied in prior works cannot achieve such optimality at finite $M$. The authors also showed that the optimal one GD step solution can be obtained by training the parameters with gradient flow.

Strengths

In the past years, there have been many works that examined the ICL ability of linear transformers in the linear regression setting, where it is typically shown that transformers implement one preconditioned GD step on the squared error in-context. This submission goes beyond the prior results by showing that linear transformers can adapt to the non-zero mean of the task distribution with the aid of a linear attention block. This yields a separation in the ICL efficiency between transformers with or without an attention block, which is an interesting message in my opinion.

Weaknesses

I have the following concerns and questions: 1. Firstly, the obvious limitations of the analysis: the problem setting is still linear, and the authors only showed adaptivity of ICL to the task mean but not the full prior (see point below). Given the additional linear block, the approximation gap is not really surprising. The optimization analysis directly assumes the GD-$\beta$ parameterization, and it is not clear if gradient descent on the standard LTB reaches the same solution. Also, convergence is only shown at a population level with no finite-sample guarantees. 2. The Bayes optimal estimator in Lemma 6.1 is a generalized ridge regression estimator where the $\ell_2$ regularization is anisotropic and adapts to the prior. On the other hand, when $\beta^*=0$, the GD-$\beta$ solution in Theorem 5.2 reduces to the one-step GD estimator studied in prior works, which does not incorporate the prior covariance. It is not clear to me why these two estimators always have the same statistical efficiency, as it is well known that anisotropic shrinkage provides a substantial advantage (Wu and Xu 2020) (Li et al. 2023). I guess the conclusion in Corollary 6.2 relies on $tr(H\Psi)$ being bounded. Can the authors comment on the optimality of GD-$\beta$ if this assumption is not satisfied? (Wu and Xu 2020) *On the optimal weighted $\ell_2$ regularization in overparameterized linear regression*. (Li et al. 2023) *Transformers as algorithms: generalization and stability in in-context learning*. 3. Some minor questions: * On page 5, why does the "scratchpad" take the specific form of (3.1)? * For the experiment reported in Table 1, what is the size of the trained model? Is the curriculum learning strategy in (Garg et al. 2022) also used? --- **Post-rebuttal update:** the authors’ response addressed most of my concerns. I found that one of my criticisms is based on a misunderstanding of the optimal solution that the authors analyzed. I have increased my score accordingly.

Questions

See Weaknesses above.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

See Weaknesses above.

Summary

In this paper, the authors investigate the in-context learning (ICL) of a linear transformer block (LTB), i.e., a single block comprised of a linear self-attention layer and a MLP layer. The task considered is ICL of linear regression with a Gaussian prior. Unlike earlier works that studied this problem, a Gaussian prior with a non-zero mean is used for generating the ICL tasks. For this setting, it is shown that the MLP layer is essential, as using only a self-attention layer results in a weaker approximation. Furthermore, the LTB at minimum is demonstrated to have a specific structure akin to one-step gradient descent with a learnable initialization, referred as GD-$\beta$. The optimisation and statistical properties of the GD-$\beta$ estimator is also widely studied.

Strengths

a) In-context learning of statistical tasks particularly the linear regression has received particularly interest. The authors have identified an interesting and relevant modification to this problem setting when the ICL tasks are generated from a more general setup Gaussian prior with non-zero mean. b) In this setting, an MLP layer and a skip connection are needed for approximation which is a novel finding. Even with these additional parameters it is shown that minimum of the task can be still be characterised and it corresponds to GD with learnable initialisation. c) The statistical and non-convex optimization aspects of the GD-$\beta$ are interesting from a technical standpoint.

Weaknesses

The paper attributes that property of learning the MLP layer, however the skip connection is also equally important. Hence it is a bit of an misleading claim to attribute this to the MLP layer. Overall, the model restricted to only linear attention and also a single layer.

Questions

1) What happens if we include positional encoding (P) are also included? It is possible to get an additional $P^{\top} W_K^{\top} W_QP$ term can come from the attention and can this overcome the need of positional encoding? 2) Does the optimisation dynamics of GD-$\beta$ offer any insights on the optimisation of the LTB ? 3) Does the results qualitatively change if the input sequence is generated from uncentered Gaussian $\mathcal{N}( \mu, H )$ instead of zero mean Gaussian prior ? Minor corrections: 1) In l.251, it should be $ \Tau^{\top} $ in the expression of $W_K^{\top} W_Q$. 2) In l.263, an additional bracket in the expression of $\Omega$.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

This paper studies the in-context learning of linear regression with a Gaussian prior that has a finite, non-zero mean across tasks. Specifically, the authors show that a linear transformer block (linear attention with an MLP layer) can achieve near-optimal Bayes risk for this task, whereas a linear attention-only block suffers from an approximation error. They further demonstrate that an LTB hypothesis class can be well-represented by a subset class of one-step gradient descent estimators, GD-$\beta$, for the ICL task.

Strengths

This work delves into theoretically understanding the very relevant and hot topic of ICL. The problem motivation is clear and concise, and I really like how the authors essentially redo the previous results of LSA and GD-$0$ estimators with LTB and GD-$\beta$, demonstrating the effective correspondence between them and showing that they achieve near Bayes optimal ICL risk. Overall, I think the authors do a great job of portraying the benefits of the MLP component in transformers, and it’s a good theoretical contribution, as most papers involve theory only from the LSA perspective.

Weaknesses

I don’t see any major weaknesses in the paper as such, but I do have a few questions (see questions).

Questions

Have the authors seen experimental evidence that trained LTB models converge to GD-$\beta$ models, similar to how LSA converges to GD-$0$? I believe the correspondence established between the two is excellent, and theoretically studying the convergence is probably another work in itself, but seeing some experimental evidence would have been really cool. I believe the restrictions imposed on $d_k, d_v$ for the hypothesis classes are to ensure sufficient overparameterization, which is pretty common for theoretical works. However, I think in practice these dimensions are typically set to be less than $d$. Are these necessary conditions, or do the authors believe there are ways to override them for their analysis? The constructions for Lemma $5.1$ and Theorem $5.3$ both operate on pairs of matrices composed together (if not more), which is also the common practice in the analysis that I’ve seen. Any reasons why the authors didn’t opt for analyzing with merged matrices, especially when the inner dimension of the product $W_p^TW_v​$, that is $d_v$​, is assumed to be larger than $d$? It doesn’t even enforce any additional rank constraints on the product matrix. It was interesting to see that in your GPT-2 experiments, the full transformer model (with all the non-linearities) also relied heavily on the MLPs. How much of this do you think can be attributed to the presence of softmax only? Did you try experiments where you used LSA in GPT-2 training with MLPs?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have addressed the limitations.

Official Comment by Reviewer XdPR

Thank you for providing the rebuttal. The authors adequately addressed my concerns, so my confidence level regarding my review of this paper has increased.

Reply to Author Rebuttal

Thank you for the rebuttal and for addressing my questions. The rebuttal did not change my initial impression of the paper: it presents an interesting point but also has its share of limitations. Therefore, I maintain my rating.