Pretrained Transformers are Deep Optimizers: Provable In-Context Learning for Deep Model Training

Authors: Weimin Wu, Maojiang Su, Jerry Yao-Chieh Hu, Zhao Song, Han Liu
Primary area: foundation or frontier models, including LLMs
Keywords: foundation model, transformer, in-context learning

Abstract:
We investigate the transformer's capability for in-context learning (ICL) to simulate the training process of deep models. 
Our key contribution is providing a positive example of using a pretrained transformer to train a deep neural network by gradient descent in an implicit fashion via ICL. 
Specifically, we provide an explicit construction of a $(2N+4)L$-layer transformer capable of simulating $L$ gradient descent steps of an $N$-layer ReLU network through ICL.
We also give the theoretical guarantees for the approximation within any given error and the convergence of the ICL gradient descent.
Additionally, we extend our analysis to the more practical setting using Softmax-based transformers. 
We validate our findings on synthetic datasets for 3-layer, 4-layer, and 6-layer neural networks.
The results show that ICL performance matches that of direct training.

Decision: reject
Reviewer scores: 6, 6, 3, 6

Reviews and discussion:
<thread>
creation date:
1730484076175

thread_type:
official_review

<thread_object>
creator: reviewer_hjr1

<content>
summary:
**Summary:** In this paper, the authors show that there exists a choice of weights, under which a Transformer can implement gradient descent on $f(w,\cdot)$ over a dataset $\\{x_i,y_i\\}_{i=1...n}$. The paper focuses on the setting where $f$ is the loss of a $N$-layer neural network, and $w$ is its parameters. The authors provide the explicit construction, along with error bounds on the approximation error to a real $L$-step gradient descent.

I lean towards recommending reject for this paper. The reason is that I am unclear about the purpose of the main result. I will elaborate below.

strengths:
1. The construction is *mostly* explicit, and is clearly described in the various proofs. The proofs seem to be correct, though I have not carefully verified each line.
2. To simulate $L$ gradient descent steps of a $N$ layer neural network, one only needs a $(2N+4)L$-layer Transformer.
3. The recursive expression Lemma 1 seems to work well with a multi-layer Transformer.
4. Being able to learn a neural network in-context with a Transformer seems interesting.

weaknesses:
Transformers have been shown to have high representation power. **If one does not care about the complexity of the construction,** [1] claims that the Transformer can implement arbitrary algorithms by literally simulating a computer; Lemma 14 of [1] claims that there is a 13 layer 1-head Transformer that implements T steps of gradient descent on a 2-layer sigmoid network.

In the other extreme, [2] (and numerous follow-up works) considered **simple problems**, but also provided **simple constructions** which are provably learnable during training. 

My main criticism of the results of this paper is that the results here falls somewhere in-between the two extremes above -- the target problem of "can Transformers implement gradient descent in-context for a neural network" is challenging, and thus required a rather **complex construction**. If so, why would I not simply use the general construction in [1]? 

The constructions in this paper are likely not being learned, in theory or in experiments, so what makes it superior to the generic construction in [1]? This is further compounded by the fact that the **error bounds in this paper are large**, growing exponentially in number of layers and steps. Additionally, the dimensions of a single Transformer layer need to be large enough to accomodate the weight parameters of the entire $N$-layer neural network, **requiring a lot of memory**, and this requirement seems unavoidable. 

A motivation of this paper was to have foundation models train other deep models, but the various issues listed above make me question whether the approach in this paper will lead to meaningful results.

Below I provide further details on some of the issues I have with the paper

1. The construction is not entirely explicit, Lemmas 2,3,4 appeal to the existence of some combination of ReLU functions to approximate sufficiently smooth functions via the $(\epsilon, R, H, C)$-approximable assumption. 

3. Theorem 1 requires the activation $r(t)$ and $r'(t)$ to be Lipschitz continuous, but ReLU does not satisfy this. It would be helpful for the authors to state a few concrete examples of N-layer neural networks (along with specific activations), and state what the Lipschitz constants/dimensions/error bounds end up being. 
4. The error bounds are quite large. Corollary 1.1 has exponential error in $L$, Lemma 6 has error growing as $d^N$. Some of the error is unavoidable -- divergence in gradient desecent over non-convex functions, and exponentially growing lipschitz constant of a $N$-layer Transformer. It would be useful if the authors can clearly distinguish the error terms inherent due to gradient descent over a non-convexity of the problem, vs the error terms that are due to the construction.
5. Experiments are a little unclear. What is the exact Transformer architecture used? Are the lines for 3/4/6 layer NNs in Figure 1 obtained from training a 3/4/6 layer NN over the in-context examples? Or are they the output of the "true network parameters for $f_i$" on line 2393?
6. The Transformer performance in experiments seems quite bad for 0.7/0.3 data. Can the authors also show the plot for 0.5/0.5?
7. If possible, it would be useful to see if scaling up the Transformer, as described in the paper, will indeed enable it to implement more gradient descent steps. (a crude way is perhaps to simply plot loss against number of Transformer layers.)

questions:
1. Can the authors please compare your construction to the one in Algorithm 10 of [1]? Are there any similar ideas used?
2. If you are already assuming that various parts of the neural network are approximable by relu, and if you are already incurring exponentially-growing error terms, why do you not just treat the entire neural network as a $(\epsilon, R, H, C)$-approximable function, and simply use the Transformer to learn relu coefficients? Will this lead to worse bounds?
3. The dimension of the Transformer is very large, as each layer needs to contain all the parameters of the N-layer neural network. Can the authors please comment on what the exact dimension of the Transformer needed? I think it is useful to state this value in Theorem 1.

**minor suggestions that did not affect score**
1. $L$ is used for both loss, number of gradient descent steps, and Lipschitz constant. This gets confusing at times, and I suggest for the authors to use different letters.

[1] [Looped Transformers as Programmable Computers](https://arxiv.org/abs/2301.13196), Giannou et el 2023

[2] [Transformers learn in-context by gradient descent](https://arxiv.org/abs/2212.07677), von Oswald et el 2023
</content>
</thread_object>

</thread>
<thread>
creation date:
1730528644488

thread_type:
official_review

<thread_object>
creator: reviewer_r2xk

<content>
summary:
This paper studies transformer's in-context learning capability. It shows that given an $N$-layer neural networks, there exists a $(2N+4)L$-layer pre-trained transformer that simulates $L$ gradient descent steps. The analysis is also extended to softmax-based transformers. The findings are validated on synthetic datasets.

strengths:
ICL is an important phenomenon, the theoretical explanation of which is still far from complete.

weaknesses:
1. The main contribution of the paper is demonstrating the ability of transformers to simulate N-layer neural networks via gradient descent. However, this has been shown in reference [A]. The submission does not cite [A] and there is no discussion about what's novel compared to [A]. 

      [A] Wang, Jiang, Li, "In-Context Learning on Function Classes Unveiled for Transformers", ICML 2024. 

2. Question 1 asks "Is it possible to train one deep model with another pretrained foundation model?". This is not the same as ICL and the present submission does not answer this question. 

3. The organization and presentation requires significant improvement. Instead of focusing on a lot of notational details, the main paper should have more discussions about the significance and implication of the presented results, and the relationship between the presented results and existing results.

questions:
1. What's the true contribution of the paper?

2. Is there any novelty in the proof techniques? If yes, it should be highlighted. 

3. What's the connection between the experimental results and the theory? Besides, the R-squared results do not seem very good. How can the experimental results validate the theory?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730668039837

thread_type:
official_review

<thread_object>
creator: reviewer_99cj

<content>
summary:
This paper investigates the theoretical underpinnings of transformers as deep optimizers, specifically focusing on their capacity for in-context learning (ICL) to simulate gradient descent steps. The authors provide a construction of a transformer capable of performing gradient-based optimization implicitly without updating its own parameters. Through a series of lemmas and Theorem 1, the paper establishes that a pretrained transformer can approximate gradient descent steps of deep neural networks with convergence guarantees.

strengths:
The general topic of theoretical understanding of transformers and ICL is interesting and novel.

weaknesses:
1. The organization could be improved. The main result, Theorem 1, is only presented at the end of the paper. The purpose of preceding Lemmas 1-6 is unclear; it would help to clarify whether they serve solely as components of Theorem 1’s proof or have independent significance.

2. The studied problem lacks sufficient motivation. Transformers are known to have a complicated structure and may contain many parameters. Therefore it is not surprising that a transformer can have enough expressiveness to approximate the gradient descent steps.

3. There is no overview of the empirical validation in the main body of the paper. At least, the authors should use a few sentences to indicate which part of the conclusion is validated by empirical results.

4. The title says "pretrained transformers are deep optimizers", while no pretraining process is analyzed in the main body of the paper.

questions:
See Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699066229

thread_type:
official_review

<thread_object>
creator: reviewer_tzww

<content>
summary:
This paper provides construction and proof on a ReLU-based Transformer are able to perform gradient descent steps for a N-layer fully connected neural network.

strengths:
This paper presents clear and rigorous proof in showing Transformer can implement gradient descent for (x, y) ~ N-layer neural network.

weaknesses:
1. The paper frequently uses the term "pretrained Transformer." Typically, this refers to a Transformer model trained on a large corpus for next-token prediction. However, in this paper, "pretrained Transformer" actually refers to a Transformer with specific weight constructions. I could not find an explanation of a "pretraining task" that leads to these constructed weights. This terminology may be confusing, especially in the introduction, where the authors emphasize that "one trained foundation model could lead to many others without resource-intensive pretraining".

2. The paper also discusses using the Transformer to estimate the score function for diffusion, which is presented as a major contribution. However, this is accomplished by reducing the score function to a fully-connected neural network. It might be helpful to adjust the emphasis on the diffusion score function and clarify that the main contribution is to provide theoretical support for the Transformer's capability to implement gradient descent for an N-layer neural network.

3. While the idea of demonstrating that a Transformer can implement an N-layer neural network is interesting, the paper only establishes the existence of such a network without addressing optimality or uniqueness. This makes the result slightly less compelling to the community, especially as the construction serves primarily as a proof of concept of the (modified) Transformer's capability rather than a more extensive finding. I understand that this limitation arises from the nature of the problem itself, and it may only be addressed through further results that demonstrate additional insights or practical applications.

questions:
I found the necessity of the element-wise multiplication layer intriguing, particularly as discussed in Remark 3. Could the Transformer architecture circumvent the need for this element-wise multiplication by increasing its depth and width? In other words, is it possible to achieve the same functional capabilities without the element-wise multiplication layer by scaling the model’s size?

---
Updated after the discussion:
Thank you to the authors for taking the time to address my concerns. However, I still find the motivation behind training an N-layer neural network via in-context learning in another model questionable. Beyond the theoretical FLOPs and actual computational requirements, there is also the storage trade-off to consider. Unfortunately, I don’t see the benefits of replacing the direct training of an N-layer neural network with the forward pass of the construction model.

I am content with the story being about proposing a construction method to implement gradient descent for an N-layer neural network. Therefore, I will keep my rating at 6.

Reason for not a higher score: This work provides proof-of-concept results regarding the Transformer’s capacity to implement gradient descent for an N-layer neural network. However, these results lack practical insights for improving real-world applications, making this primarily a proof-of-concept contribution.

Reason for not a lower score: This work contributes to the literature on understanding the Transformer’s capacity. I appreciate the theoretical merit embedded in this study.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732427916043

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Summary

comment:
Dear Reviewers,

Thank you for your insightful feedback and the time invested in reviewing our work. We have addressed your inquiries and concerns in our rebuttal and the updated manuscript.

Following your suggestions, we have enhanced the readability of the paper through extensive proofreading, detailed polishing, and correction of typos. Specifically, we have highlighted the main results in the Introduction (`line 089–090`) and provided a more informal version in Appendix A for clarity. We have also added a more detailed description of the experimental part in the Introduction (`line 085-088`). Additionally, we have clarified the independent significance of Lemmas 1–6 to ensure their purposes are well understood. Furthermore, we have provided a precise value for the hidden dimensions of the Transformer as required in Theorem 1.

For the convenience of reviewers/future readers, this response provides a high-level overview of our contributions.

---

**Revision Details**

In response to the reviewers' suggestions, we have made several key modifications, summarized as follows:

**Major revisions** include:

1. Replace “Pretrained Transformer” with “Transformer”. [Reviewers `tzww`, `99cj`, and `R2xK`]
2. Add an overview of the empirical part in the Introduction (`line 089–090`) and provide a detailed description in Appendix G.1, which covers aspects such as the model architecture. [Reviewers `99cj`, `R2xK`, and `hJr1`]
3. Add the experiments to show that scaling up the transformer enables it to implement more ICGD steps (Appendix G.2).  [Reviewer `hJr1`]
4. Add a detailed comparison with [Wang24] (In the Introduction and Related Works). [Reviewer `R2xK`]
5. Provide a precise specification of the Transformer dimensions required in Theorem 1. [Reviewer `hJr1`]

**Minor revisions** include:

1. Proofread the manuscript and fix all identified typos and grammatical errors.
2. Improve the readability of the main text. 

**A Concise Summary/Emphasis of Our Contribution**

We appreciate the opportunity to clarify the scope of our contributions. In the initial draft, we may not have described the breadth and depth of our work as effectively as we intended. Here, we have summarized our contributions:

* **Structured and efficient transformer construction** We offer a detailed characterization of the ICL capabilities of a transformer model in approximating the gradient descent training process of an $N$-layer network. Utilizing a term-by-term approximation technique and the selection of specific terms, we demonstrate that only $(2N+4)L$ transformer layers are required (Theorem 1) to implement $L$ gradient descent steps. We also detail the approximation error in Corollary 1.1.

* **Less restrictive input and output dimensions for $N$-layer neural networks.** We extend our analysis to the case where output dimensions exceed one, as detailed in Appendix E. This includes scenarios where input and output dimensions differ.

* **More practical transformer model.** We extend our analysis to include Softmax-based transformers, reflecting more realistic applications as detailed in Appendix F.

* **Experimental validation.** We validate our theory with ReLU- and Softmax-based transformers, specifically, ICGD for the $N$-layer networks. The numerical results in Appendix G show that the performance of ICL matches that of training $N$-layer networks.

* **Advanced and complicated applications.** We explore the application of our results to approximate the score function in diffusion models, as discussed in Appendix H.

[Wang24] In-context Learning on Function Classes Unveiled for Transformers, ICML 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733111986812

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Clarifications on Related Work and Contributions

comment:
Dear Reviewers and Future Readers,

We appreciate the additional reference [Wang24] suggested by [Reviewer `R2xK`](https://openreview.net/forum?id=ZIFkrT1GwM&noteId=idJUnb4cWr). It is relevant and was unfortunately overlooked in our initial literature review. Below, we provide clarifications on its relation to our work and address potential misunderstandings about our contributions.

[Wang24] In-context Learning on Function Classes Unveiled for Transformers, ICML 2024.

---

### **Related but Seemingly Irrelevant Work [Wang24]**

* While [Wang24] is relevant, **the shared topics (ICGD of N-layer NN) and results are not mentioned in both its title and abstract.** This makes it appear unrelated at first glance. 
* Moreover, **[Wang24] was not readily visible during the preparation of this work.** It's very recent. It is not available on arXiv, and the ICML 2024 paper only became accessible around July. At the time we submitted our work, this paper had no citations.

**We believe it is very likely that most authors in this ICLR submission cycle may also overlook [Wang24], just as we did.**

Despite this, we apologize to the authors of [Wang24] for not attributing proper credit in our initial draft. We acknowledge its relevance and appreciate the opportunity to discuss its connection to our work. However, we believe that **overlooking this particular reference does not diminish the originality or merits of our contributions**, as explained below.

---

### **Clarifications on Our Contributions**
In the most concise manner, our main contributions are as follows:
1. Explicit construction for ICGD of N-layer NN.  
2. Error bounds.  
3. Results for both ReLU and Softmax transformers.  

We have discussed the differences between our work and [Wang24] extensively in both our rebuttal and revisions. Specifically:  
- Contributions (2) and (3) are solid and novel, as acknowledged by the reviewers.  
- For contribution (1), **while there may be superficial similarities with [Wang24] (both showing ICGD of N-layer NN), our approach uses entirely *different techniques*, operates under *less restrictive assumptions*, and achieves *stronger results*.**  Please see the comparison below for details.

Even if [Wang24] is considered prior work in this area, **our contributions remain complementary and significant.**

---

### **Key Differences**
For clarity, we summarize the key distinctions between our work and [Wang24] below, and direct the readers to our rebuttal and latest revision for further details:

- **More structured and efficient transformer architecture.**  
  [Wang24] uses an **$O(N^2L)$**-layer transformer to approximate $L$ gradient descent steps on $N$-layer neural networks. In contrast, our approach adopts entirely different construction strategies, such as inductively approximating intermediate chain-rule terms, and requires only **$(2N + 4)L$** layers.

- **Less restrictive input and output dimensions for $N$-layer neural networks.**  
  Unlike [Wang24], which limits the output of $N$-layer networks to a scalar, our results hold for outputs of any dimension.

- **More practical transformer model.**  
  We extend our analysis to include Softmax transformers, reflecting more realistic applications, whereas [Wang24] does not consider this.

- **More advanced and complex applications.**  
  [Wang24] discusses applications to indicators, linear, and smooth functions. In contrast, we explore more advanced cases, such as the score function in diffusion models.

We understand that, despite our best efforts, it is possible that we may have overlooked or misunderstood certain parts of [Wang24] given the limited time. However, we assure that we have no intention of disregarding or discrediting [Wang24] in any way. The contributions of [Wang24] are solid, interesting, and important. The above discussions are solely for clarification, as the two works have nuanced and different focuses. We are open to any further discussion.

---

We hope these clarifications help readers understand the unique contributions of our work and its relation to [Wang24]. While both works explore related topics, we believe our approach offers distinct theoretical and practical advancements.

We thank all reviewers for their time, constructive feedback, and engagement. Your involvement has helped us improve this work.

Lastly, we thank the authors of [Wang24] for their good work.

Best,  

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores the transformer's ability for in-context learning (ICL) to simulate the training process of deep models. The key contribution is demonstrating how a pretrained transformer can train a deep neural network via ICL in an implicit manner. Specifically, the authors construct a(2N+4)L-layer transformer that simulates L gradient descent steps of an N-layer ReLU network. They provide theoretical guarantees for approximation error and ICL convergence. The analysis is extended to Softmax-based transformers in practical settings. Experiments on synthetic datasets for 3-layer, 4-layer, and 6-layer neural networks show that ICL performance matches direct training. The primary concern raised by the reviewer is that some of the contributions in this submission are considered insignificant in light of a recent paper [Wang24] published at ICML 2024. While other contributions, such as the results for softmax transformers, are deemed novel, their presentation is largely deferred to the appendix. I acknowledge that this paper may contain important contributions. However, it requires significant and careful revisions to meet the standard for acceptance. Therefore, I recommend rejection.

[Wang24] In-context Learning on Function Classes Unveiled for Transformers, ICML 2024.