On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability

Autoregressively trained transformers have brought a profound revolution to the world, especially with their in-context learning (ICL) ability to address downstream tasks. Recently, several studies suggest that transformers learn a mesa-optimizer during autoregressive (AR) pretraining to implement ICL. Namely, the forward pass of the trained transformer is equivalent to optimizing an inner objective function in-context. However, whether the practical non-convex training dynamics will converge to the ideal mesa-optimizer is still unclear. Towards filling this gap, we investigate the non-convex dynamics of a one-layer linear causal self-attention model autoregressively trained by gradient flow, where the sequences are generated by an AR process $x_{t+1} = W x_t$. First, under a certain condition of data distribution, we prove that an autoregressively trained transformer learns $W$ by implementing one step of gradient descent to minimize an ordinary least squares (OLS) problem in-context. It then applies the learned $\widehat{W}$ for next-token prediction, thereby verifying the mesa-optimization hypothesis. Next, under the same data conditions, we explore the capability limitations of the obtained mesa-optimizer. We show that a stronger assumption related to the moments of data is the sufficient and necessary condition that the learned mesa-optimizer recovers the distribution. Besides, we conduct exploratory analyses beyond the first data condition and prove that generally, the trained transformer will not perform vanilla gradient descent for the OLS problem. Finally, our simulation results verify the theoretical results.

Paper

References (63)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer c1Lc5/10 · confidence 4/52024-06-19

Summary

This paper studies the behavior of a single unit of linear self-attention trained via gradient flow on an autoregressive (AR) loss on token sequences generated by a particular class of noiseless (deterministic) linear systems. The first theoretical results state that under some assumptions, the result of gradient flow from a particular class of initializations is a model whose in-context learning (ICL) predictions can be interpreted as those of another model that has been updated with one step of gradient descent (GD) using the context examples. The second part of the theoretical contribution is to say that a stronger assumption is necessary and sufficient for the result of this GD to give accurate predictions. These results are supported by numerical simulations of the studied setting.

Strengths

- This paper aims to address an important issue in the ICL theory literature: most works study the emergence of ICL in transformers trained on an ICL objective (referred to as few-shot pertaining in the paper), whereas in practice transformers are trained on an autoregressive objective. - The theoretical results are correct as far as I can tell from skimming the proofs. The notations are clearly defined and consistent. - The proof sketch is well-written. - The experiments nicely verify the theoretical results. Sufficient details are provided.

Weaknesses

1. The results do not suggest that “ICL by AR pretraining is more difficult than ICL by few-shot pretraining”, as the authors claim. Instead, they suggest that ICL to predict the $T+1$-th token given the sequence $\mathbf{W}^0\mathbf{x}\_1, \dots, \mathbf{W}^{T-1}\mathbf{x}\_1$ is harder than ICL to predict $\mathbf{w}^\top \mathbf{x}_{T+1}$ given $\{(\mathbf{x}\_i, \mathbf{w}^\top \mathbf{x}\_i)\}\_{I=1}^T$ and $\mathbf{x}\_{T+1}$. These are two very different types of tasks and it is not fair to make conclusions about pretraining strategies based on their differing performance on these different tasks. Rather, one needs to compare the performance of a transformer trained in an autoregressive fashion and one trained on a few-shot learning objective under the same data/task model. In the data model in this, few-shot pertaining would use the same loss as the AR loss (equation 2) except that only the $t=T-1$ terms would remain in the loss. It is not clear whether this objective leads to better or worse solutions, in fact, this objective contains the same amount of information and therefore intuition suggests that it leads to solutions of the same quality. This is a key point as arguably the main message of the paper is that AR pretraining behaves differently than few-shot pertaining and thus the ICL community should focus on AR pretraining. I would suggest that the authors re-frame their discussion around the inadequacy of the regression data model considered by most of the existing ICL literature, as changing the data model is really the key contribution of this paper. 2. The setting is very artificial such that the significance of the results seems dubious. The model consists only of one, linear attention unit, and the data is assumed to follow a deterministic linear system, i.e. $\mathbf{x}\_{i+1}=\mathbf{Wx}\_i$ where $\mathbf{W}$ is diagonal with diagonal entries having magnitude 1. The learning model and data model both being linear means that the ICL prediction of any learning model with parameters $\mathbf{W}^{KQ}\_{32}=a\mathbf{I}_d$ and $\mathbf{W}^V\_{12} = b\mathbf{I}\_d$ for $ab>0$ and all other parameters being zero can be interpreted as the prediction of a linear model $\hat{\mathbf{W}}$ trained directly to approximate $\mathbf{W}$ with one step of GD starting from $\mathbf{0}\_{d\times d}$. It is not clear at all how such a conclusion would follow if the softmax activation, shown to be crucial in practice, is reinstated, or if other components of the transformer (MLP layer, multiple attention heads, multiple attention layers, etc.) are added. Also, the token embedding is strange and redundant, with consecutive inputs $\mathbf{x}\_i, \mathbf{x}\_{i-1}$ concatenated in the same token. Another, though less significant, issue is that only the population loss is considered. The authors repeatedly claim that “We believe that our analysis can give insight to the establishment of theory in settings with more complex models” but never explain why. 3. The model is already initialized as a model that implements gradient descent during ICL, and under the scalar-times-identity covariance assumption (4.1), or the enforced off-diagonal masking, it is not hard to show that the gradient flow dynamics reduce to the dynamics of two scalars ($a$ and $b$) which parameterize a set of models that always does GD during ICL. It is not clear whether the trained transformer can learn to perform GD during ICL if it is not initialized as a model that does this, and the experiments do not address this question. 4. “Gradient clipping” is used to refer to the masking procedure that zeros out all off-diagonal elements of the gradient. This undersells the severity of this procedure — clipping is a non-zero-thresholding operation applied equally to all gradient elements and is commonly used in practice for optimization and privacy purposes, and does not use any prior knowledge of how the gradient should be structured. The masking procedure described here sends a particular subset of elements of the gradient to zero, to drive the dynamics towards a particular desired solution based on prior knowledge of the data-generating process, which is a much stronger, and impractical, procedure. 5. Assumption 4.2 and Theorem 4.2 are tautological and not insightful. They should identify the properties of $\mathbf{x}\_1$ and $\mathbf{W}$ that lead to such behavior. 6. $\kappa\_3$ has $\Omega(d)$ growth, which means the effective step size of the gradient descent step is $\tilde{O}(T/d)$, meaning $T$ (number of examples in the context length) must be $\tilde{\Omega}(d)$ for accurate predictions, which makes sense for this problem setting, but suggests the proposed setting is far from capturing ICL ability in practice, where the number of in-context examples is far fewer than the token embedding dimension. In the experiments, $T=20d$, which means the experiments do not evaluate whether the results can extend to smaller-context scenarios. 7. Proposition 4.2 is not meaningful because it does not rule out whether a different initialization could lead to a model that performs GD. 8. Even if we accept all of the simplifications and assumptions (including very long context) made to obtain a trained model that performs GD during ICL, this trained model does not learn the correct step size for GD, so it can’t perform ICL. To me, this suggests that GD is *not* the mechanism that transformers use to perform ICL — even in such an idealized setting, GD cannot explain the ability of transformers to perform ICL. However, the messaging around the paper is not consistent with this, as it continues to push the narrative that trained transformers are mesa-optimizers.

Questions

Why is the y-axis in Figure 1-b defined as the ratio of two vectors?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors discuss limitations.

Authorsrebuttal2024-08-11

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.

Reviewer WaqB7/10 · confidence 4/52024-07-10

Summary

This paper studies the emergence of mesa-optimization/in-context learning capabilities of Transformers. In particular, it attempts to fill the gap of our understanding of training dynamics and specifically on non-convex dynamics where the sequences are autoregressively generated as $\vec{x}\_{t+1} = W \vec{x}\_{t}$. There are many contributions of this paper. First, it shows, under Assumption 4.1, transformer converges to the gradient-descent conjecture for ICL. Under the same assumption, the authors show moments of the data is necessary & sufficient for ICL. Finally, this papers study data distributions beyond Assumption 4.1, and shows in this case, transformers are *not* doing gradient-descent for ICL.

Strengths

Presentation is very clear and easy to follow. Contribution is significant. It bridges the gap between two main stream hypothesis of how Transformers perform mesa-optimization/ICL: (1) Transformer is doing gradient descent [1]; (2) Transformers is doing something beyond vanilla gradient descent, such as preconditioned GD [2] or higher-order optimization methods [3,4]. The authors concluded the divergence between these two hypotheses are due to the data assumption (Assumption 4.1). [1] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In ICML, volume 202, pages 35151–35174, 2023. [2] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In NeurIPS, 2023. [3] Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models, 2023. [4] Giannou, Angeliki, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D. Lee. How Well Can Transformers Emulate In-context Newton's Method?. arXiv preprint arXiv:2403.03183 (2024).

Weaknesses

When contrasting with vanilla gradient descent conjecture, the authors only mentions the alternative *preconditioned GD* hypothesis. However, there are many other propositions, such as higher-order optimization proposition, such as a *Newton's method* [3,4,5]. I would recommend the authors have a deeper discussion there. For example, in [5], they argued the GD++ method in [1] is actually a higher-order optimization method with superlinear convergence rate as discussed in [3,4]. [5] Vladymyrov, Max, Johannes von Oswald, Mark Sandler and Rong Ge. Linear Transformers are Versatile In-Context Learners. ArXiv abs/2402.14180 (2024)

Questions

Theories are built upon a one-layer linear causal self-attention model, what happens if we extend to (1) more layers or (2) non-linear activations?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

Theories are only built upon a one-layer linear causal self-attention model, where there's is still a big gap compared to what's used in practice -- many-layer non-linear activation self-attention and feed-forward model with skip connections. But for theory, this is good enough.

Authorsrebuttal2024-08-13

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

Reviewer dsNs5/10 · confidence 4/52024-07-12

Summary

- The authors study the problem of in-context learning an autoregressive (AR) process (defined with a uniformly drawn diagonal unitary transformation) with a one-layer linear causal self-attention model, trained by gradient flow on square loss. - Under a specific parameter initialization scheme and a distributional assumption on the initial token, they show that a trained model can implement one step of gradient descent (=mesa-optimizer) applied to a least-square problem estimating the transformation matrix of the AR process. - Unfortunately, they show that the obtained mesa-optimizer can fail to learn the AR process with next-token prediction, even with infinitely many training samples and infinitely long test context length. - They further propose a stronger assumption being a necessary and sufficient condition for the success of mesa-optimizer. They also provide an example of the distribution of the initial token satisfying the necessary and sufficient condition. - They show that, without the distributional assumption, the trained model may not implement a gradient descent step. - Lastly, they verify their theory with simulations.

Strengths

- The paper is well-written overall. - Most of the theoretical analyses seem correct and rigorous. The proof sketch in Section 5 provides a clear and plausible picture of the proof. - The analyses do not require/assume the diagonal structure of submatrices of the parameters. Instead, the authors found a sufficient condition (a distributional assumption on the initial token) so that they can prove the diagonal structure.

Weaknesses

- W1. The analyses and experiments cannot provide insight into the general cases of autoregressive training of Transformers. - This is mainly because the theoretical analyses greatly depend on the particular initialization scheme (Assumption 3.1). The initializations of weight matrices are too sparse: at initialization (and throughout training), they only focus on the relationship between very particular token embeddings. This seems very unnatural, thinking of the training transformers in practice. Quite ironically, the paper’s result corroborates that the initialization scheme is inappropriate for solving in-context learning with next-token prediction, under most of the plausible data distribution. - This assumption might be inevitable for a tractable analysis. However, the experiments are also confined to the framework under Assumption 3.1. I think the authors should have demonstrated many more simulations with various initialization schemes and even with diverse architectures (linear transformers with non-combined weight matrices, non-linear transformers…) to support their arguments (”the data distribution matters for the emergence of mesa-optimization”, “there is a capability limitation of the mesa-optimizer”…). - W2. There are some weaknesses in theoretical analyses. - Even though the authors found a condition for the emergence of mesa-optimization in a very particular problem setting, the analysis does not go beyond the “diagonal structure” of weight matrices. Thus, for me, although the contribution is novel, it does not look significant. - For the comments and questions on Proposition 4.1, see Q3. - W3. Minor comments on typos & mistakes. - Line 130: $W^{KQ} = {W^K}^\ast W^Q$ (rather than transpose) - Equation (2): I guess the sum should be applied over $t=1, …, T-1$. - Line 182: “$x_1 \in \mathbb{R}^d$” - Lines 207, 289: ‘$W$’ must be bold. - Figures 1b, 1d, and Line 330: “$T_{te} -1$”

Questions

- Q1. In lines 186-187, you mentioned that any random vectors with i.i.d. sampled coordinates satisfy Assumption 4.1. Is it true? I think they only satisfy the first part of the assumption (before “In addition”). As a counterexample of the second part, if the distribution of each coordinate is heavy-tailed (e.g., infinite second moment), $\kappa_1$ and $\kappa_2$ are infinite. - Q2. In Theorem 4.1, does different initialization $(a_0, b_0)$ “may” lead to different $(\tilde{a}, \tilde{b})$, or “always” lead to different point? - Q3. Questions on Proposition 4.1 - Isn’t $T_{te} \rightarrow \infty$ unnecessary? This is because we can show that for any $i\in [T_{te}-1]$, $\mathbb{E}\_{x_1} [x_i x_i^\ast] = W^{i-1} \mathbb{E}_{x_1} [x_1 x_1^\top] (W^\ast)^{i-1} = \sigma^2 W^{i-1}(W^\ast)^{i-1} = \sigma^2 I_d $ because $W$ is unitary and $x_1 \in \mathcal{N}(0_d, \sigma^2 I_d)$. - Isn’t $T_{tr} \rightarrow \infty$ also unnecessary? This is because it seems enough to show that $\tilde{a} \tilde{b} < \tfrac{1}{5}$. - Most importantly, is the last statement **correct**? Shouldn’t you prove that $\mathbb{E}_{x_1} [\hat{y}_T - Wx_T]$ converges to a nonzero vector? This is because $x_T = W^{T-1} x_1$ is also a random vector and depends on $T$. - From what I read, it seems like the data distribution $\mathcal{D}_{x_1}$ is the main cause of the failure of the mesa-optimizer. In my opinion, however, this is also because the “single-layer” linear self-attention model can only simulate a “single step” of GD for the OLS problem. Then isn’t the number of layers the problem? I think there are several reasons why the mesa-optimizer fails, and the authors should mention and analyze these factors as many as possible. - Q4. I guess Theorem 4.2 requires Assumption 3.1. Or does it? If the authors address all my concerns, I will definitely raise my score.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The limitations are provided in Section 7.

Authorsrebuttal2024-08-10

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.

Reviewer ykMD5/10 · confidence 3/52024-07-15

Summary

This paper studies the autoregressive training of a one-layer linear attention model for in-context learning of first-order AR processes. It is shown that under certain distribution of the initial point, the gradient flow on the population next-token prediction loss will converge to a model that makes the prediction based on the estimate given by one step of gradient descent over the in-context loss. Furthermore, a sufficient and necessary condition is provided to characterize learning capability of the trained model. Besides, it is shown that for more general data distribution, the trained model does not perform one step of gradient descent over the in-context OLS loss. Numerical simulations are conducted to support the theory.

Strengths

The paper is well written and easy to follow. The explanation of the results is very clear, and logic flow of the organization of the paper is smooth. The current results present a solid contribution to the theoretical understanding of the in-context learning by transformers, especially for AR training.

Weaknesses

It would be helpful to further highlight the technical contribution. In particular, how does the technique and analysis differ from those existing papers on training dynamics of transformers for in-context learning of regression problems?

Questions

What's the value of the normalization factor $\rho_t$? This seems to be a non-standard design as $\rho_t$ depends on $t$. From my reading of line 634, it requires that $\rho_t=t-1$. Please clarify this and explain why this is reasonable.

Rating

5

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have adequately addressed the limitations.

Reviewer dsNs2024-08-10

Thank you for the rebuttal

Thank you for your detailed response. - The general response and the additional experiments on standard initialization have addressed my concerns (W1 & W2). I also agree that finding Assumption 4.1 is non-trivial. - Also, I am happy to see that the refined Proposition 4.1 and its proof sketch have satisfactorily addressed my concern on the theoretical contribution. - Regarding the response to Q1, I think it is important to mention which orders of moments are finite. Overall, most of my concerns are addressed. Although I am still not absolutely certain, now I think the paper is above the acceptance bar. Hence I will raise my score to 5.

Reviewer c1Lc2024-08-10

I thank the reviewers for their thorough response, which has helped to alleviate my concerns. I especially appreciate the experimental results with non-diagonal initialization, pointers to related work and clarification on Proposition 4.2. I am still not overly impressed by the results, especially as they start with diagonal initialization (Theorem 4.1 and Theorem 4.3), and use gradient masking (Theorem 4.3). Nevertheless, as the authors point out, even in these simplified settings characterizing the limiting points of gradient flow is a novel contribution. I would recommend that in the revision, the authors put more emphasis on their results contributing a closed-form characterization of the convergence points of AR pre-training, rather than showing that AR pre-training results in a mesa-optimizer, since the latter point is weakened by the assumption that the initialization is already a mesa-optimizer. As a minor point, the empirical results with T=3 and T=10 appear to be missing from the attached pdf.

Reviewer ykMD2024-08-11

Reply to rebuttal

I thank the authors for the explanation. I don't have further questions.

Authorsrebuttal2024-08-12

Thanks for the helpful comments and positive score. We are glad that we have addressed the reviewer's concerns.

Authorsrebuttal2024-08-13

Summary of the revision

# Summary of the revision We sincerely thank all reviewers for their valuable comments and positive support, which help us further improve the quality of this paper. We also thank ACs for organizing this nice review process. We have thoroughly addressed the detailed comments and summarized the revision in the final version as follows. ## New experimental results 1. We add complete experimental results under standard initialization with different scales ($\sigma_w = 0.001/0.01/0.1$). Representatives of them have been presented in the latest uploaded PDF (Fig. a-f). 2. We add complete experimental results in small-context scenarios ($T = 3/5/10,d=5$). Representative results have been presented in the latest uploaded PDF (Fig. g-i). ## Refined theoretical result We will refine the theory and proof of Proposition 4.1 as we presented in the answer to Reviewer dsNs. ## Writing 1. We fix the typos pointed out by Reviewer dsNs. 2. We clarify that Proposition 4.2 does not depend on diagonal initialization 3. We change "gradient clipping" to "gradient masking" in Section 4.3 to avoid confusion. 4. We clarify the ratio calculation between prediction and true vector in Fig. 1b. ## Discussion ### Contribution & Impact 1. We add more discussion about the key contribution of this paper in Section 4. 2. We add more discussion about our technical improvement in Sections 4 and 5, compared with existing studies on ICL. 3. We add more discussion about the contribution and the necessity of finding Assumption 4.1 to keep the diagonal structure, compared with existing studies. We also discuss its limitations. 4. We add more discussion about the negative results of one-layer linear attention, which can be solved when considering more complex transformers in the future. 5. We add more discussion about the impact of our work on more complex autoregressive transformers (e.g., multi-layer attention, non-linear activations). ### Setting & Assumptions 1. We add more discussion about the rationality of our setting (model, data distribution, token construction, and normalization factor). 2. We add more discussion about the necessity and limitation of the diagonal initialization near Assumption 3.1 and Theorem 4.1. 3. We discuss how the experimental results under standard initialization extend our theory and support the rationality of the diagonal initialization in Section 6. 4. We add more discussion about the properties of Assumption 4.2. ### Related work 1. We cite and discuss the additional related papers mentioned by Reviewer WaqB.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC