Response to Reviewer c1Lc

We thank Reviewer c1Lc for the valuable comments.

Weakness1: key contribution

Thanks for the insightful suggestion. We agree that we do not have direct evidence that "ICL by AR pretraining is more difficult than ICL by few-shot pretraining" because the data model is changed. We will revise it to "ICL by AR pretraining is quite different from that by few-shot pretraining" following your suggestion.

However, we emphasize that our key contribution is understanding the ICL ability of AR pretraining, rather than changing the data model. First, adopting the new data model is not our target but a necessary approach to investigate AR pretraining because the few-shot regression dataset in the previous ICL theory is not suitable since each input $x_{i+1}$ does not have a relation with the context. Second, we additionally need to scale the embedding strategy, the attention model, and the loss function to perform the full AR pretraining. All these parts are not well studied in the literature. Third, we are the first to investigate the training dynamics and the mesa-optimization hypothesis in this setting, which deepens people's understanding of autoregressive transformers.

Besides, due to the non-trivial nature of the AR pretraining, we think our results will highlight the role of other components in practical transformers (more attention layers, MLP, etc), and boost researchers to study autoregressive transformers.

Weakness2: artificial setting

Thanks for the advice. We note that all settings adopted in this paper are reasonable and have attracted attention from the recent ICL theory, which is also agreed by Reviewer WaqB. We clarify key points here, and we will discuss these in more detail in the final version.

1. Single-layer one-head linear attention has been a popular setting in the recent ICL theory for few-shot regression [24,25,28,29,27]. As for the AR pretraining, the model in this paper is more complex and has also been adopted in recent works [16,17]. However, they do not investigate its training dynamics thus the convergence properties are still unclear.
2. Linear data model is proposed in [16] and simplified in [17] for tractable analysis. We use the same data model as [17], while they only analyze the loss landscape under the diagonal assumption for model parameters.
3. Token embeddings in this paper have been used in [16,17]. It is a natural extension of the embeddings in existing ICL theory because we need to predict each historical token here. In fact, practical transformers do learn similar token construction in the first softmax attention layer (see Fig. 4B in [16]).
4. Population loss is a common setting in ICL theory which considers loss landscape and training dynamics [17,24,28,49,51,52].

Implications on more complex models. We will revise the claim to "We believe that our analysis can give insight into the establishment of theory in settings with autoregressive transformers". Besides. analogous to the development of ICL theory for few-shot regression (MLP [52], softmax attention [30], etc), future works considering more complex transformers can adopt the same data model, token embeddings in this paper, and try to use the similar proof technique to derive the training dynamics.

Weakness3: initialization & keeping diagonal structure

First, the diagonal initialization is reasonable and we further conduct simulations with standard initialization. Second, we note that the findings of Assumption 4.1 and derivation are not trivial. Please see the details in the common concern 1 and 2, respectively.

Weakness4: gradient clipping

We note that our target in Theorem 4.3 is to complement the theoretical result in [17] from the perspective of optimization, rather than provide a theory for practical gradient clipping. We will change "gradient clipping" to "gradient masking" in the final version to avoid confusion.

Weakness5: Assumption 4.2 and Theorem 4.2

We think the sufficient and necessary condition Assumption 4.2 might not be further simplified and we have discussed some properties of $x_1$ in line 238-245. For Theorem 4.2, we will follow the reviewer's suggestion to move Theorem 4.2 to the Appendix.

Weakness6: effect of $d$

In the final version, we will clarify that we fix $d$ as a finite number and focus on the asymptotic property w.r.t. $T$. Besides, we conduct additional experiments in smaller-context scenarios $(d=5, T=3/5/10, x_1\sim N(0, I_d))$. As a result, the product $ab$ does converge to the results in Theorem 4.1. The representative results can be found in the latest uploaded PDF.

Weakness7: Proposition 4.2

This is a misunderstanding. In the final version, we will clarify that Proposition 4.2 does not depend on diagonal initialization. As we have mentioned in line 272-274, parameters that perform GD are not the critical points of the loss function. Therefore, the trained transformer will not perform vanilla gradient descent with any initialization.

Weakness8: rationality of mesa-optimization

We disagree with this with our highest respect. From the result of this paper, we can only conclude that one-step GD is not enough to learn the data distribution. However, even for more complex data ($W$ is not diagonal), [16] has empirically verified that multi-layer linear attention can perform multi-step GD to learn the data distribution. The theory for the multi-layer attention in AR pretraining setting is a meaningful future work.

Question 1: ratio of two vectors

As we have mentioned in line 326, we want to verify the correctness of Proposition 4.1, which claims that the prediction will converge to 1/5 of the true data. Therefore, we present the ratio and find it does converge to 1/5. In practice, we estimate this ratio by calculating the mean of the element-wise divide between two vectors.

We will clarify this detail more carefully in the final version.

Thanks for the helpful comments and positive feedback, as well as for increasing the score. We are glad that we have addressed the reviewer's main concerns.

For the claim of theoretical contribution, we agree with the reviewer's suggestion. We promise that we will put more emphasis on our results contributing a closed-form characterization of the limiting points of AR pre-training. We will also discuss the necessity and the limitation of the diagonal initialization more clearly in the final version.

For the complete experimental results in small-context ($T=3/10$) and standard-initialization ($\sigma_w=0.001/0.1$) scenarios, the claims in the common response still hold, and we will add them to the final version. We note that due to the space limitation, we only present the representative result in each scenario during the rebuttal period.

We thank the reviewer once again for the valuable suggestions and insightful comments, which definitely improve the quality of this paper a lot.

Response to Reviewer WaqB

We thank Reviewer WaqB for the positive support and valuable comments, which can definitely improve the quality of this paper.

Weakness 1: additional related work

Thanks for the helpful suggestion. We will cite and discuss the mentioned papers in the final version. We believe that the higher-order optimization hypothesis is very important for understanding practical ICL, especially in multi-layer cases.

Question 1: multi-layer and non-linear activations

Thanks for the nice suggestion. We discuss two parts respectively.

Multiple layers

We think we can explore this setting from expressive power and optimization/generalization perspectives, respectively.

1. Expressive power. We can try to construct the cases where multi-layer transformers perform multi-layer higher-order optimization (e.g., GD++, Newton's method) in the AR pretraining setting. We think [16] can provide many practical insights, and [a,b] can provide many theoretical insights.
2. Optimization/generalization. In this case, our work can be seen as a foundation. We can try to investigate the training dynamics and the generalization ability of the trained transformers. Here, we think [c] can also give some theoretical insights, which discuss the multi-layer attention in the regression setting.

Non-linear activations

We only discuss the softmax activation here because it is used most widely in practice.

1. [16] shows that there exists a strong connection between softmax attention and GD in the AR pretraining setting, but the mechanism is still unclear.
2. [16] also shows that in the AR pretraining, the first softmax layer can learn the copy behavior and obtain the embedding construction in this paper, but the mechanism is also still unclear. To this end, we think [51] can give some insights, which prove similar copy behavior of the softmax attention.

We will discuss these in detail in the final version.

[a] Transformers learn higher-order optimization methods for in-context learning: A study with linear models, 2023.

[b] How Well Can Transformers Emulate In-context Newton's Method?, 2024

[c] Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?, ICML, 2024

We thank the reviewer once again for the insightful comments and positive score, which give us much confidence to finish this rebuttal process.

Response to Reviewer dsNs

We thank Reviewer dsNs for the positive feedback and valuable comments, which can improve the quality of this paper.

Weakness1: initialization & negative results

Thanks for the nice suggestion. We discuss two concerns respectively.

Initialization

We argue the diagonal initialization is reasonable and further conduct simulations with standard initialization. Please see the details in common concern 1.

Negative results

The results show that one-step GD learned by the transformer can not recover the distribution, but this can be solved by more complex models. Even for more complex data ($W$ is not diagonal), [16] has empirically verified that multi-layer linear attention can perform multi-step gradient descent to learn the data distribution. Our work can be seen as a foundation for the theory of more practical transformers.

Weakness2.1: diagonal structure

Thanks for the suggestion. The findings of Assumption 4.1 and analysis are not trivial. Please see the details in the common concern 2.

Weakness2.2 & Q3: questions on Proposition 4.1

Thanks for the nice comment, which definitely helps us clarify Proposition 4.1!

Refined Proposition 4.1

First, we clarify the original intention of Proposition 4.1. We want to prove that in expectation w.r.t. data ($x_1$ and $W$), when the training sequence length $T_{tr}$ is large, the prediction $\hat{y}_ {T_ {te}}=W(\tilde{a}\tilde{b}\frac{\sum_ {i=1}^{T_ {te}-1}x_ ix_ i^*}{T_ {te}-1})x_ {T_ {te}}$ of the trained transformer converges to $1/5 * Wx_ {T_ {te}}$, rather than the true target $Wx_{T_{te}}$. In the submitted version, we prove this by claiming $E_{x_1}[\tilde{a}\tilde{b}\frac{\sum_{i=1}^{T_{te}-1}x_ix_i^*}{T_{te}-1}] \to \frac{1}{5}I_d$, which is not rigorous as the reviewer points. We modify the proposition rigorously as follows.

Refined Proposition 4.1: Let $D_{x_1}$ be the normal distribution $N(0_d, \sigma^2I_d)$ with $\sigma^2 > 0$, then the AR process can not be recovered by the trained transformer even with long context. Formally, when the training sequence length $T_{tr}$ is large, for any fixed test context length $T_{te}$ and dimension $j\in [d]$, the prediction satisfies $E_ {x_ 1,W}[\frac{(\widehat{y}_ {T_ {te}})_ j}{(Wx_ {T_ {te}})_ j}] \to \frac{1}{5}$. Therefore, the prediction will not converge to the true next token.

Proof skeleton: By directly calculating, we have $(Wx_ {T_ {te}})_ j = (W^{T_ {te}}x_ 1)_ j = \lambda_ j^{T_ {te}} x_ {1j}$, and $(\hat{y}_ {T_ {te}})_ j = \frac{\tilde{a}\tilde{b}}{T_ {te}-1} \sum_ {i=1}^{T_ {te}-1} \sum_ {k=1}^d \lambda_ j^i \lambda_ k^{T_ {te}-i} x_ {1j}x_ {1k}^2$. Therefore, we have

Since $\tilde{a}\tilde{b} < \frac{1}{5\sigma^2}$ and converges to $\frac{1}{5\sigma^2}$ when $T_{tr}$ is large, the proof is completed.

Answers to Q3

Q3.1: Yes, though for the new proof, the $T_{te}$ does not need to be large.

Q3.2: Yes, $\widetilde{a}\widetilde{b} < \frac{1}{5\sigma^2}$ and thus we can prove the prediction will not converge to the true next token. However, we consider the long-context scenarios ($T_{tr}$ is large) to improve the readability and make the experimental verification more convenient. Furthermore, we conduct additional experiments in small-context scenarios and the representative results can be found in the latest uploaded PDF.

Q3.3: This is a very helpful comment that inspires us to refine the proposition. We note that $E_ {x_ 1,W}[\widehat{y}_ {T_ {te}} - Wx_ {T_ {te}}] = \frac{1}{5}*0 - 0 = 0$ in this case, which fails to interpret the phenomenon. Thus, we choose to calculate the expectation of the ratio. The experiments in the main paper are also conducted to estimate the ratio with 10$k$ test examples (randomly sampled $x_1, W$), which verify the refined theoretical results here.

Q3.4: Yes, the number of layers influences the effect of the mesa-optimizer largely. However, the main message from our paper is that given a fixed number of layers, the obtained mesa-optimizer has limited capability, and we investigate this rigorously. For multi-layer attention, [16] has empirically found it performs multi-step (accelerated) GD to solve the problem well. Theory for this case is a meaningful future work.

Weakness3 & Q4: typos

Thank you for the careful reading! We will modify these typos in the final version. Here, we only clarify that the summation in Eq. (2) is applied over 2 to $T-1$, which is also adopted in recent AR theory [17]. The reasons are as follows.

1. Given only $x_1$, we do not have any information to predict $x_2$.
2. Since we set $\rho_t = t-1$ according to the convention in ICL theory, $\rho_1 = 0$ will lead to the exploration of the transformer's output.

We will discuss these in more detail in the final version.

Q1: distributions that satisfy the Assumption 3.1

Yes, thank you for the careful reading. We will modify the claim to "any random vectors $x_1$ whose coordinates are i.i.d. random variables with zero mean and finite moments satisfy this assumption". This class still includes many common distributions like the normal distribution, thus it is still meaningful.

Q2: convergence of (a,b)

Different initialization $(a_0, b_0)$ always lead to different limiting points $(\tilde{a},\tilde{b})$. For example, in the experiments with $x_1$ from $N(0_d, I_d)$, we approximately have

1. $(a_0,b_0) = (0.1,0.1) \to (\tilde{a},\tilde{b}) = (0.44,0.44)$
2. $(a_0,b_0) = (0.5,1.5) \to (\tilde{a},\tilde{b}) = (1.41, 0.14)$
3. $(a_0,b_0) = (1.5,0.5) \to (\tilde{a},\tilde{b}) = (0.14, 1.41)$

We will definitely discuss these more clearly in the final version.

We are glad that we have addressed most of the reviewer's concerns. For Q1, we will follow the reviewer's suggestion and further modify the claim to "any random vectors whose coordinates are i.i.d. random variables with zero mean and finite moments of order 2, 4, and 6 satisfy this assumption".

We thank the reviewer once again for the valuable suggestions and insightful comments, which improve the quality of this paper.

Response to Reviewer ykMD

We thank Reviewer ykMD for the positive score and valuable comments, which can definitely improve the quality of this paper.

Weakness 1: further highlight the technical contribution

Thanks for the nice suggestion. We consider the novel AR pertaining which makes our setting more complex. First, compared to ICL for regression, our data model breaks the independence between data at different times, which causes difficulty in decomposing and estimating the gradient terms. Second, we additionally modify the embedding strategy (more dimensions), scale the attention model (much more parameters), and change the loss function (more terms) to perform the full AR pertaining. All these parts are not well studied in the literature and make the gradients more complicated.

We will discuss this in detail in the final version.

Question 1: normalization factor $\rho_t$

Yes, $\rho_t = t-1$ as we defined in line 134. This setting follows the convention in ICL theory for regression problems. We need this length-aware normalization factor to balance the outputs along different times because each element in $E_tE_t^*$ is an inner product of two vectors of size $t$, which scales as the $t$ increases. In ICL theory for regression problems, we only need to predict the last token, where the width of the embedding matrix $E$ is $T+1$ and the $\rho$ is defined as $T$. In this paper, the width of $E_t$ is $t$ thus we define $\rho_t = t-1$.

In practice, this length-aware normalization factor is a non-standard design mainly due to the existence of the softmax operator, which naturally normalizes the outputs. The theory for softmax attention in the AR setting is a meaningful direction.

We will definitely discuss this in detail in the final version.