In-Context Learning with Representations: Contextual Generalization of Trained Transformers

Response to Reviewer vG3D

Thank you for your insightful review. We've added experiments to validate our theory in the supplementary pdf and our global response. Below we address your concerns. If our responses resolve your questions, we'd highly appreciate your consideration in raising the score. Certainly, please don't hesitate to request any further clarification.

(W1) Assumption $H\geq N$ too strong. The single-head works in Table 1 do not need this condition for few-shot ICL.

We recognize the importance of weakening this assumption and will explore this in future work. Meanwhile, we want to make the following clarifications.

• As explained in Section 3.3, it is generally impossible to train a single-head shallow transformer to succeed in the task we consider. Our task involves more complex non-linear regression with representation learning, which fundamentally requires multi-head attention.

• Note that $H\geq N$ is a necessary and sufficient condition for the initialization condition (Assumption 3) to hold. To see why it is necessary, note that $B_k^{(0)}$ is an $N$-by-$H$ matrix, and has full row rank only when $H\geq N$.

• Assumption 3 is crucial to our proof: it guarantees that $\zeta_0$ (defined in (12)) is positive, which in turn ensures $\sigma$ (the PL parameter in (44)) is positive.

• However, it's important to note that Assumption 3 may not be a necessary condition for the convergence of the transformer. To weaken (or avoid) Assumption 3, other proof frameworks are needed. But due to the highly non-convexity of (even shallow) transformers, it's somewhat unavoidable to make other strong assumptions (see the examples we list in the next bullet point).

• Though the single-head works in Table 1 do not need $H\geq N$, they often rely on other stronger assumptions even for tasks without representation learning, see Section 1 and 3.1 for details. To summarize:

  • Huang et al. [2023], Li et al. [2024], Nichani et al. [2024] assume that the tokens are pairwise orthogonal, which implies the the token dimension $d$ is no smaller than the dictionary size $K$, while in practice $K\gg d$.

  • Huang et al. [2023] and Chen et al. [2024] require the context length to be sufficiently large, so they are not analyzing few-shot ICL.

  • Li et al. [2024] requires the width of their model to be larger than the square of the number of data.

(W1) This means that Transformers cannot handle long contexts without a large number of heads/(W2) No experiments are conducted to verify the theory...can you show if $H<N$, the learning will fail, while if $H\geq N$, the learning becomes successful by a one-layer Transformer?

• Our above response indicates $H\geq N$ is only a sufficient (but may not be necessary) condition used to guarantee the convergence of the transformer. Therefore, our results don't indicate transformers cannot handle long contexts without a large number of heads in practice. However, shallow transformers have limited expressive power, which necessitates more attention heads to solve complex tasks. To support this, we have put empirical results in the supplementary pdf, please refer to Figure 3 and Figure 5 and the corresponding description in our global response.

(W3) Statements on [Li et al., 2024] are wrong.

• You are correct, and we apologize for the mistake and will fix it in the updated paper. Thank you for your careful review.

(W4/Q1)typo in (63)

• Thank you for your sharp observation. You are right -- $\eta_Q$ should be replaced by 1 in (63). We apologize for the typo and will fix it in the updated paper.

• Let us emphasize that this typo does not affect the subsequent results, and the rest of our proof remain valid:

  The revised (63) reads:

Combining this with (38), we have

When $\eta_Q\leq 1/L$, we have $L\eta_Q/2-1\leq -1/2$, which gives the first part of (64): $\ell(\xi^{(t)})-\mathcal{L}^\star\leq \ell(\xi^{(t-1)})-\mathcal{L}^\star - \frac{\eta_Q}{2}||\nabla_\xi \ell (\xi^{(t-1)})||_F^2.$ Note that this is the only place where (63) is used.

(W4) derive linear convergence of the loss via only PL condition and the smoothness of the loss function

• The combination of PL condition and smoothness is indeed sufficient for linear convergence of gradient descent. See for example Theorem 13.2 in this note: https://www.stat.cmu.edu/~siva/teaching/725/lec13.pdf

• This approach is commonly used in deep learning optimization literature, as PL condition has been proven to hold for over-parameterized neural networks. See [1-3] for example.

(Q2) initialization of $w_{hk}$

• $w_{hk}$ is initialized to be 0, as stated in line 226 of our paper.

Is your bound derived for the worst case of all the possible $f$?

• our bounds apply for any possible $f$, which includes the worst case one. The key insight is that our theoretical guarantees depend on the function's outputs on the dictionary $\mathcal{V}$, i.e., in our analysis, we work with the finite-dimensional matrix $\widehat Z=(f_i(v_j))_{i\in[m],j\in[K]}\in R^{m\times K}$ defined in line 485, rather than the full functional form of $f$.

[1] Liu, C., Zhu, L., and Belkin, M. (2020). Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.

[2] Karimi et al. (2016). Linear Convergence of Gradient and Proximal-Gradient Methods Under the Polyak-Łojasiewicz Condition.

[3] Loizou et al. (2021). Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence.

Thank you very much for your quick response! Below we change $N$ and find the smallest $H$ that achieves a training loss $\leq$ 0.5 within 100 iterations. We set $K=200, d=100, m=20, \tau=0.01$. Our result is as follows:

This result verifies our theoretical finding that with a larger $N$, $H$ should be larger to ensure a small loss.

If our response resolves your questions, we will highly appreciate your consideration in raising the score. Certainly, we are more than happy to answer your further questions.

Response to Reviewer MzSs

Thank you for your time in reviewing our paper. Below We address your points. If our responses adequately address your concerns, we would be grateful if you could consider increasing your current score. And we are happy to answer your additional questions.

(W1) The optimal transformer configuration $(\theta)$ is not explicitly stated; only the functional form of the output is given, which makes it difficult to the interpret the result in the parameter space.

• As is stated in line 316 and line 482-485, $\theta=${$Q_h,w_h$} is optimal if and only if it satisfies:

Dear Reviewer MzSs,

We've taken your initial feedback into careful consideration in our response. Could you kindly confirm whether our responses have appropriately addressed your concerns?

If you find that we have properly addressed your concerns, could you please kindly consider increasing your initial score accordingly? Please let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors

Dear Reviewer MzSs,

Thank you very much for your response!

In addition to what you mentioned as 'memorize' the basis vectors, our results contain two more points: (i) after 'memorizing/learning' basis vectors, the ICL implements ridge regression to use such basis information, which is a new characterization; (ii) such tranformers having those properties (memorizing basis vectors and implementing ridge regression) can be learned naturally by gradient descent.

Many thanks, Authors

Response to Reviewer TyrV

Thank you for your valuable feedback. We agree that empirical validation is crucial and have included experimental results in the supplementary PDF and the global response. Note that doing in-context training on real datasets is highly demanding and usually requires very large models trained on a lot of different datasets [3]. Therefore, it's standard practice for the line of works (e.g., [4,5]) that analyze transformers' in-context ability to experiment on the synthetic dataset. Below we address your other points. If these clarifications resolve your primary concerns, we would highly appreciate your consideration of increasing your score. Certainly, please don't hesitate to request any further clarification.

(Weaknesses/Q1) discussion on the practical implications of the theoretical findings

Besides those discussed in the global response, one interesting implication is that our paper suggests assigning different learning rates for different parameter blocks, which has been shown to work well by some very recent work (e.g. [1,2]), where the authors show sgd (which assigns the same learning rate to all blocks) doesn't work well for transformers, but it's also not necessary to keep different learning rates for all parameters as Adam does.

What are the potential impacts of the remaining assumptions and simplifications (e.g., specific initialization conditions, structure of the transformer) on the generalizability of your results?

Thank you for bringing up this important question. We'll incorporate the following discussion in our paper:

• While our current analysis focuses on shallow transformers, we believe our methodology has the potential to be extended to deeper architectures. The core of our proof technique, which combines the Polyak-Łojasiewicz (PL) condition with smoothness to demonstrate convergence (as detailed in the proof of Theorem 1), could potentially be generalized to deeper transformers. Therefore, we expect that our key results, such as the limit point of the output of the transformer (see (19) in Theorem 2), the explanation of the transformer's ICL ability and the analysis of its generalization ability would also apply to deep Transformers.

• The main challenge that we need to overcome in extending the proofs to deeper architectures lies in estimating the PL coefficient $\sigma$ and smoothness coefficient $L$. To this end, it is possible to leverage our techniques for one-layer transformer and carefully conduct a block-by-block analysis potentially via induction arguments.

• As we comment in line 217-218 and Proposition 1 in the paper, our initialization condition can be easily achieved via random sampling from Gaussian distribution.

• It's worth noting that we assume Gaussian noise for simplicity and show that the transformer decides its generating template by performing ridge regression (see Section 1.1 and Theorem 2). Our technique could potentially be extended to analyze other types of noise. For example, we could hypothesize that if the noise is Laplacian, the transformer might learn to perform Lasso regression instead. We leave this exploration as future work.

[1] Y Zhang et al. (2024). Adam-mini: Use Fewer Learning Rates To Gain More.

[2] Y Zhang et al. (2024). Why Transformers Need Adam: A Hessian Perspective.

[3] S Min (2022). MetaICL: Learning to Learn In Context.

[4] S Garg et al. (2022). What Can Transformers Learn In-Context? A Case Study of Simple Function Classes.

[5] S Chen et al. (2024). Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality.

Dear Reviewer TyrV,

We've taken your initial feedback into careful consideration in our response. Could you kindly confirm whether our responses have appropriately addressed your concerns?

If you find that we have properly addressed your concerns, could you please kindly consider increasing your initial score accordingly? Please let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors

Response to Reviewer dubr

Thank you for your review and positive comments. We've included some experiments to verify our theoretical findings in the supplementary pdf and the global response. Below we address your other points. If our responses resolve your concerns, we'd highly appreciate your consideration of increasing your current score. Certainly, please also let us know if you have further questions.

simplified model

• Our analysis indeed focuses on shallow transformers, but we believe our methodology lays a foundation for extending to deeper architectures. The core of our proof technique, which combines the Polyak-Łojasiewicz (PL) condition with smoothness to demonstrate convergence (as detailed in the proof of Theorem 1), provides a framework that could potentially be adapted for deeper transformers. We anticipate that our key results, such as the limit point of the transformer's output (see (19) in Theorem 2), the explanation of the transformer's ICL ability, and the analysis of its generalization capability, would extend to deeper models, albeit with more complex derivations.

• However, extending the proofs to deeper architectures presents significant challenges, particularly in estimating the PL coefficient $\sigma$ and smoothness coefficient $L$. As evident from our proofs, even for a one-layer transformer, this computation is highly non-trivial and requires a meticulous block-by-block analysis. For deeper models, this analysis would become substantially more intricate, requiring novel mathematical techniques to handle the increased complexity.

• To the best of our knowledge, there's no existing work analyzing deep transformers' training dynamics from an optimization perspective. Most theoretical works, including those in Table 1, focus on shallow transformers due to these challenges. Developing rigorous analytical methods for deep transformers remains an open and exciting direction for future research in this field.

focus on regression

• Our framework considers non-linear regression tasks, which are significantly more complex than the linear tasks analyzed in previous theoretical works (as highlighted in Table 1 of our paper). We expect that such a contribution within regression framework is valuable to the community.

• A key aspect of our work is the analysis of how transformers learn representations (see Section 3.3). This is fundamental to many other machine learning tasks, including feature extraction, translation, and generalization. Our insights into how transformers extract and memorize "inherent information" of basic function maps during training (lines 312-317) could be valuable for understanding other more complex tasks. Besides, our novel insights into how transformers acquire contextual generalization ability with limited data and handle underdetermined templates is relevant to a broader range of in-context learning tasks beyond regression. We leave the extensions of our settings as important directions for future work.

Dear Reviewer dubr,

We've taken your initial feedback into careful consideration in our response. Could you kindly confirm whether our responses have appropriately addressed your concerns?

If you find that we have properly addressed your concerns, could you please kindly consider increasing your initial score accordingly? Please let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors