Trainable Transformer in Transformer

We thank the reviewer for appreciating the idea behind TINT and the effort to improve its parameter efficiency. We further appreciate the reviewer's suggestions to improve the paper's presentation. Please find our responses below.

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.

“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.

"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.

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.

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).

We thank the reviewer for appreciating the idea with the comment "Using a transformer to perform in-context-learning to fine-tune a transformer sounds like a quite fancy idea." Please find a summary of our work and our responses below.

Summary of our work

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.

Central to the novelty of TINT is the emphasis on the size of the resulting construction. While past works have required billions of parameters to tune even small MLPs, we show that a TINT model with fewer than two billion parameters can tune a 125M parameter auxiliary transformer (e.g., OPT-125m or GPT-2). TINT can be constructed so efficiently in part because of the approximations we make (see Section 2.2).

So, we perform experiments to ensure that the approximations in the construction do not hurt the ability of TINT to tune an auxiliary model. These experiments, described in Section 3.1, quantify how well TINT can perform dynamic evaluation (Definition 3.1) and in-context learning.

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.

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).

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.

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.

We thank the reviewer for appreciating the idea behind fine-tuning via forward passes. Please find our responses below.

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.

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).

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.

"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

We thank the reviewer for their comment on TINT and for pointing out the possible usefulness of TINT in designing new language models for other applications. Please find our responses below.

Issues in writing.

a) Def 2.1: what are relevant auxiliary model weights?

b) Def 2.3: $p_t$ refer to position embeddings?

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.

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.

Why do we use Linear attention in TINT? Can any softmax attention be approximated by linear attention?

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.

Def 2.8 uses finite difference to approximate gradient, why can’t we do the same end-to-end? That is, simulate backpropagation by finite difference with forward passes.

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.

(Our) experiments aren’t directly comparable to Akyurek et al., since they conduct experiments on simulated tasks like linear regression.

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.

No comparisons to Bai et al. who use approximation theory for ICL using transformers. How many parameters in the construction are necessary with the framework of Bai et al.?

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.

References:

1: Information-theoretic lower bounds on the oracle complexity of stochastic convex optimization. Agarwal et al.’12

2: Zeroth-Order Algorithms for Nonconvex-Strongly-Concave Minimax Problems with Improved Complexities. Wang et al.’22

3: Fine-tuning language models with just forward passes. Malladi et al. NeurIPS’23