Paper: Trainable Transformer in Transformer (ICLR-2024)

Reviewer comment (Theory):
The theory part, although the proof seems correct, it is better to have a theorem to include all results of approximating each part of a forward/backward pass of fine-tuning a Transformer. For example, this should look like 'there is a transformer with XXX layers and XXX heads, and XXX dimensions such that, when you have an auxiliary model with XXX layers/heads/dimension and a sequence of tokens, it can approximation the objective function (something you want to approximate in the forward pass of the simulator) with an error of at most epsilon'. Adding a result like this in the main result section (section 2) after the basic setup is considered essential.

Author response:
**A clear theoretical statement is missing. Something along the lines “there is a transformer with XXX layers and XXX heads, and XXX dimensions that can simulate GD on a small auxiliary transformer with error epsilon”** 

The difficulty with such a statement is that we have introduced a lot of approximations to be able to simulate gradient descent. These errors can compound in unpredictable ways that affect performance, so it is difficult to quantify the entire approximation error. We also note that people who use TINT may decide to remove an approximation and suffer the extra parameter cost to ensure a more faithful simulation of gradient descent. So, instead of providing a formal error bound, we instead use the experiments to validate that the approximations don’t hurt performance.

Reviewer comment (Theory):
In the theorem suggested above, the most important issue is determining how many parameters are needed in the simulator Transformers to approximate the in-context learning fine-tuning objective of the auxiliary Transformer (with XX parameters). The proportion of the number of parameters (parameter size of big TF / parameter size of the auxiliary one) matters since the claim is that the construction is more parameter-efficient than those in previous works [1,2].

Author response:
**“The size of the simulator for different auxiliary models isn’t clear.”**

In Table 3 in the appendix, we provided the sizes of different TINT architectures to simulate OPT models at different scales. We agree that this is an important figure to include in the main paper and we will move it in the subsequent version.

Reviewer comment (Experiments):
Based on the proportion of parameter size, one can determine in the experiments, in order to train an OPT model with 125M parameters in a single forward pass of the simulator, what size model is needed to be used as a simulator.

Author response:
**“The size of the simulator for different auxiliary models isn’t clear.”**

In Table 3 in the appendix, we provided the sizes of different TINT architectures to simulate OPT models at different scales. We agree that this is an important figure to include in the main paper and we will move it in the subsequent version.

Reviewer comment (Presentation):
The 'key component' section is especially hard to follow, and it might be better to illustrate the structure of simulators in more detail in this subsection instead of providing formal definitions (like 2.3, 2.6, 2.7) in the main text, as these are well-known definitions.

Author response:
**Presentation**
In the subsequent version, we will include some of the suggestions to improve the presentation of the paper. We will add more details to the simulator design in Figure 1. We will bring some of the details from the appendix to give more insights into the design of simulator modules for a linear layer (please refer to section B in the appendix which currently outlines the modules).

Reviewer comment (Experiments):
There is a question regarding the motivation: Why do you want to implement/approximate the forward/backward pass of Transformers in-context? The understanding is that this aims to achieve 'in-context fine-tuning' for Transformers, which avoids actual parameter updates. However, the paper does not assess which method is better and more efficient — in-context fine-tuning or direct fine-tuning.

Author response:
**"Why do we want to approximate the forward/backward pass of Transformers in context?" “Unclear whether it is more computationally efficient to do in-context fine-tuning than direct fine-tuning. Comparisons between the two methods aren’t present.”**

Our message behind TINT is two-fold:

(a) The existence of TINT indicates that existing large-scale transformers may encode intricate sub-routines to solve complex tasks in-context.

(b) We provide a framework that can explicitly encode a desirable implicit bias (i.e., the transformer should implicitly tune another smaller model) in the architecture. In particular, our framework shows that we can warm-start a large transformer model from a small transformer. The experiments further show that the large transformer improves over the small transformer at initialization.

Reviewer comment (Experiments):
Another motivational question is why Transformers are used as the simulators. Are there obvious reasons making Transformers better than other neural network architectures in this context? Since the task is to compute the forward/backward pass of the auxiliary Transformer in a single forward pass of the simulator model — a very complex task — why not consider other types of models as the outer simulator? Although Transformers or attention models have shown better in-context learnability than CNN/RNN in some tasks, those tasks differ significantly from the current one, so it is uncertain whether that reason is sufficient.

Author response:
**Why do we use transformers as simulators and not any other neural network like CNN/RNN?**

Transformers have demonstrated unprecedented success in solving language-related tasks with little to no additional gradient updates. TINT particularly studies the ability of the model to adapt to new inputs during inference (e.g., in-context learning). Such phenomena have not been commonly observed in CNNs or RNNs, so encoding such a tuning procedure in those architectures is not of as much interest. Ideas from TINT can inform the construction of CNN or RNN simulators.

Reviewer comment (Writing):
The writing of this paper is not completely clear, which makes it hard to understand the exact architecture of TinT.

Author response:
Please see Figure 1 and Section 2.1 for a straightforward illustration of the architecture of TINT. TINT is a transformer architecture that can simulate GD on an auxiliary transformer (e.g., GPT-2) during inference. The first few layers simulate the forward pass, the next few perform backpropagation while updating the parameters, and the final layers perform a forward pass using the updated parameters. The TINT architecture consists of 12 core operations that perform the forward, backward, and update operations on the 4 basic operations in a transformer: linear layer, self-attention, and layer normalization. The appendix contains all of the details and illustrations of how to simulate these 12 operations using standard transformer operations.

Reviewer comment (Presentation):
In Figure 1, it is not clear what are the dimensions of V_k, e_i, \partial y_j. Are there multiple back-propagation and gradient update steps in the TinT forward pass?

Author response:
**In figure 1, what are the dimensions of v_k, e_i, and \partial y_j?**

The size of the TINT models for different auxiliary models was in Table 3 in the appendix. We will move this table to the main paper in the final version.

**Are there multiple backpropagations through the auxiliary model?**

We can perform multiple gradient computations and updates for the auxiliary model. The size of the TINT model increases proportionally to the product of the number of auxiliary layers that we backpropagate through and the number of gradient steps.

Reviewer comment (Writing):
What is the meaning of notation $\partial$? There are many of them in Sections 2.5 and 2.6, but I don't understand the equations that contain this symbol.

Author response:
**Notations: Meaning of $\partial$ unclear.**

$\partial$ of a variable refers to the gradient of the loss function with respect to the variable (e.g., Definition 2.7). We will clarify this in the next version.

Reviewer comment (Reproducibility):
What are the parameters of the auxiliary transformer, and what are the trainable parameters of the TinT?

Author response:
**What are the parameters of the auxiliary model and what are the trainable parameters of TINT?**

The parameters of the auxiliary model are stored in the embedding layer of TINT. They are fed in as prefix tokens when necessary to TINT modules. The trainable parameters of TINT include all the parameters in the TINT modules and the embeddings (which include the auxiliary model parameters as well).

Reviewer comment (Reproducibility):
I cannot see from the description of Section 2.7 PARAMETER SHARING IN THE TINT how the parameters of TinT are shared.

Author response:
**How are the parameters shared in TINT?**

Let’s take the example of simulating forward propagation through an MLP layer in the auxiliary model (Appendix G). An MLP layer’s first layer dimension is 4 times the input dimension. However, it can be split into 4 sub-MLPs, whose first layer dimension equals the input dimension. Each sub-MLP is composed of 2 linear layers, separated by an activation layer. To simulate a sub-MLP, we can stack 3 TINT modules that simulate forward propagation through linear and activation layers. However, to simulate all 4 sub-MLPs, we use the same stack, with the contents of prefix tokens containing parameters of the 4 sub-MLPs in turns.

Reviewer comment (Reproducibility):
Could the authors write down more concretely what is the architecture of the TinT?

Author response:
Please see Figure 1 and Section 2.1 for a straightforward illustration of the architecture of TINT. TINT is a transformer architecture that can simulate GD on an auxiliary transformer (e.g., GPT-2) during inference. The first few layers simulate the forward pass, the next few perform backpropagation while updating the parameters, and the final layers perform a forward pass using the updated parameters. The TINT architecture consists of 12 core operations that perform the forward, backward, and update operations on the 4 basic operations in a transformer: linear layer, self-attention, and layer normalization. The appendix contains all of the details and illustrations of how to simulate these 12 operations using standard transformer operations.

Reviewer comment (Evaluation):
Despite its parameter efficiency, the TINT model's reliance on prefix embeddings to access auxiliary weights may lead to longer input sequence, which could potentially reduce computational and memory efficiency during inference due to the transformer's sequence length dependency.

Author response:
**"Prefix embeddings weren’t present in previous constructions, where the model learned gradient descent implicitly.”, “Computational overhead with sequence length."**

We employ prefix embeddings to represent relevant auxiliary parameters for access in each layer (as defined in 2.1). This aligns closely with the concept of prefix-tuning in large language models [2]. This approach has been utilized in prior studies like [1] to construct transformers to simulate gradient descent on 2-layer MLPs. We minimize the increase in sequence length by stacking (section 2.4).


An alternative solution was distributing auxiliary model parameters into TINT's token embeddings, eliminating the need for explicit prefix embeddings. However, this would have necessitated an increase in embedding dimension, resulting in a quadratic increase in parameters in TINT. So even though this solution could hint towards implicit gradient descent in large LLMs, it will be computationally expensive to implement and simulate this model. Exploring potential constructions that bypass the need for explicit prefix embeddings could be an interesting avenue for future research.

 

1:  Looped transformers as programmable computers. Giannou et al.'23

2: Prefix-Tuning: Optimizing Continuous Prompts for Generation. Li et al.'21

Reviewer comment (Experiments):
The capacity of the TINT structure in mimicking complex models like pre-trained transformers is well-documented, but is it equally adaptable to simpler auxiliary models, such as 2-layer MLPs?

Author response:
**Is TINT adaptable to simpler auxiliary models like 2-layer MLPs?**

Yes, the framework of TINT can be used to build a transformer to simulate any auxiliary model ranging from simple models like 2-layer MLP to the full transformer architecture. We create TINT modules for simulating forward, backward, and descent operations for 4 basic operations, linear, self-attention, layer norm, and activation. These can be composed according to the given auxiliary model.

Reviewer comment (Experiments):
Is the TINT's structure inherently dependent on the configuration of the auxiliary models it simulates, for instance in terms of number of layers, activation functions, or attention mechanisms?

Author response:
**Is TINT’s structure dependent on the configuration of auxiliary model?**

Yes, TINT’s structure depends on the configuration of the auxiliary models. We first build TINT modules to simulate the forward, backward, and descent operations on basic operations in the auxiliary model. We then stack these TINT modules, according to the configuration of the auxiliary model. The width, depth, and other dimensions of the auxiliary model will factor into the size of TINT (please refer to Table 3 in the appendix).

Reviewer comment (Presentation):
The details of the Backward Module depicted in Figure 1 are unclear. Can you provide more insight into it, particularly regarding the connections between the Backward Module's inputs/outputs across layers and their relationship to the Forward Module's output?

Author response:
**Details of Backward module in Figure 1 are unclear.**

Backpropagation uses the gradients of the (L+1)th layer to compute gradients with respect to the Lth layer. TINT does the same, using the activation of the previous TINT layer (i.e., the gradient of the (L+1)th layer) to compute the gradient with respect to the Lth layer. In order to save parameters, in the TINT architecture, the weight update for the (L+1)th layer is applied in parallel with the gradient computation of the Lth layer. Figure 4 and section B in the appendix [point to section]  illustrates this procedure on a linear layer.

Reviewer comment (Writing):
Writtings could be improved in some places. For two examples, in definition 2.1, what are the "relevant" auxiliary model weights? The current definition is a bit difficult for me to interpret.

Author response:
Each TINT layer is responsible for simulating a particular operation in the auxiliary model (e.g., the forward pass through a linear layer), which requires access to the auxiliary model weights. We provide the parameters relevant to that particular operation (e.g., the weight matrix of the linear layer) as prefix tokens to the layer. As a result, each layer in TINT is given a unique set of prefix tokens. This is similar to prefix-tuning [1], which adds a new prefix to each layer of the transformer.

Reviewer comment (Writing):
In definition 2.3, are $p_t$'s referring to positional embedding? Could you explain why there aren't positional embeddings in definition 2.10.

Author response:
In def 2.3, p_t should have been $p_t^{\textrm{TINT}}$, defined as a set of one-hot positional encodings in TINT. $p_t^{\textrm{TINT}}$ has been defined above def 2.3. We slightly modify TINT's attention definition to include $p_t^{\textrm{TINT}}$  in the computation of attention scores. On the other hand, def 2.10 outlines the self-attention definition in the auxiliary transformer.

Reviewer comment (Theory):
Theorem 2.5 shows linear attention could be approximated by softmax attention. Can softmax attention also be approximated by linear attention? If not, I feel Theorem 2.5 alone does not suffice to justify the claim that "Thus, we often use linear attention in TINT". Let me know if I have misunderstood anything. In addition, is the claimed parameter saving based on linear attention or self-attention?

Author response:
Linear attention modules are used to simulate the operations of linear layers in the auxiliary model. Theorem 2.5 shows that any linear attention can be simulated to arbitrary accuracy by a softmax attention module with a few additional parameters. So, our construction uses linear attention modules throughout to minimize the number of parameters in TINT, but one could instead use softmax attention modules without much degradation.

Reviewer comment (Theory):
Definition 2.8 uses finite difference to approximate gradient. I am wondering if we can do this from end to end. That is, can we simulate a backward pass by doing finite-difference and two forward-pass? What's the disadvantage of doing so?

Author response:
Classical analyses [1, 2] show that the variance of finite difference gradient estimates grows proportionally with the number of parameters, so if we use this method to train the entire model, we will likely have too noisy gradient estimate to successfully optimize. Recent work, dubbed MeZO [3], has shown that the end-to-end zeroth order optimization approach (i.e., multiple forward passes) can succeed in fine-tuning language models. However, MeZO relies on the inclusion of a meaningful prompt, and succeeds on relatively simple tasks. It is unknown if it can perform dynamic evaluation or in-context learning.

Reviewer comment (Experiments):
This work provides experiments on language tasks, while prior works provide experiments on simulated tasks (e.g., Akyurek et al 2022 did ICL for linear regression). So the empirical results are not directly comparable with prior works.

Author response:
Our goal is different from Akyurek et al., so their experimental setting would not make sense in our work. Our goal is to exhibit that TINT can tune a pre-trained language model (e.g., GPT-2 and OPT-125m) during inference. This allows us to demonstrate the complex behaviors that moderately sized transformers could exhibit on naturally occurring data. Our setting is very different from Akyurek et al. because their simulated models are simpler (e.g., linear regression) and trained on synthetic data. Our experiments are focused on validating the correctness of the approximations in our construction, not on how the in-context learning capability emerges during pre-training.

Reviewer comment (Novelty):
I feel an important prior work [1] is missed. Specifically, [1] also did approximation theory for ICL using transformers. How would the required number of parameters in the construction in this work compare to theirs?

Author response:
The work of Bai et al. is complementary to ours, and we will add a discussion of it to our related works section. Our construction is more general than Bai et al. because it can implement backpropagation on many complex architectures. Also, while Bai et al. choose the number of layers in the construction to ensure the prediction error is below some threshold, we cannot straightforwardly do the same, since our construction involves a much more complicated learning algorithm (i.e., fine-tuning a transformer). On the other hand, TINT performs better on standard natural language tasks than the simple statistical algorithms in Bai et al. would.