# Trainable Transformer in Transformer

Abhishek Panigrahi  Sadhika Malladi  Mengzhou Xia  Sanjeev Arora

###### Abstract

Recent works attribute the capability of in-context learning (ICL) in large
pre-trained language models to implicitly simulating and fine-tuning an
internal model (e.g., linear or 2-layer MLP) during inference. However, such
constructions require large memory overhead, which makes simulation of more
sophisticated internal models intractable. In this work, we propose a new
efficient construction, Transformer in Transformer (in short, TinT), that
allows a transformer to simulate and fine-tune more complex models during
inference (e.g., pre-trained language models). In particular, we introduce
innovative approximation techniques that allow a TinT model with less than 2
billion parameters to simulate and fine-tune a 125 million parameter
transformer model within a single forward pass. TinT accommodates many common
transformer variants and its design ideas also improve the efficiency of past
instantiations of simple models inside transformers. We conduct end-to-end
experiments to validate the internal fine-tuning procedure of TinT on various
language modeling and downstream tasks. For example, even with a limited one-
step budget, we observe TinT for a OPT-125M model improves performance by
$4-16\%$ absolute on average compared to OPT-125M. These findings suggest that
large pre-trained language models are capable of performing intricate
subroutines. To facilitate further work, a modular and extensible codebase
111https://github.com/abhishekpanigrahi1996/transformer_in_transformer for
TinT is included.

Machine Learning, ICML

<EMAIL_ADDRESS>
Department of Computer Science, Princeton University

##  1 Introduction

TinT can efficiently perform simulated gradient descent of an auxiliary model.

######  Theorem 1.1.

Consider an auxiliary transformer with $L$ layers, $D_{\text{aux}}$ embedding
dimension, $H_{\text{aux}}$ attention heads, and a maximum sequence length of
$T_{\text{aux}}$. Given a hyperparameter $S$ (see Section 3.1), TinT can
perform an efficient forward pass (Section 3), compute the simulated gradient
(Section 4), and evaluate the updated auxiliary model with a total of
$\left(\frac{(c_{1}S^{2}+c_{3})D_{\text{aux}}^{2}}{\min(H_{\text{aux}},S^{2})}\cdot
D_{\text{aux}}^{2}+c_{2}SD_{\text{aux}}\min(S^{2},H_{\text{aux}})+c_{3}\frac{T_{\text{aux}}D_{\text{aux}}S}{\min(H_{\text{aux}},S^{2})}\right)L$
parameters, with constants $c_{1},c_{2},c_{3}<150$. The TinT model has
$D_{\text{sim}}=SD_{\text{aux}}$ embedding dimension and
$H_{\text{sim}}=\min(S^{2},H_{\text{aux}})$ attention heads. See Table 3 for a
detailed breakdown of the parameters.

Large transformers (Vaswani et al., 2017) have brought about a revolution in
language modeling, with scaling yielding significant advancements in
capabilities (Brown et al., 2020; Chowdhery et al., 2022). These capabilities
include performing in-context learning or following natural language
instructions at inference time.

Researchers have tried to understand how these models can learn new tasks
without parameter updates (Garg et al., 2022; von Oswald et al., 2023; Xie et
al., 2022; Nanda et al., 2023). A popular hypothesis is that in-context
learning corresponds to the transformer (referred to as the simulator from now
on) simulating gradient-based learning of a smaller model (called _auxiliary_
model) that is embedded within it.

From perspective of AI safety and alignment (Amodei et al., 2016; Leike et
al., 2018; Askell et al., 2021), the ability of a larger model to use input
data (which could be arbitrary in a deployed setting) to implicitly train an
auxiliary model feels worrisome. This concern felt minor due to efficiency
considerations: previous analyses and experiments required the auxiliary model
to be quite tiny compared to the simulator. For instance, simulating and
training an auxiliary model that is a linear layer requires tens of millions
of parameters in the simulator (Akyurek et al., 2022). This scaling is even
more dramatic if the auxiliary model is a multi-layer fully-connected net
(Giannou et al., 2023).

Our primary contribution is an explicit and nontrivial construction of a
simulator called TinT that explicitly adapts to the context without parameter
updates. In particular, we show that a forward pass through a modestly sized
TinT can involve gradient-based training of an auxiliary model that is itself
a large transformer. For example, we show that TinT with 2B parameters can
faithfully simulate fine-tuning a 125M parameter auxiliary transformer in a
single forward pass. (Prior constructions would have required trillions of
parameters in the simulator for a far simpler auxiliary model.)

Our main result is described in Theorem 1.1, which details how the size of
TinT depends on the auxiliary model. Our construction is generally applicable
to diverse variants of pre-trained language models. The rest of the paper is
structured to highlight the key design choices and considerations in TinT.

  1. 1.

Section 2 discusses the overall design decisions required to make TinT,
including how the simulator can read from and write to the auxiliary model and
how the data must be formatted.

  2. 2.

Section 3 uses the linear layer as an example to describe how highly
parallelized computation and careful rearrangement of activations enable TinT
to efficiently simulate the forward pass of the auxiliary model.

  3. 3.

Section 4 describes how TinT uses first-order approximations and stop
gradients to compute the _simulated gradient_ of the auxiliary model.

  4. 4.

Section 5 performs experiments comparing TinT to suitable baselines in
language modeling and in-context learning settings. Our findings validate that
the simulated gradient can effectively update large pre-trained auxiliary
models. Notably, we instantiate TinT in a highly extensible codebase, making
TinT the first such construction to undergo end-to-end evaluation.

Due to the complexity of the construction, we defer the formal details of TinT
to the appendix.

Figure 1: The overall structure of TinT (see Section 2 for an overview). Each
forward, backward, and descent module is represented using combinations of
linear, self-attention, layernorm, and activation layers. The input consists
of prefix embeddings (Definition 2.1) that represent relevant auxiliary model
parameters in each layer followed by natural language input. A prefix mask
separates the train and test segments of the input (§2.1).  Figure 2: TinT
simulates the forward pass of a linear layer with a $H_{\text{sim}}$-head
attention layer ($H_{\text{sim}}=6$ here). We stack $S$ weights per prefix
embedding to reduce the number of prefix embeddings required ($S=2$ here). We
furthermore shard each weight and token embedding ${\bm{x}}_{t}$ into
$S^{\prime}$ shards and compute inner products of each shared in parallel
using $S\times S^{\prime}$ attention heads ($S^{\prime}=3$ here). Please see
Section 3.1.

##  2 Design Considerations

Our goal is to construct a simulator that can train an auxiliary model over
the course of an inference pass. This procedure requires four steps:

  1. 1.

Forward Pass: A forward pass to compute the auxiliary model output
$f(\xi;{\bm{\theta}}_{\text{aux}})$ on training input $\xi$ and a loss
$\mathcal{L}$.

  2. 2.

Backward Pass: Backpropagation to compute the gradient of the auxiliary model
$\nabla_{{\bm{\theta}}_{\text{aux}}}\mathcal{L}(f(\xi;{\bm{\theta}}_{\text{aux}}))$.

  3. 3.

Parameter Update: Update the auxiliary model using gradient descent, setting
${\bm{\theta}}_{\text{aux}}^{\prime}={\bm{\theta}}_{\text{aux}}-\eta\nabla_{{\bm{\theta}}_{\text{aux}}}\mathcal{L}(f(\xi;{\bm{\theta}}_{\text{aux}}))$.

  4. 4.

Output: Output next-token predictions
$f(\xi^{\prime};{\bm{\theta}}_{\text{aux}}^{\prime})$ on a test input
$\xi^{\prime}$ using the updated auxiliary model.

Note that steps 1-3 can be looped to train the auxiliary model for a few
steps222Looping steps 1-3 scales the depth of the simulator model., either on
the same training data or on different training data for each step, before
evaluating it on the test input (Giannou et al., 2023). The above method
highlight two crucial features of the simulator: (1) it has access to some
amount of training data, and (2) it can use (i.e., read) and update (i.e.,
write) the auxiliary model. Below, we discuss how to design a modest-sized
simulator around these two considerations.

###  2.1 Input structure

For simplicity, we describe only one update step on a single batch of training
data $\xi$ but note that our formal construction and our experiments handle
multiple training steps (see Definition 5.1). Steps 1 and 4 show that the
simulator must access some training data $\xi$ to train the auxiliary model
and some testing data $\xi^{\prime}$ on which it evaluates the updated
auxiliary model. For the sake of illustration we consider the following simple
setting: given a sequence of input tokens ${\bm{e}}_{1},...,{\bm{e}}_{T}$, we
split it into training data $\xi={\bm{e}}_{1},...,{\bm{e}}_{r}$ and testing
data $\xi^{\prime}={\bm{e}}_{r+1},...,{\bm{e}}_{T}$.

Suppose $\xi$ contains an in-context input-output exemplar and $\xi^{\prime}$
contains a test input. Then, the simulator performs a very natural operation
of training the auxiliary model on a task-specific example and outputs results
for the test example.

On the other hand, if the input is not specially formatted, $\xi$ and
$\xi^{\prime}$ may simply contain some natural language tokens. In this case,
the simulator is using the first part of the context tokens to do a quick
fine-tune of the auxiliary for some task before outputting the subsequent
tokens with the auxiliary model. In a worst-case scenario, users might provide
harmful contents, leading the model to implicitly fine-tune on them and
potentially output even more harmful content.

Our experiments consider many options for splitting a sequence into $\xi$ and
$\xi^{\prime}$, and we defer a more detailed discussion of possible setups to
Section 5.

#### Accessing Training Labels.

The simulator must be able to see the labels of the training tokens in order
to compute the loss $\mathcal{L}$ (usually, the autoregressive cross-entropy
loss) in step 1. For example, in Figure 1, when we compute the loss for the
token ${\bm{e}}_{2}$ in the second position, we need to use its label
${\bm{e}}_{3}$ in the third position. However, this is not possible if the
simulator uses strictly autoregressive attention (Appendix H contains a more
general discussion). We thus use a bidirectional attention mask on the
training tokens and autoregressive attention on the evaluation portion. We
note that encoding relevant (e.g., retrieved) context with bidirectional
attention is a popular way to improve autoregressive capabilities in language
modeling and natural language tasks (Raffel et al., 2020; Borgeaud et al.,
2022; Izacard & Grave, 2020; Izacard et al., 2023; Wang et al., 2023a; Tay et
al., 2022). This empirical approach is similar in motivation to how TinT uses
a few context tokens to adapt the auxiliary model to a given input. Having
established the training and testing data, we can now move to discussing how
the simulator can access (i.e., read) and update (i.e., write to) the
auxiliary model at inference time.

###  2.2 Read and write access to auxiliary model

As discussed in the start of this section, the simulator must have read and
write access to the parameters of the auxiliary model. Crucially, the
simulator must do at least two forward passes through the auxiliary model, one
with the current parameters ${\bm{\theta}}_{\text{aux}}$ and one with the
updated parameters ${\bm{\theta}}_{\text{aux}}^{\prime}$.

The straightforward way to simulate the forward pass of the auxiliary model
would be to store its weights in the simulator’s weights and run a forward
pass as usual. One can analogously simulate the backward pass according to the
loss $\mathcal{L}$ to compute the gradients. However, the simulator cannot
update its own weights at inference time, so this strategy would not permit
the model to write the updated parameters
${\bm{\theta}}_{\text{aux}}^{\prime}$ and later read them when simulating the
second forward pass. Therefore, the auxiliary model
${\bm{\theta}}_{\text{aux}}$ must be available in the activations of the
simulator.

To this end, Wei et al. (2021); Perez et al. (2021) model the simulator after
a Turing machine, where the activation
${\bm{e}}_{t}^{(\ell)}\in\mathbb{R}^{D_{\text{sim}}}$ in each layer acts as a
workspace for operations, and computation results are copied to and from
memory using attention operations. In this paradigm, if $D_{\text{aux}}=768$,
computing a dot product $\langle{\bm{w}},{\bm{x}}_{t}^{(\ell)}\rangle$ with
weight ${\bm{w}}\in\mathbb{R}^{768}$ requires at least $6.4$ million
parameters in the simulator333 Using a feedforward module to mimic the dot
product (as in Akyurek et al. (2022), see thm. C.4), where the simulator
embedding comprises $[{\bm{w}},{\bm{x}}_{t}]\in\mathbb{R}^{1536}$,
necessitates a minimum of $4.7$ million parameters. Using an attention module
to copy the weight from memory adds another $1.7$ million parameters.. Given
the pervasiveness of dot products in neural network modules, this strategy
would yield a simulator with trillions of parameters.

Alternatively, one can store parameters in the first few context tokens and
allow the attention modules to attend to those tokens (Giannou et al., 2023).
This removes the need for copying and token-wise operations. Then, the same
dot product requires only a self-attention module with $1.7$ million
parameters. We thus adopt this strategy to provide relevant auxiliary model
weights as _prefix embeddings_.

######  Definition 2.1 (Prefix Embeddings).

$\\{{\bm{v}}_{j}^{(\ell)}\\}_{j=1}^{K}$ denotes the $K$ prefix embeddings at
the $\ell$th layer in TinT. These contain _relevant_ auxiliary model weights
or simulated activations.

We now consider how to efficiently simulate the building block of neural
networks: matrix-vector multiplication. In the next section, we demonstrate
that a careful construction of the prefix embeddings enables efficient
parallelizaton of matrix-vector products across attention heads.

##  3 Efficient Forward Propagation

We now discuss how TinT performs a highly efficient forward pass through the
auxiliary model. Here, we focus on the linear layer because it is repeated
many times in various transformer modules (e.g., in self-attention), so
improving the efficiency dramatically reduces TinT’s size.

######  Definition 3.1 (Linear layer).

For a weight ${\bm{W}}\in\mathbb{R}^{D_{\text{aux}}\times D_{\text{aux}}}$, a
linear layer takes ${\bm{x}}\in\mathbb{R}^{D_{\text{aux}}}$ as input and
outputs ${\bm{y}}={\bm{W}}{\bm{x}}$.444Linear layers are applied token-wise,
so we can consider a single position $t$ without loss of generality.

We compute ${\bm{y}}$ coordinate-wise, i.e.,
$\langle{\bm{w}}_{i},{\bm{x}}_{t}\rangle$ for all $i\in[D_{\text{aux}}]$,
where ${\bm{w}}_{i}$ is the $i$th row of ${\bm{W}}$. The simulator represents
$\langle{\bm{w}}_{i},{\bm{x}}_{t}\rangle$ as an attention score between the
row ${\bm{w}}_{i}$ and the input ${\bm{x}}_{t}$. So, the input embeddings
${\bm{e}}_{t}$ contain ${\bm{x}}_{t}$ in the first $D_{\text{aux}}$
coordinates, and the rows $\\{{\bm{w}}_{i}\\}$ of the weight matrix ${\bm{W}}$
are in prefix embeddings $\\{{\bm{v}}_{j}\\}$ (def. 2.1).

We strategically distribute the weights (§3.1) and aggregate the parallelized
computation results (§3.2). As we briefly mentioned in the previous section, a
straightforward construction of the linear layer would use the context and
attention heads inefficiently. Our construction instead parallelizes the
computation across attention heads in such a way that aggregating the output
of the linear operation can also be conducted efficiently.

###  3.1 Stacking and Sharding

We partition the inner product computation across attention heads by carefully
rearranging the weights and activations via stacking and sharding (Figure 2).

Instead of representing each weight ${\bm{w}}_{i}$ as its own prefix token
${\bm{v}}_{i}$, we _stack_ $S$ weights on top of each other to form each
prefix embedding ${\bm{v}}_{i}$. $S$ drives a trade-off between the embedding
dimension of the TinT, $D_{\text{sim}}=D_{\text{aux}}S$, and the context
length to the TinT, $T_{\text{sim}}=K+T_{\text{aux}}$. We set $S=4$.

A simple strategy now would be to use different attention heads to operate on
different rows; however, this would still use only $S$ attention heads whereas
we could parallelize across many more heads. We instead parallelize across
more attention heads, where each head is responsible for computing the inner
product on a subset of the coordinates. We _shard_ each individual weight and
the activation into $S^{\prime}$ parts and compute the inner product on each
of the $S^{\prime}$ parts in parallel We set $S$ and $S^{\prime}$ such that
$H_{\text{sim}}=S\times S^{\prime}$, thereby using all of TinT heads to
efficiently compute the dot products.

###  3.2 Efficient Aggregation

The attention module outputs a sparse matrix with shape
$(D_{\text{sim}}/H_{\text{sim}})\times H_{\text{sim}}$ containing the inner
products on various subsets of the coordinates in its entries. To complete the
linear forward pass, we need to sum the appropriate terms to form a
$D_{\text{sim}}$-length vector with ${\bm{W}}{\bm{x}}$ in the first
$D_{\text{aux}}$ coordinates. Straightforwardly summing along an axis
aggregates incorrect terms, since the model was sharded. On the other hand,
rearranging the matrix would require an additional $D_{\text{sim}}\times
D_{\text{sim}}$ linear layer. Instead, TinT saves a factor of
$H_{\text{sim}}\times$ parameters by leveraging the local structure of the
attention output. We illustrate this visually in Section D.1. This procedure
requires $D_{\text{sim}}^{2}/H_{\text{sim}}+D_{\text{sim}}H_{\text{sim}}$
parameters. This efficient aggregation also compresses the constructions for
the TinT’s backpropagation modules for layer normalization and activations
(Appendices F and G).

##  4 Simulated Gradient

TinT adapts backpropagation to compute gradients (Figure 1). We aim to train a
capable (i.e., pre-trained) auxiliary model for just a few steps, so high
precision gradients may be unnecessary. Instead, TinT performs an approximate
backpropagation. TinT then uses this _simulated gradient_ to update the
auxiliary model. Prior works computed similar approximate gradients in hopes
of more faithfully modeling neurobiology (Scellier & Bengio, 2017; Hinton,
2022) or improving the efficiency of training models (Hu et al., 2021; Malladi
et al., 2023). We note that the approximations in the simulated gradients can
be made stronger at the cost of enlarging TinT. Indeed, one could construct a
simulator to _exactly_ perform the procedure outlined in §2, though it would
be orders of magnitude larger than TinT. For brevity’s sake, we focus on the
key approximations and design choices and defer formal details to the
appendix.

###  4.1 First-order approximations

We use first-order approximations of gradients to backpropagate through the
layer normalization layer.555We discuss a layer normalization layer
$f_{\text{ln}}$ without scale and bias parameters, but Appendix F contains a
general construction. It normalizes the input using its mean and standard
deviation across the input dimensions. Since the operation is token-wise, we
can consider a single position $t$ without loss of generality.

######  Definition 4.1 (Layer normalization).

A layer normalization layer $f_{\text{ln}}$ takes input
${\bm{x}}\in\mathbb{R}^{D_{\text{aux}}}$ and outputs
${\bm{y}}=({\bm{x}}-\mu)/\sigma$, where $\mu$ and $\sigma$ denote its mean and
standard deviation.

High precision gradients: Formally, for input-output pair
$({\bm{x}},{\bm{y}})$, we can compute the gradients $\partial_{{\bm{y}}}$,
$\partial_{{\bm{x}}}$ with chain rule:

$\displaystyle\partial_{{\bm{x}}}=\left(\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}\right)^{\top}\partial_{{\bm{y}}}$
$\displaystyle=\frac{1}{\sigma}\left(\langle\partial_{{\bm{y}}},{\bm{y}}\rangle{\bm{y}}+\partial_{{\bm{y}}}-\frac{1}{D_{\text{aux}}}\sum_{i=1}^{D_{\text{aux}}}\partial_{y_{i}}\right).$
(1)

Inefficiency of exact computation: A TinT layer simulating backpropagation
through an auxiliary’s layer normalization layer receives
$\partial_{{{\bm{y}}}_{t}}$ and ${\bm{x}}_{t}$ in its input embeddings. We go
through the exact gradient and why it is inefficient.

For exact computation one could first compute ${\bm{y}}_{t}$ using a
normalization layer and store in the embeddings. However, inefficiency arises
from computing the term
$\langle\partial_{{{{\bm{y}}}_{t}}},{\bm{y}}_{t}\rangle{\bm{y}}_{t}$. To
calculate $\langle\partial_{{{{\bm{y}}}_{t}}},{\bm{y}}_{t}\rangle{\bm{y}}_{t}$
at each token position $t$, we could either: (1) use a two-layer MLP that
focuses on each token separately, or (2) a single self-attention module to
treat the operation as a sequence-to-sequence task.

For (1) we could initially compute
$\langle\partial_{{{\bm{y}}}_{t}},{\bm{y}}_{t}\rangle$ via an MLP, followed by
computation of
$\langle\partial_{{{\bm{y}}}_{t}},{\bm{y}}_{t}\rangle{\bm{y}}_{t}$ using
another MLP. The element-wise multiplication in embeddings would be
facilitated with a nonlinear activation function like GeLU (Akyurek et al.,
2022) (refer to thm. C.4 for details). However, this approach would need
substantial number of simulator parameters to represent the MLPs.

Alternatively, we could use a single self-attention module. Constructing such
a module would require careful engineering to make sure the input tokens only
attend to themselves while keeping an attention score of $0$ to others. If we
used a linear attention, we would need to space out the gradient
$\partial_{{{\bm{y}}}_{t}}$ and ${\bm{x}}_{t}$ in each position $t$, such that
the attention score is $0$ between different tokens. This would require an
embedding dimension proportional to the context length. On the other hand, if
we used a softmax attention module, we would need an additional superfluous
token in the sequence. Then, a token at position $t$ would attend to itself
with attention $\langle\partial{{\bm{y}}}_{t},{\bm{y}}_{t}\rangle$ and to the
extra token with an attention score of
$1-\langle\partial{{\bm{y}}}_{t},{\bm{y}}_{t}\rangle$. The extra token would
return a value vector $0$. To avoid such inefficiency, we opt for a first-
order approximation instead.

Efficient approximation: Instead of explicitly computing each term in the
chain rule of $\left(\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}\right)^{\top}\partial_{{\bm{y}}}$
in Eq. 1, we instead use a first order Taylor expansion of $f_{\text{ln}}$.

$\displaystyle
f_{\text{ln}}({\bm{x}}+\epsilon\partial_{{\bm{y}}})=f_{\text{ln}}({\bm{x}})+\epsilon\left(\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}\right)\partial_{{\bm{y}}}+\mathcal{O}(\epsilon^{2}).$

Rearranging allows us to write

$\displaystyle\left(\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}\right)\partial_{{\bm{y}}}=\frac{1}{\epsilon}\left(f_{\text{ln}}({\bm{x}}+\epsilon\partial_{{\bm{y}}})-f_{\text{ln}}({\bm{x}})\right)+\mathcal{O}(\epsilon).$

Similar to the computation of Eq. 1, we can show

$\displaystyle\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}=\frac{1}{\sigma}\left((1-{D_{\text{aux}}}^{-1}){\bm{I}}-f_{\text{ln}}({\bm{x}})f_{\text{ln}}({\bm{x}})^{\top}\right).$

Because $\partial f_{\text{ln}}({\bm{x}})/\partial{\bm{x}}$ is symmetric666For
a linear function $f$ with matrix ${\bm{W}}$, $\frac{\partial
f({\bm{x}})}{\partial{\bm{x}}}={\bm{W}}$. Since ${\bm{W}}$ may not be a
symmetric matrix, this method can’t be generally applied to approximately
backpropagate linear layers or causal self-attention layers., we can write

$\displaystyle\left(\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}\right)^{\top}\partial_{{\bm{y}}}$
$\displaystyle=\left(\frac{\partial
f_{\text{ln}}({\bm{x}})}{\partial{\bm{x}}}\right)\partial_{{\bm{y}}}$
$\displaystyle=\frac{1}{\epsilon}\left(f_{\text{ln}}({\bm{x}}+\epsilon\partial_{{\bm{y}}})-f_{\text{ln}}({\bm{x}})\right)+\mathcal{O}(\epsilon).$

Then, ignoring the small error term, we can use just two linear layers,
separated by a normalization layer, to simulate the approximation.

###  4.2 Fuzzy backpropagation via stop gradients

Self-attention is inherently quadratic, because it uses the keys and queries
to compute attention scores between every possible pair of tokens in the
sequence. These scores then linearly combine the value vectors (see def. B.1).
Computing the gradient exactly is thus a very complex operation. Instead, we
stop the gradient computation through attention scores in the self-attention
layer. For similar reasons, we only update the value parameter in the self-
attention module.

Gradient backpropagation: For an input, output sequence pair
$\\{{\bm{y}}_{t}\\},\\{{\bm{y}}_{t}\\}$, if
$\\{{\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}\\}$ denote the intermediate query,
key, value vectors, on gradients $\\{\partial_{{\bm{y}}_{t}}\\}$,
$\\{\partial_{{\bm{x}}_{t}}\\}$ is given via the chain rule:

$\displaystyle\partial_{{\bm{x}}_{t}}$
$\displaystyle={\bm{Q}}^{\top}\partial_{{\bm{q}}_{t}}+{\bm{K}}^{\top}\partial_{{\bm{k}}_{t}}+{\bm{V}}^{\top}\partial_{{\bm{v}}_{t}}.$
(2)

Here, ${\bm{V}},{\bm{K}},{\bm{Q}}$ denote the query, key, and value matrices.

Inefficiency in exact computation: Here, we demonstrate that simulating
computation of the three terms in Eq. 2 is inefficient, because
$\partial_{{\bm{q}}_{t}},\partial_{{\bm{k}}_{t}}$ depend on the derivatives
w.r.t. the attention scores. As an example, we focus on
$\partial_{{\bm{k}}_{t}}$:

$\displaystyle\partial_{{\bm{k}}_{t}}=\sum_{j}a_{t,j}((\partial_{{\bm{y}}_{t}})^{\top}{\bm{v}}_{j})({\bm{k}}_{j}-\sum_{j^{\prime}}a_{t,j^{\prime}}{\bm{k}}_{j^{\prime}}).$

Computing this term would require us at least 2 self-attention layers and an
MLP layer. The first attention layer would compute
$(\partial_{{\bm{y}}_{t}})^{\top}{\bm{v}}_{j}$ for different token pairs,
similar to the forward simulation of a linear layer with linear attention
(§3). These would be then multiplied to the pair-wise attention scores
$a_{t,j}$ with an MLP to compute
$a_{t,j}((\partial_{{\bm{y}}_{t}})^{\top}{\bm{v}}_{j})$, with elementwise
product would be facilitated by GeLU non-linearity (thm. C.4). These would be
finally used by an attention layer to combine the different key vectors. A
similar simulation would be necessary to compute $\partial_{{\bm{q}}_{t}}$.

Stop gradients through query and key vectors: In order to reduce the necessary
resources, we ignore the query and key gradients in Eq. 2. When we ignore
these gradient components, $\\{\partial_{{\bm{x}}_{t}}\\}$ can be simplified
as

$\displaystyle\partial_{{\bm{x}}_{t}}\approx{\bm{V}}^{\top}\partial_{{\bm{v}}_{t}}={\bm{V}}^{\top}\sum_{j}a_{j,t}\partial_{{\bm{y}}_{t}}.$
(3)

A single self-attention layer can compute this by using the attention scores
to combine the token-wise gradients.

Why won’t it hurt performance? Estimating $\partial_{{\bm{x}}_{t}}$ as
described is motivated by recent work (Malladi et al., 2023) showing that
fuzzy gradient estimates don’t adversely affect fine-tuning of pre-trained
models. Furthermore, we theoretically show that when the attention head for
each position pays a lot of attention to a single token (i.e., behaves like
hard attention (Perez et al., 2021)), the approximate gradient in Eq. 3 is
entry-wise close to the true gradients (thm. E.5).

The other approximation is to update only the value parameters ${\bm{V}}$ of
the auxiliary model (§E). This is motivated by parameter efficient fine-tuning
methods like LoRA (Hu et al., 2021) and IA3 (Liu et al., 2022), which restrict
the expressivity of the gradient updates without degrading the quality of the
resulting model. We similarly show in the next section that the simulated
gradients in TinT can effectively tune large pre-trained transformers.

Table 1: Language modeling results on WikiText-103. We use $30\%,50\%,70\%$
and $90\%$ of sequences for training in the language modeling setting (§5.2).
TinT improves the auxiliary model perplexities by $0.3-0.7$ absolute on
average. The small perplexity difference between the TinT and explicitly
updating the auxiliary model suggests that the simulated gradient (Section 4)
can still effectively fine-tune the auxiliary model.

|  | Training proportion
---|---|---
| Evaluating with | $30\%$ | $50\%$ | $70\%$ | $90\%$
GPT-2 | Auxiliary Model | 25.6 | 24.9 | 24.5 | 23.3
Fine-tuning | 24.9 | 24.0 | 23.5 | 22.2
TinT | 25.1 | 24.3 | 23.8 | 22.6
OPT-125M | Auxiliary Model | 29.6 | 28.8 | 28.0 | 28.0
Fine-tuning | 29.0 | 28.2 | 27.4 | 27.4
TinT | 29.3 | 28.4 | 27.5 | 27.4

Figure 3: Different settings in few-shot learning ($k=3$) using TinT. The
Single mode (left) treats each example as a training datapoint, and the
auxiliary model is updated with a batch of inputs (see def. 5.1). The Multi.
mode (right) concatenates all examples to form a single input and uses batch
size $1$ in def. 5.1. For Label loss, only underlined label words are used as
training signal, while full context loss includes all tokens.

##  5 Experiments

Table 2: Zero-shot and few-shot in-context learning results across $7$
downstream tasks. All the few-shot results are averaged over three training
seeds. TinT consistently surpasses its auxiliary model and achieves comparable
performance to one-off dynamic evaluation. TinT outperforms auxiliary models
by $3-4\%$ and $12-16\%$ absolute points on average in $0$-shot and $32$-shot
experiments respectively. TinT performs competitively with a similar-sized
pre-trained model (opt-1.3b) in both $0$-shot and $32$-shot settings. We show
the standard deviation for few-shot settings in parentheses.

Model | Shots | Subj | AGNews | SST2 | CR | MR | MPQA | Amazon | Avg.
---|---|---|---|---|---|---|---|---|---
OPT-125m | $0$ | $64.0$ | $66.0$ | $70.5$ | $64.5$ | $71.0$ | $68.0$ | $76.5$ | $68.6$
OPT-1.3b | $0$ | $59.0$ | $55.5$ | $54.0$ | $50.5$ | $52.5$ | $74.0$ | $57.0$ | $57.5$
OPT-125m Fine-tuning | $0$ | $71.0$ | $67.0$ | $79.5$ | $71.5$ | $70.0$ | $68.0$ | $85.5$ | $73.2$
OPT-125m TinT | $0$ | $67.5$ | $66.0$ | $76.5$ | $69.0$ | $76.0$ | $70.5$ | $78.5$ | $72.0$
OPT-125m | $32$ | $58.7_{(4.9)}$ | $33.7_{(8.4)}$ | $50.8_{(1.2)}$ | $51.3_{(1.9)}$ | $50.0_{(0.0)}$ | $54.3_{(2.5)}$ | $55.0_{(6.7)}$ | $50.5_{(1.9)}$
OPT-1.3b | $32$ | $74.2_{(6.1)}$ | $71.3_{(5.3)}$ | $89.8_{(3.6)}$ | $71.5_{(4.5)}$ | $68.3_{(6.1)}$ | $81.7_{(3.3)}$ | $70.3_{(9.9)}$ | $75.3_{(0.4)}$
OPT-125m Fine-tuning | $32$ | $78.0_{(1.4)}$ | $66.7_{(1.6)}$ | $71.5_{(1.4)}$ | $73.7_{(3.3)}$ | $72.0_{(0.0)}$ | $80.7_{(0.6)}$ | $79.8_{(0.2)}$ | $74.6_{(2.7)}$
OPT-125m TinT | $32$ | $82.3_{(2.7)}$ | $69.3_{(0.9)}$ | $73.7_{(0.8)}$ | $75.7_{(1.9)}$ | $72.3_{(1.2)}$ | $83.2_{(1.0)}$ | $78.2_{(0.2)}$ | $76.4_{(0.7)}$

We evaluate the performance of the TinTs constructed using GPT2 and OPT-125M
as auxiliary models. The findings from our experiments in the language
modeling and in-context learning settings confirm that fine-tuning with the
simulated gradients (Section 4) still allows for effective learning in the
auxiliary model. We loop the training steps (i.e., steps 1-3) outlined in
Section 2 to accommodate solving real-world natural language tasks. We
formalize the setting below.

###  5.1 Setting: $N$-step Fine-Tuning

We formalize the procedure in Section 2 to construct a suitable setting in
which we can compare TinT to explicitly training the auxiliary model.

######  Definition 5.1 ($N$-step Fine-Tuning).

Given a batch of training datapoints $\xi_{1},\cdots,\xi_{B}$ and a validation
input $\xi^{\prime}$, we compute and apply gradient updates on the auxiliary
model ${\bm{\theta}}_{\text{aux}}$ for timesteps $t=0,...,N-1$ as

${\bm{\theta}}_{\text{aux}}^{t+1}={\bm{\theta}}_{\text{aux}}^{t}-\eta\sum_{i=1}^{B}\nabla_{\bm{\theta}}\mathcal{L}(f(\xi_{i};{\bm{\theta}}_{\text{aux}}^{t}))$

where $\eta$ is the learning rate and $\mathcal{L}$ is a self-supervised loss
function on each input $\xi_{i}$. Then, we evaluate the model
${\bm{\theta}}_{\text{aux}}^{N}$ on $\xi^{\prime}$.
${\bm{\theta}}_{\text{aux}}^{0}$ denotes the pre-trained auxiliary model.

Below, we instantiate this setting with text inputs of different formats and
different self-supervised loss functions $\mathcal{L}$. To manage
computational demands, we limit $N$ to $3$ or fewer.777Performing many
gradient steps scales the depth of TinT and makes experimentation
computationally infeasible.

###  5.2 Case Study: Language Modeling

The first case we consider is language modeling, where the input data
${\bm{e}}_{1},...,{\bm{e}}_{T}$ is natural language without any additional
formatting. We use a batch size of $1$ in def. 5.1, and delegate
$\xi_{1}={\bm{e}}_{1},...,{\bm{e}}_{t}$ and
$\xi^{\prime}={\bm{e}}_{t+1},...,{\bm{e}}_{T}$. The loss $\mathcal{L}$ is the
sum of the token-wise autoregressive cross-entropy loss in the sequence
$\xi_{1}$. For example, given an input Machine learning is a useful tool for
solving problems., we use the red part as the training data $\xi_{1}$, and the
brown part as the validation data $\xi^{\prime}$. We perform language modeling
experiments on WikiText-103 (Merity et al., 2016) and vary the number of
tokens $t$ used as training data $\xi$.

Results. In Table 1, we observe that TinT achieves a performance comparable to
explicit fine-tuning of the auxiliary model, indicating that the simulated
gradient (Section 4) is largely effective for fine-tuning. Both TinT and
explicitly fine-tuning the auxiliary model show improvement over the base
model, confirming that minimal tuning on the context indeed enhances
predictions on the test portion.

###  5.3 Case Study: In-Context Learning

For in-context learning, we consider input data to be a supervised
classification task transformed into a next-token prediction task using
surrogate labels (see Figure 3). Using binary sentiment classification of
movie reviews as an example, given an input (e.g., the review), the model’s
predicted label is computed as follows. First, we design a simple task-
specific prompt (e.g., “Sentiment:”) and select label words $c_{1},...,c_{n}$
to serve as surrogates for each class (e.g., “positive” and “negative”). Then,
we provide the input along with the prompt to the model, and the label
assigned the highest probability is treated as the model’s prediction. We
describe the zero-shot and few-shot settings below.

Zero-shot. In the zero-shot setting, we are given text with the first $T-1$
tokens as the input text and final token as the surrogate text label. Hence,
we adapt def. 5.1 to use batch size $B=1$, training data
$\xi_{1}=x_{1},...,x_{T-1}$, and testing data $\xi^{\prime}=x_{T}$. The loss
$\mathcal{L}$ is again the sum of the token-wise autoregressive cross-entropy
losses.

Few-shot. In the few-shot setting, we are given input texts that are a
concatenation of $k$ sequences $\xi_{1},\cdots,\xi_{k}$. Each sequence
contains the input text followed by the surrogate label for the in-context
exemplar. These $k$ exemplars are followed by test data $\xi^{\prime}$. In
this case, we can compute the gradient updates to ${\bm{\theta}}_{\text{aux}}$
in two different ways (Figure 3). The first setting, denoted Single, treats
the $k$ sequences as a batch of $B=k$ training datapoints
$\xi_{1},...,\xi_{B}$. The second setting, denoted Multi, treats the
concatenation of the $B$ sequences as a single training datapoint $\xi_{1}$.
Furthermore, $\mathcal{L}$ for a training datapoint can be defined in two
different ways. The first setting, denoted as Full context loss, defines
$\mathcal{L}$ for a training datapoint $\xi_{i}$ as the sum of cross entropy
loss over all tokens. The second setting, denoted as Label loss, defines
$\mathcal{L}$ for a training datapoint $\xi_{i}$ in def. 5.1 as the sum of
cross entropy loss over the surrogate label tokens.

Tasks. We evaluate 7 classification tasks for zero-shot and few-shot settings:
SST-2 (Socher et al., 2013), MR (Pang & Lee, 2004), CR (Hu & Liu, 2004), MPQA
(Wiebe et al., 2005), Amazon Polarity (Zhang et al., 2015), AGNews (Zhang et
al., 2015), and Subj (Pang & Lee, 2005).

Model. We compare a TinT model that uses an OPT-125m pre-trained model as its
auxiliary model against two alternative approaches: (1) directly fine-tuning
OPT-125m, and (2) performing standard evaluation using OPT-1.3b, which is of a
similar size to TinT.888Our construction is generally applicable to diverse
variants of pre-trained language models (Appendix K).

Observations. We observe that inferences passes through TinT perform on par
with directly fine-tuning the auxiliary model, affirming the validity of the
construction design (see Section 2). As expected, TinT outperforms the base
auxiliary model, since it simulates training the auxiliary model. More
intriguingly, TinT demonstrates performance comparable to a pre-trained model
of similar size (OPT-1.3b). This suggests that the capabilities of existing
pre-trained models may be understood via the simulation of smaller auxiliary
models. For further details and results of the experiments, please refer to
Appendix L.

##  6 Related Work

Gradient-based learning and in-context learning: Several works relate in-
context learning to gradient-based learning algorithms. Bai et al. (2023)
explicitly constructed transformers to simulate simple gradient-based learning
algorithms. Mahankali et al. (2023); Ahn et al. (2023) suggested one attention
layer mimics gradient descent on a linear layer, and Zhang et al. (2023a)
showed polynomial convergence. Cheng et al. (2023); Han et al. (2023) extended
these ideas to non-linear attentions. Experiments in Dai et al. (2022) suggest
that LLM activations during in-context learning mirror fine-tuned models.
These works focus on using a standard transformer for the simulator and hence
cannot accommodate more complex auxiliary models; on the other hand, our work
uses structural modifications and approximations to construct an efficient
simulator for complex auxiliary models. Our work in contrast attempts to build
even stronger transformers by introducing few structural modifications that
can run gradient descent on auxiliary transformers.

Transformer Expressivity: Perez et al. (2021); Pérez et al. (2019) show that
Transformers with hard attention are Turing complete, and Wei et al. (2021)
construct transformers to study statistical learnability, but the proposed
constructions are extremely large. Other works have investigated encoding
specific algorithms in smaller simulators, e.g. bounded-depth Dyck languages
(Yao et al., 2021), modular prefix sums (Anil et al., 2022), adders (Nanda et
al., 2023), regular languages (Bhattamishra et al., 2020), and sparse logical
predicates (Edelman et al., 2022). Liu et al. (2023) aim to understand
automata-like mechanisms within transformers. Ba et al. (2016) connect self-
attention and fast weight programmers (FWPs), which compute input-dependent
weight updates during inference. Follow-up works (Schlag et al., 2021; Irie et
al., 2021) use self-attention layers to update linear and recurrent networks
during inference. Clark et al. (2022) add and efficiently tune Fast Weights
Layers (FWL) on a frozen pre-trained model.

##  7 Discussion

We present a parameter-efficient construction TinT capable of simulating
gradient descent on an internal transformer model during inference. Using
fewer than 2 billion parameters, it can simulate fine-tuning a 125 million
transformer (e.g., GPT-2) internally, dramatically reducing the scale required
by previous works. Language modeling and in-context learning experiments
demonstrate that the efficient approximations still allow the TinT to fine-
tune the model. Our work emphasizes that the inference behavior of complex
models may rely on the training dynamics of smaller models. As such, the
existence of TinT has strong implications for interpretability and AI
alignment research.

While our work represents a significant improvement over previous simulations
in terms of auxiliary model complexity, similar to prior research in this
area, our insights into existing pre-trained models are limited. Furthermore,
we have not yet examined potential biases that may arise in the auxiliary
models due to one-step gradient descent. We plan to investigate these aspects
in future work.

## Impact Statements

We note that the construction of TinT does not appear to increase the
probability of harmful behavior, because the construction’s primary objective
is to implicitly tune an internal model (§2). Such tuning has been possible
for a long time and is not made more expressive by TinT.

Our findings suggest that existing transformer-based language models can
plausibly possess the ability to learn and adapt to context by internally
fine-tuning a complex model _even during inference_. Consequently, although
users are unable to directly modify deployed models, these models may still
undergo dynamic updates while processing a context left-to-right, resulting in
previously unseen behavior by the time the model reaches the end of the
context. This has significant implications for the field of model alignment.
It is challenging to impose restrictions on a model that can perform such
dynamics updates internally, so malicious content can influence the output of
deployed models.

Alternatively, we recognize the potential benefits of pre-training constructed
models that integrate explicit fine-tuning mechanisms. By embedding the
functionalities typically achieved through explicit fine-tuning, such as
detecting malicious content and intent within the models themselves, the need
for external modules can be mitigated. Pre-training the constructed model may
offer a self-contained solution for ensuring safe and responsible language
processing without relying on external dependencies.

## Acknowledgements

The authors acknowledge funding from NSF, ONR, Simons Foundation, and DARPA.
We thank Danqi Chen, Jason Lee, Zhiyuan Li, Kaifeng Lyu, Simran Kaur, Tianyu
Gao, and Colin Wang for their suggestions and helpful discussions at different
stages of our work.

## References

  * Ahn et al. (2023) Ahn, K., Cheng, X., Daneshmand, H., and Sra, S.  Transformers learn to implement preconditioned gradient descent for in-context learning.  _arXiv preprint arXiv:2306.00297_ , 2023.
  * Akyurek et al. (2022) Akyurek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D.  What learning algorithm is in-context learning? investigations with linear models.  _arXiv preprint arXiv:2211.15661_ , 2022.
  * Amodei et al. (2016) Amodei, D., Olah, C., Steinhardt, J., Christiano, P., Schulman, J., and Mané, D.  Concrete problems in ai safety.  _arXiv preprint arXiv:1606.06565_ , 2016.
  * Anil et al. (2022) Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B.  Exploring length generalization in large language models.  _arXiv preprint arXiv:2207.04901_ , 2022.
  * Askell et al. (2021) Askell, A., Bai, Y., Chen, A., Drain, D., Ganguli, D., Henighan, T., Jones, A., Joseph, N., Mann, B., DasSarma, N., et al.  A general language assistant as a laboratory for alignment.  _arXiv preprint arXiv:2112.00861_ , 2021.
  * Ba et al. (2016) Ba, J., Hinton, G., Mnih, V., Leibo, J. Z., and Ionescu, C.  Using fast weights to attend to the recent past, 2016.
  * Bai et al. (2023) Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S.  Transformers as statisticians: Provable in-context learning with in-context algorithm selection.  _arXiv preprint arXiv:2306.04637_ , 2023.
  * Bhattamishra et al. (2020) Bhattamishra, S., Ahuja, K., and Goyal, N.  On the ability and limitations of transformers to recognize formal languages.  _arXiv preprint arXiv:2009.11264_ , 2020.
  * Borgeaud et al. (2022) Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E., Millican, K., Van Den Driessche, G. B., Lespiau, J.-B., Damoc, B., Clark, A., et al.  Improving language models by retrieving from trillions of tokens.  In _International conference on machine learning_ , pp. 2206–2240. PMLR, 2022.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Chan et al. (2022) Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, F.  Data distributional properties drive emergent in-context learning in transformers.  _Advances in Neural Information Processing Systems_ , 35:18878–18891, 2022.
  * Cheng et al. (2023) Cheng, X., Chen, Y., and Sra, S.  Transformers implement functional gradient descent to learn non-linear functions in context.  _arXiv preprint arXiv:2312.06528_ , 2023.
  * Chowdhery et al. (2022) Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_ , 2022.
  * Chughtai et al. (2023) Chughtai, B., Chan, L., and Nanda, N.  A toy model of universality: Reverse engineering how networks learn group operations.  _arXiv preprint arXiv:2302.03025_ , 2023.
  * Clark et al. (2022) Clark, K., Guu, K., Chang, M.-W., Pasupat, P., Hinton, G., and Norouzi, M.  Meta-learning fast weight language models.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pp. 9751–9757, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.  URL https://aclanthology.org/2022.emnlp-main.661.
  * Conmy et al. (2023) Conmy, A., Mavor-Parker, A. N., Lynch, A., Heimersheim, S., and Garriga-Alonso, A.  Towards automated circuit discovery for mechanistic interpretability.  _arXiv preprint arXiv:2304.14997_ , 2023.
  * Dai et al. (2022) Dai, D., Sun, Y., Dong, L., Hao, Y., Sui, Z., and Wei, F.  Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers, 2022.
  * Edelman et al. (2022) Edelman, B. L., Goel, S., Kakade, S., and Zhang, C.  Inductive biases and variable creation in self-attention mechanisms.  In _International Conference on Machine Learning_ , pp. 5793–5831. PMLR, 2022.
  * Elhage et al. (2021) Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., et al.  A mathematical framework for transformer circuits.  _Transformer Circuits Thread_ , 2021.
  * Garg et al. (2022) Garg, S., Tsipras, D., Liang, P. S., and Valiant, G.  What can transformers learn in-context? a case study of simple function classes.  _Advances in Neural Information Processing Systems_ , 35:30583–30598, 2022.
  * Giannou et al. (2023) Giannou, A., Rajput, S., yong Sohn, J., Lee, K., Lee, J. D., and Papailiopoulos, D.  Looped transformers as programmable computers, 2023.
  * Gong et al. (2019) Gong, L., He, D., Li, Z., Qin, T., Wang, L., and Liu, T.  Efficient training of bert by progressively stacking.  In _International conference on machine learning_ , pp. 2337–2346. PMLR, 2019.
  * Hahn & Goyal (2023) Hahn, M. and Goyal, N.  A theory of emergent in-context learning as implicit structure induction.  _arXiv preprint arXiv:2303.07971_ , 2023.
  * Han et al. (2023) Han, C., Wang, Z., Zhao, H., and Ji, H.  In-context learning of large language models explained as kernel regression.  _arXiv preprint arXiv:2305.12766_ , 2023.
  * Hendrycks & Gimpel (2016) Hendrycks, D. and Gimpel, K.  Gaussian error linear units (gelus).  _arXiv preprint arXiv:1606.08415_ , 2016.
  * Hinton (2022) Hinton, G.  The forward-forward algorithm: Some preliminary investigations, 2022.
  * Holtzman et al. (2021) Holtzman, A., West, P., Shwartz, V., Choi, Y., and Zettlemoyer, L.  Surface form competition: Why the highest probability answer isn’t always right.  _arXiv preprint arXiv:2104.08315_ , 2021.
  * Hu et al. (2021) Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., and Chen, W.  Lora: Low-rank adaptation of large language models.  _arXiv preprint arXiv:2106.09685_ , 2021.
  * Hu & Liu (2004) Hu, M. and Liu, B.  Mining and summarizing customer reviews.  In _Proceedings of the tenth ACM SIGKDD international conference on Knowledge discovery and data mining_ , pp. 168–177, 2004.
  * Irie et al. (2021) Irie, K., Schlag, I., Csordás, R., and Schmidhuber, J.  Going beyond linear transformers with recurrent fast weight programmers.  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=ot2ORiBqTa1.
  * Izacard & Grave (2020) Izacard, G. and Grave, E.  Leveraging passage retrieval with generative models for open domain question answering.  _arXiv preprint arXiv:2007.01282_ , 2020.
  * Izacard et al. (2023) Izacard, G., Lewis, P., Lomeli, M., Hosseini, L., Petroni, F., Schick, T., Dwivedi-Yu, J., Joulin, A., Riedel, S., and Grave, E.  Atlas: Few-shot learning with retrieval augmented language models.  _Journal of Machine Learning Research_ , 24(251):1–43, 2023.  URL http://jmlr.org/papers/v24/23-0037.html.
  * Jiang (2023) Jiang, H.  A latent space theory for emergent abilities in large language models.  _arXiv preprint arXiv:2304.09960_ , 2023.
  * Kumar et al. (2022) Kumar, A., Shen, R., Bubeck, S., and Gunasekar, S.  How to fine-tune vision models with sgd, 2022.
  * Leike et al. (2018) Leike, J., Krueger, D., Everitt, T., Martic, M., Maini, V., and Legg, S.  Scalable agent alignment via reward modeling: a research direction.  _arXiv preprint arXiv:1811.07871_ , 2018.
  * Lindner et al. (2023) Lindner, D., Kramár, J., Rahtz, M., McGrath, T., and Mikulik, V.  Tracr: Compiled transformers as a laboratory for interpretability.  _arXiv preprint arXiv:2301.05062_ , 2023.
  * Liu et al. (2023) Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., and Zhang, C.  Transformers learn shortcuts to automata.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=De4FYqjFueZ.
  * Liu et al. (2022) Liu, H., Tam, D., Muqeeth, M., Mohta, J., Huang, T., Bansal, M., and Raffel, C. A.  Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning.  _Advances in Neural Information Processing Systems_ , 35:1950–1965, 2022.
  * Mahankali et al. (2023) Mahankali, A., Hashimoto, T. B., and Ma, T.  One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention.  _arXiv preprint arXiv:2307.03576_ , 2023.
  * Malladi et al. (2023) Malladi, S., Gao, T., Nichani, E., Damian, A., Lee, J. D., Chen, D., and Arora, S.  Fine-tuning language models with just forward passes.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=Vota6rFhBQ.
  * Merity et al. (2016) Merity, S., Xiong, C., Bradbury, J., and Socher, R.  Pointer sentinel mixture models.  _arXiv preprint arXiv:1609.07843_ , 2016.
  * Nanda et al. (2023) Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J.  Progress measures for grokking via mechanistic interpretability.  _arXiv preprint arXiv:2301.05217_ , 2023.
  * Olsson et al. (2022) Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., et al.  In-context learning and induction heads.  _arXiv preprint arXiv:2209.11895_ , 2022.
  * Pang & Lee (2004) Pang, B. and Lee, L.  A sentimental education: Sentiment analysis using subjectivity summarization based on minimum cuts.  In _Proceedings of the 42nd Annual Meeting of the Association for Computational Linguistics (ACL-04)_ , pp. 271–278, 2004.
  * Pang & Lee (2005) Pang, B. and Lee, L.  Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales.  In _Proceedings of the 43rd Annual Meeting of the Association for Computational Linguistics (ACL’05)_ , pp. 115–124, 2005.
  * Perez et al. (2021) Perez, J., Barcelo, P., and Marinkovic, J.  Attention is turing-complete.  _Journal of Machine Learning Research_ , 22(75):1–35, 2021.  URL http://jmlr.org/papers/v22/20-302.html.
  * Press et al. (2021) Press, O., Smith, N. A., and Lewis, M.  Train short, test long: Attention with linear biases enables input length extrapolation.  _arXiv preprint arXiv:2108.12409_ , 2021.
  * Pérez et al. (2019) Pérez, J., Marinković, J., and Barceló, P.  On the turing completeness of modern neural network architectures.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=HyGBdo0qFm.
  * Raffel et al. (2020) Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _The Journal of Machine Learning Research_ , 21(1):5485–5551, 2020.
  * Reddi et al. (2023) Reddi, S. J., Miryoosefi, S., Karp, S., Krishnan, S., Kale, S., Kim, S., and Kumar, S.  Efficient training of language models using few-shot learning.  2023\.
  * Saunshi et al. (2020) Saunshi, N., Malladi, S., and Arora, S.  A mathematical exploration of why language models help solve downstream tasks.  _arXiv preprint arXiv:2010.03648_ , 2020.
  * Scao et al. (2022) Scao, T. L., Fan, A., Akiki, C., Pavlick, E., Ilić, S., Hesslow, D., Castagné, R., Luccioni, A. S., Yvon, F., Gallé, M., et al.  Bloom: A 176b-parameter open-access multilingual language model.  _arXiv preprint arXiv:2211.05100_ , 2022.
  * Scellier & Bengio (2017) Scellier, B. and Bengio, Y.  Equilibrium propagation: Bridging the gap between energy-based models and backpropagation.  _Frontiers in computational neuroscience_ , 11:24, 2017.
  * Schlag et al. (2021) Schlag, I., Irie, K., and Schmidhuber, J.  Linear transformers are secretly fast weight memory systems.  _CoRR_ , abs/2102.11174, 2021.  URL https://arxiv.org/abs/2102.11174.
  * Shazeer (2020) Shazeer, N.  Glu variants improve transformer.  _arXiv preprint arXiv:2002.05202_ , 2020.
  * Socher et al. (2013) Socher, R., Perelygin, A., Wu, J., Chuang, J., Manning, C. D., Ng, A. Y., and Potts, C.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _Proceedings of the 2013 conference on empirical methods in natural language processing_ , pp. 1631–1642, 2013.
  * Su et al. (2021) Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., and Liu, Y.  Roformer: Enhanced transformer with rotary position embedding.  _arXiv preprint arXiv:2104.09864_ , 2021.
  * Tay et al. (2022) Tay, Y., Dehghani, M., Tran, V. Q., Garcia, X., Wei, J., Wang, X., Chung, H. W., Bahri, D., Schuster, T., Zheng, S., et al.  Ul2: Unifying language learning paradigms.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Touvron et al. (2023) Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * von Oswald et al. (2023) von Oswald, J., Niklasson, E., Schlegel, M., Kobayashi, S., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., Pascanu, R., et al.  Uncovering mesa-optimization algorithms in transformers.  _arXiv preprint arXiv:2309.05858_ , 2023.
  * Wang et al. (2023a) Wang, B., Ping, W., Xu, P., McAfee, L., Liu, Z., Shoeybi, M., Dong, Y., Kuchaiev, O., Li, B., Xiao, C., Anandkumar, A., and Catanzaro, B.  Shall we pretrain autoregressive language models with retrieval? a comprehensive study.  In Bouamor, H., Pino, J., and Bali, K. (eds.), _Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing_ , pp. 7763–7786, Singapore, December 2023a. Association for Computational Linguistics.  doi: 10.18653/v1/2023.emnlp-main.482.  URL https://aclanthology.org/2023.emnlp-main.482.
  * Wang et al. (2022) Wang, K. R., Variengien, A., Conmy, A., Shlegeris, B., and Steinhardt, J.  Interpretability in the wild: a circuit for indirect object identification in GPT-2 small.  In _NeurIPS ML Safety Workshop_ , 2022.  URL https://openreview.net/forum?id=rvi3Wa768B-.
  * Wang et al. (2023b) Wang, X., Zhu, W., and Wang, W. Y.  Large language models are implicitly topic models: Explaining and finding good demonstrations for in-context learning.  _arXiv preprint arXiv:2301.11916_ , 2023b.
  * Wei et al. (2021) Wei, C., Chen, Y., and Ma, T.  Statistically meaningful approximation: a case study on approximating turing machines with transformers.  _CoRR_ , abs/2107.13163, 2021.  URL https://arxiv.org/abs/2107.13163.
  * Weiss et al. (2021) Weiss, G., Goldberg, Y., and Yahav, E.  Thinking like transformers.  In _International Conference on Machine Learning_ , pp. 11080–11090. PMLR, 2021.
  * Wiebe et al. (2005) Wiebe, J., Wilson, T., and Cardie, C.  Annotating expressions of opinions and emotions in language.  _Language resources and evaluation_ , 39:165–210, 2005.
  * Wies et al. (2023) Wies, N., Levine, Y., and Shashua, A.  The learnability of in-context learning.  _arXiv preprint arXiv:2303.07895_ , 2023.
  * Xie et al. (2022) Xie, S. M., Raghunathan, A., Liang, P., and Ma, T.  An explanation of in-context learning as implicit bayesian inference.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=RdJVFCHjUMI.
  * Yao et al. (2021) Yao, S., Peng, B., Papadimitriou, C., and Narasimhan, K.  Self-attention networks can process bounded hierarchical languages.  _arXiv preprint arXiv:2105.11115_ , 2021.
  * Zhang & Sennrich (2019) Zhang, B. and Sennrich, R.  Root mean square layer normalization.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Zhang et al. (2023a) Zhang, R., Frei, S., and Bartlett, P. L.  Trained transformers learn linear models in-context.  _arXiv preprint arXiv:2306.09927_ , 2023a.
  * Zhang et al. (2015) Zhang, X., Zhao, J., and LeCun, Y.  Character-level convolutional networks for text classification.  _Advances in neural information processing systems_ , 28, 2015.
  * Zhang et al. (2023b) Zhang, Y., Zhang, F., Yang, Z., and Wang, Z.  What and how does in-context learning learn? bayesian model averaging, parameterization, and generalization.  _arXiv preprint arXiv:2305.19420_ , 2023b.
  * Zhou et al. (2023) Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P.  What algorithms can transformers learn? a study in length generalization.  _arXiv preprint arXiv:2310.16028_ , 2023.

###### Contents

  1. 1 Introduction
  2. 2 Design Considerations
    1. 2.1 Input structure
    2. 2.2 Read and write access to auxiliary model
  3. 3 Efficient Forward Propagation
    1. 3.1 Stacking and Sharding
    2. 3.2 Efficient Aggregation
  4. 4 Simulated Gradient
    1. 4.1 First-order approximations
    2. 4.2 Fuzzy backpropagation via stop gradients
  5. 5 Experiments
    1. 5.1 Setting: $N$-step Fine-Tuning
    2. 5.2 Case Study: Language Modeling
    3. 5.3 Case Study: In-Context Learning
  6. 6 Related Work
  7. 7 Discussion
  8. A Additional related works
  9. B Deferred defintions from main paper
  10. C Notations
    1. C.1 Simulating Multiplication from (Akyurek et al., 2022)
  11. D Linear layer
    1. D.1 $H_{\text{sim}}$-split operation
  12. E Self-attention layer
    1. E.1 Proofs of theorems and gradient definitions
  13. F Layer normalization
    1. F.1 Additional definitions
    2. F.2 Proof of theorems and gradient definitions
  14. G Activation layer
    1. G.1 Proofs of theorems
  15. H Language model head
  16. I Parameter sharing
  17. J Additional modules
    1. J.1 Root mean square normalization (RMSnorm)
    2. J.2 Attention variants
    3. J.3 Gated linear units (GLUs)
  18. K Construction of other variants of pre-trained models
  19. L Experiments

## Brief overview of the appendix

In Appendix A, we report few additional related works. In Appendix B, we
present some of the deferred definitions from the main paper. In Appendix C,
we present all the important notations used to present the design of TinT. In
Appendices D, E, F and G, we present the simulation details of all operations
on linear, self-attention, layer normalization, and activation layers
respectively for an auxiliary model. In Appendix H, we present the details for
simulating loss computation with the language model head of the auxiliary
model. In Appendix J, we discuss simulation of additional modules necessary to
simulate transformer variants like LLaMA (Touvron et al., 2023) and BLOOM
(Scao et al., 2022). Finally, in Appendix L, we discuss the deferred
experimental details from the main paper.

##  Appendix A Additional related works

Interpretability: Mechanistic interpretability works reverse-engineer the
algorithms simulated by these models (Elhage et al., 2021; Olsson et al.,
2022; Wang et al., 2022; Nanda et al., 2023; Chughtai et al., 2023; Conmy et
al., 2023). These works study local patterns, e.g. activations and attention
heads, to derive interpretable insights. Other works (Weiss et al., 2021;
Lindner et al., 2023) use declarative programs to algorithmically describe
transformer models. Zhou et al. (2023) use these to explain task-specific
length generalization of transformer models.

Alternative Explanations for ICL: Some works study ICL using a Bayesian
framework. Xie et al. (2022) model pretraining data as a mixture of HMMs and
cast ICL identifying one such component. Hahn & Goyal (2023) later modeled
language as a compositional grammar, and propose ICL as a composition of
operations. (Zhang et al., 2023b; Jiang, 2023; Wang et al., 2023b; Wies et
al., 2023) further strengthen this hypothesis by generalizing the underlying
latent space. On the other hand, careful experiments in Chan et al. (2022)
show that data distributional properties (e.g. Zipf’s law) drive in-context
learning in transformers.

Table 3: Number of parameters of TinT for the forward, backward, and gradient update operations on various modules. For simplicity, we have ignored biases in the following computation. We set $S=4$, i.e. stack $4$ weights in each prefix embedding. We set $H_{\text{sim}}=12$ for OPT-125M and $H_{\text{sim}}=16$ for the other models, $D_{\text{sim}}=4D_{\text{aux}}$ for all the models, and $T_{\text{sim}}=T_{\text{aux}}+K$, with $T_{\text{aux}}=2048$ for opt models, and $K=D_{\text{aux}}/4$. $Q=4Q_{split}+3T_{\text{sim}}D_{\text{sim}}/H_{\text{sim}}$, where $Q_{split}=\frac{1}{H_{\text{sim}}}(D_{\text{sim}})^{2}+H_{\text{sim}}D_{\text{sim}}$, denotes the number of parameters in a TinT Linear Forward module (Section 3).  | Module Size
---|---
Module Name | Forward | Backward | Descent | Total
Linear layer | $Q$ | $Q$ | $Q$ | $3Q$
Layer norms | $Q$ | $Q+2D_{\text{sim}}H_{\text{sim}}$ | $Q$ | $3Q+2D_{\text{sim}}H_{\text{sim}}$
Self-Attention | $2Q$ | $2Q$ | $2Q$ | $6Q$
Activation | $Q_{split}$ | $2D_{\text{sim}}H_{\text{sim}}$ | $0$ | $Q_{split}+2D_{\text{sim}}H_{\text{sim}}$
Self-Attention block | $4Q$ | $4Q+2D_{\text{sim}}H_{\text{sim}}$ | $4Q$ | $12Q+2D_{\text{sim}}H_{\text{sim}}$
Feed-forward block | $3Q+Q_{split}$ | $3Q+4D_{\text{sim}}H_{\text{sim}}$ | $3Q$ | $9Q+4D_{\text{sim}}H_{\text{sim}}$
Transformer block | $7Q+Q_{split}$ | $7Q+6D_{\text{sim}}H_{\text{sim}}$ | $7Q$ | $21Q+6D_{\text{sim}}H_{\text{sim}}+Q_{split}$
Transformer | $7QL+LQ_{split}$ | $(7Q+6D_{\text{sim}}H_{\text{sim}})L$ | $7QL$ | $(21Q+6D_{\text{sim}}H_{\text{sim}}+Q_{split})L$
OPT-125m | 0.4b | 0.4b | 0.4b | 1.2b
OPT-350m | 1.2b | 1.1b | 1.1b | 3.4b
OPT-1.3b | 3.7b | 3.6b | 3.5b | 10.8b
OPT-2.7b | 7.4b | 7.2b | 7.2b | 21.8b

Transfer learning: Our construction uses a pre-trained model to initialize a
larger transformer, which is similar to several other more empirically
oriented works (Gong et al., 2019; Reddi et al., 2023).

##  Appendix B Deferred defintions from main paper

For simplicity of exposition, we showcase the definition on a single head
self-attention layer (multi-head attention is in Definition E.1).

######  Definition B.1 (Auxiliary model softmax self-attention).

A self-attention layer with parameters
$\\{{\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\\}$ takes a sequence
$\\{{\bm{x}}_{t}\\}_{t\leq T_{\text{aux}}}$ and outputs a sequence
$\\{{\bm{y}}_{t}\\}_{t\leq T_{\text{aux}}}$, such that

${\bm{y}}_{t}=\sum_{j}a_{t,j}{\bm{v}}_{j},\qquad\text{with
}a_{t,j}=\mathrm{softmax}({\bm{K}}{\bm{q}}_{t})_{j},\quad{\bm{q}}_{t}={\bm{W}}_{Q}{\bm{x}}_{t},\quad{\bm{k}}_{t}={\bm{W}}_{K}{\bm{x}}_{t},\quad{\bm{v}}_{t}={\bm{W}}_{V}{\bm{x}}_{t},$

for all $t\leq T_{\text{aux}}$, and
${\bm{K}}\in\mathbb{R}^{T_{\text{aux}}\times D_{\text{aux}}}$ defined with
rows $\\{{\bm{k}}_{t}\\}_{t=1}^{T_{\text{aux}}}.$

##  Appendix C Notations

Let $D$ denote the embedding dimension for a token and $T$ denote the length
of an input sequence. $H$ denotes the number of attention heads. With the
exception of contextual embeddings, we use subscripts to indicate if the
quantity is from TinT or from the auxiliary model. For example,
$D_{\text{aux}}$ refers to the embedding dimension and $D_{\text{sim}}$ refers
to the TinT embedding dimension. For contextual embeddings, we use
${\bm{e}}_{t}^{(\ell)}\in\mathbb{R}^{D_{\text{sim}}}$ to denote activations in
TinT and ${\bm{x}}_{t}^{(\ell)}\in\mathbb{R}^{D_{\text{aux}}}$ to denote
activations in the auxiliary model, where $\ell$ is the layer and $t$ is the
sequence position. When convenient, we drop the superscript that represents
the layer index and the subscript that represents the position index. For a
matrix ${\bm{A}}$, ${\bm{a}}_{j}$ refers to its $j$th row, and for any vector
${\bm{b}}$, $b_{j}$ refers to its $j$th element. TinT uses one-hot positional
embeddings
$\\{{{\bm{p}}}^{\textsc{TinT}}_{i}\in\mathbb{R}^{T_{\text{sim}}}\\}_{i\leq
T_{\text{sim}}}$.

We differentiate the parameters of the auxiliary model and TinT by using an
explicit superscript TinT for TinT parameters, for example, the weights of a
linear layer in TinT will be represented by ${{\bm{W}}}^{\textsc{TinT}}$. We
use two operations throughout: $\textsc{Split}_{h}$ and Vectorize. Function
$\textsc{Split}_{h}:\mathbb{R}^{d}\to\mathbb{R}^{h\times\lfloor d/h\rfloor}$
takes an input ${\bm{x}}\in\mathbb{R}^{d}$ and outputs $H$ equal splits of
${\bm{x}}$, for any arbitrary dimension $d$. Function
$\textsc{Vectorize}:\mathbb{R}^{h\times d}\to\mathbb{R}^{dh}$ concatenates the
elements of a sequence $\\{{\bm{x}}_{i}\in\mathbb{R}^{d}\\}_{i\leq h}$ into
one single vector, for any arbitrary $d$ and $h$. Recall that for a matrix
${\bm{A}}$, ${\bm{a}}_{j}$ refers to its $j$th row, and for any vector
${\bm{b}}$, $b_{j}$ refers to its $j$th element. However, at a few places in
the appendix, for typographical reasons, for a matrix ${\bm{A}}$, we have also
used $({\bm{A}})_{j}$ to refer to its $j$th row, and for any vector
${\bm{b}}$, $({\bm{b}})_{j}$ to refer to its $j$th element.

####  TinTAttention Module

We modify the usual attention module to include the position embeddings
$\\{{{\bm{p}}}^{\textsc{TinT}}_{i}\in\mathbb{R}^{T_{\text{sim}}}\\}_{i\leq
T_{\text{sim}}}$. In usual self-attention modules, the query, key, and value
vectors at each position are computed by token-wise linear transformations of
the input embeddings. In TinT’s Attention Module, we perform additional linear
transformations on the position embeddings, using parameters
${\bm{W}}^{p}_{Q},{\bm{W}}^{p}_{K},{\bm{W}}^{p}_{V}$, and decision vectors
$\lambda^{Q},\lambda^{K},\lambda^{V}\in\mathbb{R}^{H_{\text{sim}}}$ decide
whether to add these transformed position vectors to the query, key, and value
vectors of different attention heads. For the following definition, we use
$\hat{e}$ to represent input sequence and $\tilde{e}$ to represent the output
sequence: we introduce these general notations below to avoid confusion with
the notations for token and prefix embeddings for TinTillustrated in Figure 1.

######  Definition C.1 (TinT’s self-attention with $H_{\text{sim}}$ heads).

For parameters
$\\{{{\bm{W}}}^{\textsc{TinT}}_{Q},{{\bm{W}}}^{\textsc{TinT}}_{K},{{\bm{W}}}^{\textsc{TinT}}_{V}\in\mathbb{R}^{D_{\text{sim}}\times
D_{\text{sim}}}\\}$,
$\\{{{\bm{b}}}^{\textsc{TinT}}_{Q},{{\bm{b}}}^{\textsc{TinT}}_{K},{{\bm{b}}}^{\textsc{TinT}}_{V}\in\mathbb{R}^{D_{\text{sim}}}\\}$,
$\\{{\bm{W}}^{p}_{Q},{\bm{W}}^{p}_{K},{\bm{W}}^{p}_{V}\in\mathbb{R}^{T_{\text{sim}}\times
D_{\text{sim}}/H_{\text{sim}}}\\}$ and
$\\{\lambda^{Q},\lambda^{K},\lambda^{V}\in\mathbb{R}^{H_{\text{sim}}}\\}$,
TinT self-attention with $H_{\text{sim}}$ attention heads and a function
$f_{\mathrm{attn}}:\mathbb{R}^{T_{\text{sim}}}\to\mathbb{R}^{T_{\text{sim}}}$
takes a sequence
$\\{\hat{{\bm{e}}}_{t}\in\mathbb{R}^{D_{\text{sim}}}\\}_{t\leq
T_{\text{sim}}}$ as input and outputs
$\\{\tilde{{\bm{e}}}_{t}\in\mathbb{R}^{D_{\text{sim}}}\\}_{t\leq
T_{\text{sim}}}$, with

$\displaystyle\tilde{{\bm{e}}}_{t}=\textsc{Vectorize}(\\{\sum_{j\leq
T_{\text{sim}}}a^{h}_{t,j}\tilde{{\bm{v}}}^{h}_{j})_{h}\\}_{h\leq
H_{\text{sim}}}),\text{ with
}a^{h}_{t,j}=f_{\mathrm{attn}}(\tilde{{\bm{K}}}^{h}\tilde{{\bm{q}}}_{t}^{h})_{j}$
$\displaystyle\tilde{{\bm{q}}}^{h}_{t}=\textsc{Split}_{H}({\bm{q}}_{t})_{h}+\lambda^{Q}_{h}{\bm{W}}^{p}_{Q}{{\bm{p}}}^{\textsc{TinT}}_{t};\quad\tilde{{\bm{k}}}^{h}_{t}=\textsc{Split}_{H}({\bm{k}}_{t})_{h}+\lambda^{K}_{h}{\bm{W}}^{p}_{K}{{\bm{p}}}^{\textsc{TinT}}_{t};$
$\displaystyle\tilde{{\bm{v}}}^{h}_{t}=\textsc{Split}_{H}({\bm{v}}_{t})_{h}+\lambda^{V}_{h}{\bm{W}}^{p}_{v}{{\bm{p}}}^{\textsc{TinT}}_{t}.$

Here, ${\bm{q}}_{t}$, ${\bm{k}}_{t}$, ${\bm{v}}_{t}$ denote the query, key,
and value vectors at each position $t$, computed as
${{\bm{W}}}^{\textsc{TinT}}_{Q}\hat{{\bm{e}}}_{t}+{{\bm{b}}}^{\textsc{TinT}}_{Q}$,
${{\bm{W}}}^{\textsc{TinT}}_{K}\hat{{\bm{e}}}_{t}+{{\bm{b}}}^{\textsc{TinT}}_{K}$,
and
${{\bm{W}}}^{\textsc{TinT}}_{V}\hat{{\bm{e}}}_{t}+{{\bm{b}}}^{\textsc{TinT}}_{V}$
respectively. $\tilde{{\bm{K}}}^{h}\in\mathbb{R}^{T_{\text{sim}}\times
D_{\text{sim}}/H_{\text{sim}}}$ is defined with its rows as
$\\{\tilde{{\bm{k}}}^{h}_{t}\\}_{t\leq T_{\text{sim}}}$ for all $h\leq
H_{\text{sim}}$.

$f_{\mathrm{attn}}$ can be either linear or softmax function.

#### Bounded parameters and input sequence:

We define a linear self-attention layer to be $B_{w}$-bounded, if the
$\ell_{2}$ norms of all the parameters are bounded by $B_{w}$. Going by
Definition C.1, this implies

$\max\\{\left\|{{\bm{W}}}^{\textsc{TinT}}_{Q}\right\|_{2},\left\|{{\bm{W}}}^{\textsc{TinT}}_{K}\right\|_{2},\left\|{{\bm{W}}}^{\textsc{TinT}}_{V}\right\|_{2}\\}\leq
B_{w},\quad\max\\{\left\|{{\bm{b}}}^{\textsc{TinT}}_{Q}\right\|_{2},\left\|{{\bm{b}}}^{\textsc{TinT}}_{K}\right\|_{2},\left\|{{\bm{b}}}^{\textsc{TinT}}_{V}\right\|_{2}\\}\leq
B_{w}$
$\max\\{\left\|{\bm{W}}_{Q}^{p}\right\|_{2},\left\|{\bm{W}}_{K}^{p}\right\|_{2},\left\|{\bm{W}}_{V}^{p}\right\|_{2}\\}\leq
B_{w},\quad\max\\{\left\|\lambda^{Q}\right\|_{2},\left\|\lambda^{K}\right\|_{2},\left\|\lambda^{V}\right\|_{2}\\}\leq
B_{w}.$

Furthermore, we define an input sequence $\\{\hat{{\bm{e}}}_{t}\\}_{t\leq
T_{\text{sim}}}$ to $B_{x}$-bounded, if
$\left\|\hat{{\bm{e}}}_{t}\right\|_{2}\leq B_{x}$ for all $t$.

Recall from the main paper (Section 3), we used Linear TinT Self-Attention
layer to represent the linear operations of the auxiliary model. In the
following theorem, we show that a linear attention layer can be represented as
a softmax attention layer that uses an additional attention head and an extra
token ${\bm{u}}$, followed by a linear layer. Therefore, replacing softmax
attention with linear attention does not deviate too far from the canonical
transformer. We use the Linear TinT Self-Attention layers in several places
throughout the model.

######  Theorem C.2.

For any $B_{w}>0$, consider a $B_{w}$-bounded linear self-attention layer that
returns
$\\{\tilde{{\bm{e}}}^{linear}_{t}\in\mathbb{R}^{D}_{\text{sim}}\\}_{t\leq
T_{\text{sim}}}$ on any input
$\\{\hat{{\bm{e}}}_{t}\in\mathbb{R}^{D}_{\text{sim}}\\}_{t\leq
T_{\text{sim}}}$. Consider a softmax self-attention layer with
$2H_{\text{sim}}$ attention heads and an additional token
${\bm{u}}\in\mathbb{R}^{2D_{\text{sim}}}$ such that for any $B_{x}$-bounded
input $\\{\hat{{\bm{e}}}_{t}\\}_{t\leq T_{\text{sim}}}$, it takes a modified
input sequence
$\\{\bar{{\bm{e}}}_{1},\cdots,\bar{{\bm{e}}}_{T_{\text{sim}}},{\bm{u}}\\}$,
and returns
$\\{\tilde{{\bm{e}}}^{softmax}_{t}\in\mathbb{R}^{2D_{\text{sim}}}\\}_{t\leq
T_{\text{sim}}}$. Each modified input token
$\bar{{\bm{e}}}_{t}\in\mathbb{R}^{2D_{\text{sim}}}$ is obtained by
concatenating additional $0$s to $\hat{{\bm{e}}}_{t}$. Then, for any
$B_{x}>0$, and
$\epsilon\leq\mathcal{O}(T_{\text{sim}}^{-2}B_{w}^{-5}B_{x}^{-5})$, there
exists ${\bm{W}}_{O}\in\mathbb{R}^{D_{\text{sim}}\times 2D_{\text{sim}}}$ and
such a softmax self-attention layer such that

$\displaystyle\left\|{\bm{W}}_{O}\tilde{{\bm{e}}}^{softmax}_{t}-\tilde{{\bm{e}}}^{linear}_{t}\right\|_{2}\leq\mathcal{O}(\sqrt{\epsilon}),$

for all $t\leq T_{\text{sim}}$.

###### Proof.

Consider an input sequence $\\{{\bm{x}}_{t}\\}_{t\leq T_{\text{sim}}}$. Let
the attention scores of any linear head $h\leq H_{\text{sim}}$ in the linear
attention layer be given by $\\{a^{h}_{t,j}\\}_{j\leq T_{\text{sim}}},$ at any
given position $t$. Additionally, let the value vectors for the linear
attention be given by ${\bm{v}}_{t}$. To repeat our self-attention definition,
the output of the attention layer at any position $t$ is given by
$\textsc{Vectorize}(\\{\tilde{{\bm{e}}}^{linear,h}_{t}\\}_{h\leq
H_{\text{sim}}})$, where

$\displaystyle\tilde{{\bm{e}}}^{linear,h}_{t}=\sum_{j\leq
T_{\text{sim}}}a^{h}_{t,j}{\bm{v}}^{h}_{j}.$

Under our assumption, $B_{w}$ denotes the maximum $\ell_{2}$ norm of all the
parameters in the linear self-attention layer and $B_{x}$ the maximum
$\ell_{2}$ norm in the input sequence, i.e. $\max_{t\leq
T_{\text{sim}}}\left\|{\bm{x}}_{t}\right\|_{2}\leq B_{x}$. With a simple
application of Cauchy-Schwartz inequality, we can show that $\max_{j\leq
T_{\text{sim}}}|a^{h}_{t,j}|\leq\mathcal{O}(B_{w}^{2}B_{x}^{2}),$ and
$\max_{t\leq
T_{\text{sim}}}\left\|{\bm{v}}^{h}_{t}\right\|_{2}\leq\mathcal{O}(B_{w}B_{x}).$

For
$\epsilon\leq\mathcal{O}(T_{\text{sim}}^{-10/9}B_{w}^{-40/9}B_{x}^{-40/9})$,
we can then use Lemma C.3 to represent for each $t,j\leq T_{\text{sim}}$,

$\displaystyle a^{h}_{t,j}$ $\displaystyle=\frac{\epsilon^{-3}e^{\epsilon
a_{t,j}}}{\sum_{t^{\prime}\leq T_{\text{sim}}}e^{\epsilon
a^{h}_{t,t^{\prime}}}+e^{-2\log\epsilon}}-\epsilon^{-1}+\mathcal{O}\left(\epsilon(T_{\text{sim}}+a^{h}_{t,j})\right)$
$\displaystyle:=\epsilon^{-3}\mathrm{softmax}\left(\\{\epsilon
a^{h}_{t,1},\epsilon a^{h}_{t,2},\cdots,\epsilon
a^{h}_{t,T_{\text{sim}}},-2\log\epsilon\\}\right)_{j}-\epsilon^{-1}+\mathcal{O}\left(\epsilon^{0.9}\right).$

#### Softmax attention construction:

We define ${\bm{u}}$, and the query and key parameters of the softmax
attention layer such that for the first $H_{\text{sim}}$ attention heads, the
query-key dot products for all the attention heads between any pairs
$\\{(\bar{{\bm{e}}}_{t},\bar{{\bm{e}}}_{j})\\}_{t,j\leq T_{\text{sim}}}$ is
given by $\\{\epsilon a^{h}_{t,j}\\}_{h\leq H_{\text{sim}}}$, while being
$-2\log\epsilon$ between ${\bm{u}}$ and any token $\bar{{\bm{e}}}_{t}$, with
$t\leq T_{\text{sim}}$. For the rest of $H_{\text{sim}}$ attention heads, the
attention scores are uniformly distributed across all pairs of tokens
(attention score between any pair of tokens is given by
$\frac{1}{T_{\text{sim}}+1}$).

We set the value parameters of the softmax attention layer such that at any
position $t\leq T_{\text{sim}}$, the value vector is given by
$\textsc{Vectorize}(\\{\epsilon^{-3}{\bm{v}}_{t},{\bm{v}}_{t}\\}).$ The value
vector returned for ${\bm{u}}$ contains all $0$s.

#### Softmax attention computation:

Consider an attention head $h\leq H_{\text{sim}}$ in the softmax attention
layer now. The output of the attention head at any position $t\leq
T_{\text{sim}}$ is given by

$\displaystyle\tilde{{\bm{e}}}^{softmax,h}_{t}$ $\displaystyle=\sum_{j\leq
T_{\text{sim}}}\mathrm{softmax}\left(\\{\epsilon a^{h}_{t,1},\epsilon
a^{h}_{t,2},\cdots,\epsilon
a^{h}_{t,T_{\text{sim}}},-2\log\epsilon\\}\right)_{j}\epsilon^{-3}{\bm{v}}_{j}^{h}$
$\displaystyle=\sum_{j\leq
T_{\text{sim}}}\left(a_{t,j}^{h}+\epsilon^{-1}+\mathcal{O}(\epsilon^{0.9})\right){\bm{v}}^{h}_{j}.$

This has an additional $\sum_{j\leq
T_{\text{sim}}}\left(\epsilon^{-1}+\mathcal{O}(\epsilon^{0.9})\right){\bm{v}}^{h}_{j}$,
compared to $\tilde{{\bm{e}}}^{linear,h}_{t}$. However, consider the output of
the attention head $H_{\text{sim}}+h$ at the same position:

$\displaystyle\tilde{{\bm{e}}}^{softmax,H_{\text{sim}}+h}_{t}=\frac{1}{T_{\text{sim}}+1}\sum_{j\leq
T_{\text{sim}}}{\bm{v}}^{h}_{j}.$

Hence, we can use the output matrix ${\bm{W}}_{O}$ to get
$\tilde{{\bm{e}}}^{softmax,h}_{t}-\frac{T_{\text{sim}}+1}{\epsilon}\tilde{{\bm{e}}}^{softmax,H_{\text{sim}}+h}_{t}=\sum_{j\leq
T_{\text{sim}}}\left(a_{t,j}^{h}+\mathcal{O}(\epsilon^{0.9})\right){\bm{v}}^{h}_{j}.$
The additional term $\mathcal{O}(\epsilon^{0.9})\sum_{j\leq
T_{\text{sim}}}{\bm{v}}^{h}_{j}$ can be further shown to be
$\mathcal{O}(\epsilon^{0.5})$ small with the assumed bound of $\epsilon$,
since each ${\bm{v}}^{h}_{j}$ is atmost $\mathcal{O}(B_{w}B_{x})$ in
$\ell_{2}$ norm with a Cauchy Schwartz inequality. ∎

######  Lemma C.3.

For $\epsilon>0$, $B>0$, and a sequence $\\{a_{1},a_{2},\cdots,a_{T}\\}$ with
each $a_{i}\in\mathbb{R}$ and $\left|a_{i}\right|\leq B$, the following holds
true for all $i\leq T$,

$\displaystyle\frac{\epsilon^{-3}e^{\epsilon a_{i}}}{\sum_{t^{\prime}\leq
T}e^{\epsilon
a_{t^{\prime}}}+e^{-2\log\epsilon}}=a_{i}+\frac{1}{\epsilon}+\mathcal{O}\left(\epsilon^{0.9}\right),$

provided $\epsilon\leq\mathcal{O}(T^{-10/9}B^{-20/9}).$

###### Proof.

We will use the following first-order Taylor expansions:

$\displaystyle e^{x}=1+x+\mathcal{O}(x^{2}).$ (4)
$\displaystyle\frac{1}{1+x}=1-\mathcal{O}(x).$ (5)

Hence, for any $x\ll 1$, $x\approx e^{x}-1.$

Simplifying the L.H.S. of the desired bound, we have

$\displaystyle\frac{\epsilon^{-3}e^{\epsilon a_{i}}}{\sum_{t^{\prime}\leq
T}e^{\epsilon a_{t^{\prime}}}+e^{-2\log\epsilon}}$
$\displaystyle=\frac{\epsilon^{-3}(1+\epsilon
a_{i}+\mathcal{O}(\epsilon^{2}a_{i}^{2}))}{\sum_{t^{\prime}\leq T}(1+\epsilon
a_{t^{\prime}}+\mathcal{O}(\epsilon^{2}a_{t^{\prime}}^{2}))+e^{-2\log\epsilon}}$
(6) $\displaystyle=\frac{\epsilon^{-1}+a_{i}+\mathcal{O}(\epsilon
a_{i}^{2})}{\sum_{t^{\prime}\leq
T}(\epsilon^{2}+\epsilon^{3}a_{t^{\prime}}+\mathcal{O}(\epsilon^{4}a_{t^{\prime}}^{2}))+1}$
(7) $\displaystyle=\left(\epsilon^{-1}+a_{i}+\mathcal{O}(\epsilon
a_{i}^{2})\right)\left(1+\mathcal{O}(\epsilon^{2}T)\right)$ (8)
$\displaystyle=\epsilon^{-1}+a_{i}+\mathcal{O}(\epsilon
T+a_{i}^{2}T\epsilon^{2}+a_{i}^{2}T\epsilon^{3}+\epsilon
a_{i}^{2})=\epsilon^{-1}+a_{i}+\mathcal{O}(\epsilon^{0.9}).$

We used taylor expansion of exponential function( Equation 4 ) in Equation 6
to get Equation 7, and taylor expansion of inverse function(Equation 5) to get
Equation 8 from Equation 7. Furthermore, with the lower bound assumption on
$\epsilon$, $\sum_{t^{\prime}\leq
T}(\epsilon^{2}+\epsilon^{3}a_{t^{\prime}}+\mathcal{O}(\epsilon^{4}a_{t^{\prime}}^{2}))$
can be shown to be atmost $3\epsilon^{2}T$, which amounts to
$\mathcal{O}(\epsilon^{2}T)$ error in Equation 8. The final error bound has
again been simplified using the lower bound assumption on $\epsilon$. ∎

###  C.1 Simulating Multiplication from (Akyurek et al., 2022)

We refer to the multiplication strategy of (Akyurek et al., 2022) at various
places.

######  Lemma C.4.

[Lemma 4 in (Akyurek et al., 2022)] The $GeLU$ (Hendrycks & Gimpel, 2016)
nonlinearity can be used to perform multiplication: specifically,

$\displaystyle\sqrt{\pi/2}(GeLU(x+y)-GeLU(y))=xy+\mathcal{O}(x^{3}y^{3}).$

Thus, to represent an element-wise product or a dot product between two sub-
vectors in a token embedding, we can use a MLP with a $GeLU$ activation.

##  Appendix D Linear layer

In the main paper, we defined the linear layer without the bias term for
simplicity (Definition 3.1). In this section, we will redefine the linear
layer with the bias term and present a comprehensive construction of the
Linear Forward module.

######  Definition D.1 (Linear layer).

For a weight ${\bm{W}}\in\mathbb{R}^{D_{\text{aux}}\times D_{\text{aux}}}$ and
bias ${\bm{b}}\in\mathbb{R}^{D_{\text{aux}}}$, a linear layer takes
${\bm{x}}\in\mathbb{R}^{D_{\text{aux}}}$ as input and outputs
${\bm{y}}={\bm{W}}{\bm{x}}+{\bm{b}}$.

In the discussions below, we consider a linear layer in the auxiliary model
with parameters $\\{{\bm{W}},{\bm{b}}\\}$ that takes in input sequence
${\bm{x}}_{1},\cdots,{\bm{x}}_{T_{\text{aux}}}$ and outputs
${\bm{y}}_{1},\cdots,{\bm{y}}_{T_{\text{aux}}}$, with
${\bm{y}}_{t}={\bm{W}}{\bm{x}}_{t}+{\bm{b}}$ for each $t\leq T_{\text{aux}}$.
Since this involves a token-wise operation, we will present our constructed
modules with a general token position $t$ and the prefix tokens
$\\{{\bm{v}}_{j}\\}.$

####  TinT Linear Forward module

Continuing our discussion from Section 3, we represent $S$ stacked rows of
${\bm{W}}$ as a prefix embedding. In addition, we store the bias ${\bm{b}}$ in
the first prefix embedding (${\bm{v}}_{1}$).

Using a set of $S^{\prime}$ unique attention heads in a TinT attention module
(Definition C.1), we copy the bias ${\bm{b}}$ to respective token embeddings
and use a TinT linear layer to add the biases to the final output.

#### Auxiliary’s backpropagation through linear layer

For a linear layer as defined in Definition D.1, the linear backpropagation
layer takes in the loss gradient w.r.t. output ($\partial_{{\bm{y}}}$) and
computes the loss gradient w.r.t. input ($\partial_{{\bm{x}}}$).

######  Definition D.2 (Linear backpropagation ).

For a weight ${\bm{W}}\in\mathbb{R}^{D_{\text{aux}}\times D_{\text{aux}}}$ ,
the linear backpropagation layer takes
$\partial_{{\bm{y}}}\in\mathbb{R}^{D_{\text{aux}}}$ as input and outputs
$\partial_{{\bm{x}}}={\bm{W}}^{\top}\partial_{{\bm{y}}}$.

####  TinT Linear backpropagation module

This module will aim to simulate the auxiliary’s linear backpropagation. The
input embedding ${\bm{e}}_{t}$ to this module will contain the gradient of the
loss w.r.t. ${\bm{y}}_{t}$, i.e. $\partial_{{\bm{y}}_{t}}$. As given in
Definition D.2, this module will output the gradient of the loss w.r.t.
${\bm{x}}_{t}$, given by
$\partial_{{\bm{x}}_{t}}={\bm{W}}^{\top}\partial_{{\bm{y}}_{t}}.$

We first use the residual connection to copy the prefix embeddings
$\\{{\bm{v}}_{j}\\}$ (i.e., the rows of ${\bm{W}}$) from the forward
propagation module. A straightforward construction would be to use the Linear
Forward module but with the columns of ${\bm{W}}$ stored in the prefix tokens,
thereby simulating multiplication with ${\bm{W}}^{\top}$. However, such a
construction requires applying attention to the prefix tokens, which increases
the size of the construction substantially.

We instead perform the operation more efficiently by splitting it across
attention heads. In particular, once we view the operation as
$\partial_{{\bm{x}}_{t}}=\sum_{i}\left(\partial_{{\bm{y}}_{t}}\right)_{i}{\bm{w}}_{i}$,
we can see that the attention score between the current token and the prefix
token containing ${\bm{w}}_{i}$ must be
$\left(\partial_{{\bm{y}}_{t}}\right)_{i}$. Using value vectors as rows of
${\bm{W}}$ returns the desired output. Similar to the Linear Forward module,
we shard the weights into $S^{\prime}$ parts to parallelize across more
attention heads. Please see Figure 4.

Figure 4: TinT simulates the backward pass of a linear layer as a $H$-head
attention layer ($H=6$ pictured), with the gradient of the loss w.r.t. linear
layer output ($\partial_{{\bm{y}}_{t}}$) as the query, the positional one-hot
vector of prefix embeddings as the key, and the parameters of the auxiliary
model stored in the prefix embeddings as the value. Similar to the Linear
Forward module (Figure 2), we distribute the dot product computations across
all attention heads by sharding the vectors into $S^{\prime}$ ($S^{\prime}=3$
here) parts. We omitted the identical transformation for query, and value
matrices, and permutation-based transformation for key matrix for illustration
purposes. Figure 5: TinT computes the parameter gradients for a linear layer
as a $H$-head attention layer ($H=6$ pictured), with the gradient of the loss
w.r.t. linear layer output ($\partial_{{\bm{y}}_{t}}$) as the query, the
positional one-hot vector of prefix embeddings as the key, and the input to
the linear layer (${\bm{x}}_{t}$) as the value. The auxiliary model parameters
in the prefix embeddings are then updated using a residual connection. Similar
to the Linear Forward module (Figure 2), we distribute the dot product
computations across all attention heads, by sharding the vectors into
$S^{\prime}$ ($S^{\prime}=3$ here) parts. We omitted the identical
transformation for query, and value matrices, and permutation-based
transformation for key matrix for simplicity.

#### Auxiliary’s linear descent update

Finally, the linear descent layer updates the weight and the bias parameters
using a batch of inputs $\\{{\bm{x}}_{t}\\}_{t\leq{T_{\text{aux}}}}$ and the
loss gradient w.r.t. the corresponding outputs
$\\{\partial_{{\bm{y}}_{t}}\\}_{t\leq{T_{\text{aux}}}}$.

######  Definition D.3 (Linear descent).

For a weight ${\bm{W}}\in\mathbb{R}^{D_{\text{aux}}\times D_{\text{aux}}}$ and
a bias ${\bm{b}}\in\mathbb{R}^{D_{\text{aux}}}$, the linear descent layer
takes in a batch of inputs
$\\{{\bm{x}}_{t}\in\mathbb{R}^{D}_{\text{aux}}\\}_{t\leq{T_{\text{aux}}}}$ and
gradients
$\\{\partial_{{\bm{y}}_{t}}\in\mathbb{R}^{D}_{\text{aux}}\\}_{t\leq{T_{\text{aux}}}}$
and updates the parameters as follows:

$\displaystyle{\bm{W}}\leftarrow{\bm{W}}-\eta\sum_{t\leq{T_{\text{aux}}}}\partial_{{\bm{y}}_{t}}{\bm{x}}_{t}^{\top};\quad\quad{\bm{b}}\leftarrow{\bm{b}}-\eta\sum_{t\leq{T_{\text{aux}}}}\partial_{{\bm{y}}_{t}}.$

####  TinT Linear descent module

The input embedding ${\bm{e}}_{t}$ to this module will contain the gradient of
the loss w.r.t. ${\bm{y}}_{t}$, i.e. $\partial_{{\bm{y}}_{t}}$.

As in the Linear backpropagation module, the prefix tokens
$\\{{\bm{v}}_{j}\\}$ will contain the rows of ${\bm{W}}$ and ${\bm{b}}$, which
have been copied from the Linear forward module using residual connections.
Since, in addition to the gradients, we also require the input to the linear
layer, we will use residual connections to copy the input $\\{{\bm{x}}_{t}\\}$
to their respective embeddings $\\{{\bm{e}}_{t}\\},$ from the Linear Forward
module. As given in Definition D.3, this module will update ${\bm{W}}$ and
${\bm{b}}$ using the gradient descent rule.

Focusing on ${\bm{w}}_{i}$, the descent update is given by
${\bm{w}}_{i}\leftarrow{\bm{w}}_{i}-\eta\sum_{t}\left(\partial_{{\bm{y}}_{t}}\right)_{i}{\bm{x}}_{t}$.
For the prefix token ${\bm{v}}_{j}$ that contains ${\bm{w}}_{i}$, the update
term $-\eta\sum_{t}\left(\partial_{{\bm{y}}_{t}}\right)_{i}{\bm{x}}_{t}$ can
be expressed with an attention head that represents the attention between the
prefix token ${\bm{v}}_{j}$ and any token ${\bm{e}}_{t}$ with score
$\left(\partial_{{\bm{y}}_{t}}\right)_{i}$ and value $-\eta{\bm{x}}_{t}$. The
residual connection can then be used to update the weights ${\bm{w}}_{i}$ in
${\bm{v}}_{j}$.

For the bias ${\bm{b}}$, the descent update is give by
${\bm{b}}\leftarrow{\bm{b}}-\eta\sum_{t}\partial_{{\bm{y}}_{t}}$. With
${\bm{b}}$ present in ${\bm{v}}_{1}$, we use one attention head to represent
the attention score between prefix token ${\bm{v}}_{1}$ and any token
${\bm{e}}_{t}$ as $1$, with the value being $-\eta\partial_{{\bm{y}}_{t}}.$
The residual connection can then be used to update the weights ${\bm{b}}$ in
${\bm{v}}_{1}$.

The above process can be further parallelized across multiple attention heads,
by sharding each weight computation into $S^{\prime}$ parts. Please see Figure
5.

###  D.1 $H_{\text{sim}}$-split operation

We leverage local structure within the linear operations of TinT to make the
construction smaller. We build two $H_{\text{sim}}$-split operations to
replace all the linear operations. We use $d_{\text{sim}}$ to denote
$D_{\text{sim}}/H_{\text{sim}}$ in the following definitions.

######  Definition D.4 (Split-wise $H_{\text{sim}}$-split Linear operation).

For weight and bias parameters
${{\bm{W}}}^{\textsc{TinT}}\in\mathbb{R}^{H_{\text{sim}}\times
d_{\text{sim}}\times
d_{\text{sim}}},{{\bm{B}}}^{\textsc{TinT}}\in\mathbb{R}^{H_{\text{sim}}\times
d_{\text{sim}}}$ , this layer takes in input
${\bm{e}}\in\mathbb{R}^{D_{\text{sim}}}$ and returns
$\tilde{{\bm{e}}}=\textsc{Vectorize}(\tilde{{\bm{S}}}+{{\bm{B}}}^{\textsc{TinT}})$,
with $\tilde{{\bm{S}}}\in\mathbb{R}^{H_{\text{sim}}\times d_{\text{sim}}}$
defined with rows
$\\{{{\bm{W}}}^{\textsc{TinT}}_{h}\textsc{Split}_{H_{\text{sim}}}({\bm{e}})_{h}\\}_{h\leq
H_{\text{sim}}}$.

######  Definition D.5 (Dimension-wise $H_{\text{sim}}$-split Linear
operation).

For weight and bias parameters
${{\bm{W}}}^{\textsc{TinT}}\in\mathbb{R}^{d_{\text{sim}}\times
H_{\text{sim}}\times
H_{\text{sim}}},{{\bm{B}}}^{\textsc{TinT}}\in\mathbb{R}^{d_{\text{sim}}\times
H_{\text{sim}}}$ , this layer takes in input
${\bm{e}}\in\mathbb{R}^{D_{\text{sim}}}$, defines
${\bm{S}}\in\mathbb{R}^{d_{\text{sim}}\times H_{\text{sim}}}$ with columns
$\\{\textsc{Split}_{H_{\text{sim}}}({\bm{e}})_{h}\\}_{h\leq H_{\text{sim}}}$,
and returns
$\tilde{{\bm{e}}}=\textsc{Vectorize}((\tilde{{\bm{S}}}+{{\bm{B}}}^{\textsc{TinT}})^{\top})$,
where $\tilde{{\bm{S}}}\in\mathbb{R}^{d_{\text{sim}}\times H_{\text{sim}}}$ is
defined with rows
$\\{{{\bm{W}}}^{\textsc{TinT}}_{d}{{\bm{s}}}^{\textsc{TinT}}_{d}\\}_{d\leq
d_{\text{sim}}}$.

We find that we can replace all the linear operations with a splitwise
$H_{\text{sim}}$-split Linear operation followed by a dimensionwise
$H_{\text{sim}}$-split Linear operation, and an additional splitwise
$H_{\text{sim}}$-split Linear operation, if necessary. A linear operation on
$D_{\text{sim}}$-dimensional space involves $D_{\text{sim}}^{2}$ parameters,
while its replacement requires
$D_{\text{sim}}^{2}/H_{\text{sim}}+2D_{\text{sim}}H_{\text{sim}}$ parameters,
effectively reducing the total number of necessary parameters by
$H_{\text{sim}}$.

We motivate the $H_{\text{sim}}$-split linear operations with an example. We
consider the Linear Forward module in Figure 2 for simulating a linear
operation with parameters ${\bm{W}}\in\mathbb{R}^{D_{\text{aux}}\times
D_{\text{aux}}}$ and no biases. For simplicity of presentation, we assume
$D_{\text{aux}}$ is divisible by $4$. We stack $2$ rows of weights per prefix
embedding. We distribute the dot-product computation across the
$H_{\text{sim}}=6$ attention heads, by sharding each weight into $3$ parts.
Since we require to have enough space to store all the sharded computation
from the linear attention heads, we require $D_{\text{sim}}=3D_{\text{aux}}$
(we get $3$ values for each of the $D_{\text{aux}}$ weights in ${\bm{W}}$).
For presentation, for a given vector ${\bm{v}}\in\mathbb{R}^{D_{\text{aux}}}$,
we represent $\textsc{Split}_{3}({\bm{v}})_{i}$ by ${\bm{v}}^{i}$ for all
$1\leq i\leq 3$.

Now, consider the final linear operation responsible for combining the output
of the attention heads. The output, after the linear operation, should contain
${\bm{W}}{\bm{x}}_{t}$ in the first $D_{\text{aux}}$ coordinates. At any
position $t$, if we stack the output of the linear attention heads as rows of
a matrix ${\bm{S}}_{t}\in\mathbb{R}^{H_{\text{sim}}\times
D_{\text{sim}}/H_{\text{sim}}}$ we get

$\displaystyle{\bm{S}}_{t}=\begin{bmatrix}\langle{\bm{w}}_{1}^{1},{\bm{x}}_{t}^{1}\rangle&\langle{\bm{w}}_{3}^{1},{\bm{x}}_{t}^{1}\rangle&\langle{\bm{w}}_{5}^{1},{\bm{x}}_{t}^{1}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}-1}^{1},{\bm{x}}_{t}^{1}\rangle\\\
\langle{\bm{w}}_{1}^{2},{\bm{x}}_{t}^{2}\rangle&\langle{\bm{w}}_{3}^{2},{\bm{x}}_{t}^{2}\rangle&\langle{\bm{w}}_{5}^{2},{\bm{x}}_{t}^{2}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}-1}^{2},{\bm{x}}_{t}^{2}\rangle\\\
\langle{\bm{w}}_{1}^{3},{\bm{x}}_{t}^{3}\rangle&\langle{\bm{w}}_{3}^{3},{\bm{x}}_{t}^{3}\rangle&\langle{\bm{w}}_{5}^{3},{\bm{x}}_{t}^{3}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}-1}^{3},{\bm{x}}_{t}^{3}\rangle\\\
\langle{\bm{w}}_{2}^{1},{\bm{x}}_{t}^{1}\rangle&\langle{\bm{w}}_{4}^{1},{\bm{x}}_{t}^{1}\rangle&\langle{\bm{w}}_{6}^{1},{\bm{x}}_{t}^{1}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}}^{1},{\bm{x}}_{t}^{1}\rangle\\\
\langle{\bm{w}}_{2}^{2},{\bm{x}}_{t}^{2}\rangle&\langle{\bm{w}}_{4}^{2},{\bm{x}}_{t}^{2}\rangle&\langle{\bm{w}}_{6}^{2},{\bm{x}}_{t}^{2}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}}^{2},{\bm{x}}_{t}^{2}\rangle\\\
\langle{\bm{w}}_{2}^{3},{\bm{x}}_{t}^{3}\rangle&\langle{\bm{w}}_{4}^{3},{\bm{x}}_{t}^{3}\rangle&\langle{\bm{w}}_{6}^{3},{\bm{x}}_{t}^{3}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}}^{3},{\bm{x}}_{t}^{3}\rangle\end{bmatrix}$

Note that for each $j\leq D_{\text{aux}}$, we have
$\langle{\bm{w}}_{j},{\bm{x}}_{t}\rangle=\sum_{i=1}^{3}\langle{\bm{w}}_{j}^{i},{\bm{x}}_{t}^{i}\rangle$.
Thus, with a column-wise linear operation on ${\bm{S}}_{t}$, we can sum the
relevant elements in each column to get

$\displaystyle{\bm{S}}^{col}_{t}=$
$\displaystyle\begin{bmatrix}\langle{\bm{w}}_{1},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{3},{\bm{x}}_{t}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}/2-1},{\bm{x}}_{t}\rangle&0&0&\cdots&0\\\
\langle{\bm{w}}_{2},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{4},{\bm{x}}_{t}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}/2},{\bm{x}}_{t}\rangle&0&0&\cdots&0\\\
0&0&\cdots&0&\langle{\bm{w}}_{D_{\text{aux}}/2+1},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}/2+3},{\bm{x}}_{t}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}-1},{\bm{x}}_{t}\rangle\\\
0&0&\cdots&0&\langle{\bm{w}}_{D_{\text{aux}}/2+2},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}/2+4},{\bm{x}}_{t}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}},{\bm{x}}_{t}\rangle\\\
0&0&\cdots&0&0&0&\cdots&0\\\ 0&0&\cdots&0&0&0&\cdots&0\end{bmatrix}$

A row-wise linear operation on ${\bm{S}}^{col}_{t}$ can space out the non-zero
elements in the matrix and give us

$\displaystyle{\bm{S}}^{row}_{t}=$
$\displaystyle\begin{bmatrix}\langle{\bm{w}}_{1},{\bm{x}}_{t}\rangle&0&\langle{\bm{w}}_{3},{\bm{x}}_{t}\rangle&0&\cdots&\langle{\bm{w}}_{D_{\text{aux}}/2-1},{\bm{x}}_{t}\rangle&0\\\
0&\langle{\bm{w}}_{2},{\bm{x}}_{t}\rangle&0&\langle{\bm{w}}_{4},{\bm{x}}_{t}\rangle&\cdots&0&\langle{\bm{w}}_{D_{\text{aux}}/2},{\bm{x}}_{t}\rangle\\\
\langle{\bm{w}}_{D_{\text{aux}}/2+1},{\bm{x}}_{t}\rangle&0&\langle{\bm{w}}_{D_{\text{aux}}/2+3},{\bm{x}}_{t}\rangle&0&\cdots&\langle{\bm{w}}_{D_{\text{aux}}-1},{\bm{x}}_{t}\rangle&0\\\
0&\langle{\bm{w}}_{D_{\text{aux}}/2+2},{\bm{x}}_{t}\rangle&0&\langle{\bm{w}}_{D_{\text{aux}}/2+4},{\bm{x}}_{t}\rangle&\cdots&0&\langle{\bm{w}}_{D_{\text{aux}}},{\bm{x}}_{t}\rangle\\\
0&0&\cdots&0&\cdots&0&0\\\ 0&0&\cdots&0&\cdots&0&0\end{bmatrix}$

Finally, a column-wise linear operation on ${\bm{S}}^{row}_{t}$ helps to get
the non-zero elements in the correct order.

$\displaystyle\bar{{\bm{S}}}^{col}_{t}=$
$\displaystyle\begin{bmatrix}\langle{\bm{w}}_{1},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{2},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{3},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{4},{\bm{x}}_{t}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}/2-1},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}/2},{\bm{x}}_{t}\rangle\\\
\langle{\bm{w}}_{D_{\text{aux}}/2+1},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}/2+2},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}/2+3},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}/2+4},{\bm{x}}_{t}\rangle&\cdots&\langle{\bm{w}}_{D_{\text{aux}}-1},{\bm{x}}_{t}\rangle&\langle{\bm{w}}_{D_{\text{aux}}},{\bm{x}}_{t}\rangle\\\
0&0&0&0&\cdots&0&0\\\ \vdots&\vdots&\vdots&\vdots&\cdots&\vdots&\vdots\\\
0&0&0&0&\cdots&0&0\end{bmatrix}$

The desired output is then given by
$\textsc{Vectorize}(\\{\bar{{\bm{s}}}^{col}_{t,j}\\}\\}_{j=1}^{D_{\text{aux}}})$,
which contains ${\bm{W}}{\bm{x}}_{t}$ in the first $D_{\text{aux}}$
coordinates. The operations that convert ${\bm{S}}_{t}$ to
${\bm{S}}_{t}^{col}$ and ${\bm{S}}_{t}^{row}$ to $\bar{{\bm{S}}}_{t}^{row}$
represents a split-wise $6$-split linear operation, while the operation that
converts ${\bm{S}}^{col}_{t}$ to ${\bm{S}}_{t}^{row}$ represents a dimension-
wise $6$-split linear operation. A naive linear operation on the output of the
attention heads would require $D_{\text{sim}}^{2}$ parameters, while its
replacement requires $D_{\text{sim}}^{2}/6$ parameters to represent a
dimension-wise $6$-split linear operation, and an additional
$12D_{\text{sim}}$ parameters to represent the split-wise $6$-split linear
operations.

##  Appendix E Self-attention layer

Figure 6: TinT simulates the forward pass of a self-attention layer of the
auxiliary model with a Linear Forward module (Figure 2) and a TinT softmax
attention layer (Definition C.1). The Linear Forward module computes the
query, key, and value vectors using a Linear Forward module on the current
embeddings, changing the prefix embeddings to correspond to ${\bm{W}}_{Q}$,
${\bm{W}}_{K}$, and ${\bm{W}}_{K}$ respectively. Figure 7: The gradient w.r.t.
the value vectors $\\{\partial_{{\bm{v}}_{t}}\\}$ (Definition E.2) forms the
integral component for both TinT self-attention backward and descent update
modules. TinT computes $\\{\partial_{{\bm{v}}_{t}}\\}$ using a softmax
attention and a linear attention layer. We first use residual connections to
copy the query and key vectors to the current embeddings from the TinT Self-
attention Forward module (Figure 6). The softmax attention layer re-computes
the attention scores $\\{a^{h}_{t,j}\\}$ between all token pairs $\\{(t,j)\\}$
and stores them in the token embeddings. The linear attention layer uses the
one-hot position embeddings of the input tokens as the query to use the
transposed attention scores $\\{a^{h}_{j,t}\\}$ for all token pairs
$\\{(t,j)\\}$ and use the gradients $\\{\partial_{{\bm{y}}_{t}}\\}$ as the
value vectors to compute $\\{\partial_{{\bm{v}}_{t}}\\}$.

We first introduce multi-head attention, generalizing single-head attention
(Definition B.1).

######  Definition E.1 (Auxiliary self-attention with $H_{\text{aux}}$ heads).

For query, key, and value weights
${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\in\mathbb{R}^{D_{\text{aux}}\times
D_{\text{aux}}}$ and bias
${\bm{b}}_{Q},{\bm{b}}_{K},{\bm{b}}_{V}\in\mathbb{R}^{D_{\text{aux}}}$, a
self-attention layer with $H_{\text{aux}}$ attention heads and a function
$f_{\mathrm{attn}}:\mathbb{R}^{T_{\text{aux}}}\to\mathbb{R}^{T_{\text{aux}}}$
takes a sequence $\\{{\bm{x}}_{t}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq
T_{\text{aux}}}$ as input and outputs $\\{{\bm{y}}_{t}\\}_{t\leq
T_{\text{aux}}}$, with

$\displaystyle{\bm{y}}_{t}=\textsc{Vectorize}(\\{\sum_{j\leq
T_{\text{aux}}}a^{h}_{t,j}{\bm{v}}^{h}_{j}\\}_{h\leq H_{\text{aux}}}).$ (9)

$a^{h}_{t,j}$ is defined as the attention score of head $h$ between tokens at
positions $t$ and $j$, and is given by

$\displaystyle
a^{h}_{t,j}=\mathrm{softmax}({\bm{K}}^{h}{\bm{q}}^{h}_{t})_{j}.$ (10)

Here, ${\bm{q}}_{t}$, ${\bm{k}}_{t}$, ${\bm{v}}_{t}$ denote the query, key,
and value vectors at each position $t$, computed as
${\bm{W}}_{Q}{\bm{x}}_{t}+{\bm{b}}_{Q}$,
${\bm{W}}_{K}{\bm{x}}_{t}+{\bm{b}}_{K}$, and
${\bm{W}}_{V}{\bm{x}}_{t}+{\bm{b}}_{V}$ respectively. In addition,
${\bm{q}}^{h}_{t},{\bm{k}}^{h}_{t},{\bm{v}}^{h}_{t}$ denote
$\textsc{Split}_{H_{\text{aux}}}({\bm{q}}_{t})_{h}$,
$\textsc{Split}_{H_{\text{aux}}}({\bm{k}}_{t})_{h}$, and
$\textsc{Split}_{H_{\text{aux}}}({\bm{v}}_{t})_{h}$ respectively for all
$t\leq T_{\text{aux}}$, and $h\leq H_{\text{aux}}$.
${\bm{K}}^{h}\in\mathbb{R}^{T_{\text{aux}}\times D_{\text{aux}}}$ is defined
with its rows as $\\{{\bm{k}}^{h}_{t}\\}_{t\leq T_{\text{aux}}}$ for all
$h\leq H_{\text{aux}}$.

In the discussions below, we consider a self-attention layer in the auxiliary
model with parameters
$\\{{\bm{W}}_{Q},{\bm{b}}_{Q},{\bm{W}}_{K},{\bm{b}}_{K},{\bm{W}}_{V},{\bm{b}}_{V}\\}$
that takes in input sequence ${\bm{x}}_{1},\cdots,{\bm{x}}_{T_{\text{aux}}}$
and outputs ${\bm{y}}_{1},\cdots,{\bm{y}}_{T_{\text{aux}}}$, with
$\\{{\bm{y}}_{t}\\}_{t=1}^{T_{\text{aux}}}$ given by (9). As in the
definition, ${\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}$ denote the query, key,
and value vectors for position $t$. We will use TinT self-attention modules in
order to simulate the operations on the auxiliary’s self-attention layer. To
do so, we will need $H_{\text{sim}}\geq H_{\text{aux}}$ in the corresponding
TinT self-attention modules.

####  TinT Self-attention forward module

The input embedding to this module ${\bm{e}}_{t}$ at each position $t$ will
contain ${\bm{x}}_{t}$ in its first $D_{\text{aux}}$ coordinates. The self-
attention module can be divided into four sub-operations: Computation of (a)
query vectors $\\{{\bm{q}}_{t}\\}_{t\leq T}$, (b) key vectors
$\\{{\bm{k}}_{t}\\}_{t\leq T}$, (c) value vectors $\\{{\bm{v}}_{t}\\}_{t\leq
T}$, and (d) $\\{{\bm{y}}_{t}\\}_{t\leq T}$ using (9). Please see Figure 6.

  * •

Sub-operations (a): The computation of query vector
${\bm{q}}_{t}:={\bm{W}}_{Q}{\bm{x}}_{t}+{\bm{b}}_{Q}$ at each position $t$ is
a linear operation involving parameters ${\bm{W}}_{Q},{\bm{b}}_{Q}$. Thus, we
can first feed in the stacked rows of ${\bm{W}}_{Q}$ and ${\bm{b}}_{Q}$ onto
the prefix embeddings $\\{{\bm{v}}_{j}\\}$. We use a Linear Forward module
(Appendix D) on the current embeddings and the prefix embeddings to get
embedding ${\bm{e}}^{q}_{t}$ at each position $t$ that contains ${\bm{q}}_{t}$
in the first $D_{\text{aux}}$ coordinates.

  * •

Sub-operations (b, c): Similar to (a), we feed in the stacked rows of the
necessary parameters onto the prefix embeddings $\\{{\bm{v}}_{j}\\}$, and call
two Linear Forward Modules (Appendix D) independently to get embeddings
${\bm{e}}^{k}_{t}$, and ${\bm{e}}^{v}_{t}$ containing ${\bm{k}}_{t}$ and
${\bm{v}}_{t}$ respectively.

We now combine the embeddings ${\bm{e}}^{q}_{t}$, ${\bm{e}}^{k}_{t}$, and
${\bm{e}}^{v}_{t}$ to get an embedding ${\bm{e}}_{t}$ that contain
${\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}$ in the first $3D_{\text{aux}}$
coordinates.

  * •

Sub-operation (d): Finally, we call a TinT self-attention module (Definition
C.1) on our current embeddings $\\{{\bm{e}}_{t}\\}_{t\leq T}$ to compute
$\\{{\bm{y}}_{t}\\}_{t\leq T}$. The query, key, and value parameters in the
self-attention module contain sub-Identity blocks that pick out the relevant
information from ${\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}$ stored in
${\bm{e}}_{t}$.

Remark: Sub-operations (a), (b), and (c) can be represented as a single linear
operation with a weight ${\bm{W}}\in\mathbb{R}^{3D_{\text{aux}}\times
D_{\text{aux}}}$ by concatenating the rows of
$\\{{\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\\}$ and a bias
${\bm{b}}\in\mathbb{R}^{3D_{\text{aux}}}$ that concatenates
$\\{{\bm{b}}_{Q},{\bm{b}}_{K},{\bm{b}}_{V}\\}$. Thus, they can be simulated
with a single Linear Forward Module, with ${\bm{W}},{\bm{b}}$ fed into the
prefix embeddings. However, we decide to separate them in order to limit the
number of prefix embeddings and the embedding size. E.g. for GPT-2,
$D_{\text{aux}}=768$. This demands either a $3\times$ increase in the
embedding size in TinT or a $3\times$ increase in the number of prefix
embeddings. Hence, in order to minimize the parameter cost, we call Linear
Forward Module separately to compute ${\bm{q}}_{t}$, ${\bm{k}}_{t}$, and
${\bm{v}}_{t}$ at each position $t$.

#### Auxiliary’s backpropagation through self-attention

For an auxiliary self-attention layer as defined in Definition E.1, the
backpropagation layer takes in the loss gradient w.r.t. output
($\\{\partial_{{\bm{y}}_{t}}\\}_{t\leq T_{\text{aux}}}$) and computes the loss
gradient w.r.t. input token ($\\{\partial_{{\bm{x}}_{t}}\\}_{t\leq
T_{\text{aux}}}$).

######  Definition E.2.

[Auxiliary self-attention backpropagation] For query, key, and value weights
${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\in\mathbb{R}^{D_{\text{aux}}\times
D_{\text{aux}}}$ and bias
${\bm{b}}_{Q},{\bm{b}}_{K},{\bm{b}}_{V}\in\mathbb{R}^{D_{\text{aux}}}$, the
backpropagation layer corresponding to a self-attention layer with
$H_{\text{aux}}$ attention heads takes a sequence
$\\{\partial_{{\bm{y}}_{t}}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq
T_{\text{aux}}}$ and $\\{{\bm{x}}_{t}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq
T_{\text{aux}}}$ as input and outputs $\\{\partial_{{\bm{x}}_{t}}\\}_{t\leq
T_{\text{aux}}}$, with

$\displaystyle\partial_{{\bm{x}}_{t}}$
$\displaystyle={\bm{W}}_{Q}^{\top}\partial_{{\bm{q}}_{t}}+{\bm{W}}_{K}^{\top}\partial_{{\bm{k}}_{t}}+{\bm{W}}_{V}^{\top}\partial_{{\bm{v}}_{t}},\quad\text{
with }$ $\displaystyle\partial_{{\bm{q}}_{t}}$
$\displaystyle=\textsc{Vectorize}(\\{\sum_{j}a^{h}_{t,j}((\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j})[{\bm{k}}^{h}_{j}-\sum_{j^{\prime}}a^{h}_{t,j^{\prime}}{\bm{k}}^{h}_{j^{\prime}}]\\}_{h\leq
H_{\text{aux}}});$ $\displaystyle\partial_{{\bm{k}}_{t}}$
$\displaystyle=\textsc{Vectorize}(\\{\sum_{j}a^{h}_{j,t}{\bm{q}}^{h}_{j}[(\partial_{{\bm{y}}_{j}^{h}})^{\top}({\bm{v}}^{h}_{t}-\sum_{j^{\prime}}a^{h}_{j,j^{\prime}}{\bm{v}}^{h}_{j^{\prime}})]\\}_{h\leq
H_{\text{aux}}});$ $\displaystyle\partial_{{\bm{v}}_{t}}$
$\displaystyle=\textsc{Vectorize}(\\{\sum_{j}a^{h}_{j,t}\partial_{{\bm{y}}^{h}_{j}}\\}_{h\leq
H_{\text{aux}}})$

Here, ${\bm{q}}_{t}$, ${\bm{k}}_{t}$, and ${\bm{v}}_{t}$ refer to query, key,
and value vectors at each position $t$, with the attention scores
$\\{a_{t,j}^{h}\\}_{t,j\leq T_{\text{aux}},h\leq H_{\text{aux}}}$.

#### Complexity of true backpropagation

The much-involved computation in the above operation is due to the computation
of $\partial_{{\bm{q}}_{t}}$ and $\partial_{{\bm{k}}_{t}}$ at each position
$t$. For the following discussion, we assume that our current embeddings
${\bm{e}}_{t}$ contain ${\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}$, in addition
to the gradient $\partial_{{\bm{y}}_{t}}$. The computation of
$\partial_{{\bm{q}}_{t}}$ (and similarly $\partial_{{\bm{k}}_{t}}$) at any
position $t$ involves the following sequential computations and the necessary
TinT modules.

  * •

$\\{\\{\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j}\\}_{j\leq
T_{\text{aux}}}\\}_{h\leq H_{\text{aux}}}$ with a TinT linear self-attention
module (Definition C.1), with atleast $H_{\text{aux}}$ attention heads that
represent the attention score between ${\bm{e}}_{t}$ and any other token
${\bm{e}}_{j}$, by
$\\{(\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j}\\}_{h\leq
H_{\text{aux}}}$.

  * •

Attention scores $\\{a^{h}_{t,j}\\}_{h\leq H_{\text{aux}}}$, which requires a
TinT softmax self-attention module (Definition C.1), with at least
$H_{\text{aux}}$ heads, that uses the already present
$\\{{\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}\\}$ in the current embeddings
${\bm{e}}_{t}$ to re-compute the attention scores.

  * •

$\\{a_{t,j}^{h}(\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j}\\}_{h\leq
H_{\text{aux}}}$ for all $j\leq T_{\text{aux}}$ by multiplying the attention
scores $\\{a_{t,j}^{h}\\}_{h\leq H_{\text{aux}}}$ with
$\\{(\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j}\\}_{h\leq
H_{\text{aux}}}$ using an MLP layer (Lemma C.4). Furthermore,
$\\{\sum_{j}a^{h}_{t,j}{\bm{k}}^{h}_{j}\\}_{h\leq H_{\text{aux}}}$ needs to be
computed in parallel as well, with additional attention heads.

  * •

$\partial_{{\bm{y}}_{t}}$ with a TinT linear self-attention module (Definition
C.1), with atleast $H_{\text{aux}}$ attention heads that represent the
attention score between any token ${\bm{e}}_{j}$ and ${\bm{e}}_{t}$ by
$\\{a_{t,j}^{h}(\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j}\\}_{h\leq
H_{\text{aux}}}$, with value vectors given by
$\\{{\bm{k}}^{h}_{j}-\sum_{j^{\prime}}a^{h}_{t,j^{\prime}}{\bm{k}}^{h}_{j^{\prime}}\\}_{h\leq
H_{\text{aux}}}$.

The sequential computation requires the simulator to store
$\\{\\{\partial_{{\bm{y}}_{t}^{h}})^{\top}{\bm{v}}^{h}_{j}\\}_{j\leq
T_{\text{aux}}}\\}_{h\leq H_{\text{aux}}}$ and $\\{a^{h}_{t,j}\\}_{h\leq
H_{\text{aux}}}$ in the token embedding ${\bm{e}}_{t}$, which requires an
additional $2T_{\text{aux}}H_{\text{aux}}$ embedding dimension size. To avoid
the much-involved computation for the true gradient propagation, we instead
only use the gradients w.r.t. ${\bm{v}}_{t}$.

#### Approximate auxiliary self-attention backpropagation

We formally extend the definition of approximate gradients
$\\{\partial_{{\bm{x}}_{t}}\\}_{t=1}^{T_{\text{aux}}}$ from Definition E.3 to
multi-head attention in Definition E.3.

######  Definition E.3.

For query, key, and value weights
${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\in\mathbb{R}^{D_{\text{aux}}\times
D_{\text{aux}}}$ and bias
${\bm{b}}_{Q},{\bm{b}}_{K},{\bm{b}}_{V}\in\mathbb{R}^{D_{\text{aux}}}$, the
approximate backpropagation layer corresponding to a self-attention layer with
$H_{\text{aux}}$ attention heads takes a sequence
$\\{\partial_{{\bm{y}}_{t}}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq{T_{\text{aux}}}}$
and $\\{{\bm{x}}_{t}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq{T_{\text{aux}}}}$
as input and outputs
$\\{\partial_{{\bm{x}}_{t}}:=\textsc{Vectorize}(\\{\partial_{{\bm{x}}^{h}_{t}}\\}_{h\leq
H_{\text{aux}}})\\}_{t\leq{T_{\text{aux}}}}$, with

$\displaystyle\hat{\partial_{{\bm{x}}_{t}}}={\bm{W}}_{V}^{\top}\partial_{{\bm{v}}_{t}},\quad\text{
where
}\partial_{{\bm{v}}_{t}}=\textsc{Vectorize}(\\{\sum_{j}a^{h}_{j,t}\partial_{{\bm{y}}^{h}_{j}}\\}_{h\leq
H_{\text{aux}}})$

Here, ${\bm{q}}_{t}$, ${\bm{k}}_{t}$, and ${\bm{v}}_{t}$ refer to query, key,
and value vectors at each position $t$, as defined in Definition E.1, with the
attention scores $\\{a_{t,j}^{h}\\}_{t,j\leq{T_{\text{aux}}},h\leq
H_{\text{aux}}}$ defined in Equation 10.

Figure 8: TinT simulates the backward pass of a self-attention layer of the
auxiliary model using a Linear Backward module (Figure 4). The input
embeddings contain the gradient of the loss w.r.t. the value vectors
($\partial_{{\bm{v}}_{t}}$) computed in Figure 7. The value matrix
${\bm{W}}_{V}$ is encoded in the prefix embeddings. We call the Linear
Backward module on this sequence. Figure 9:  TinT simulates the backward pass
of the self-attention layer in the auxiliary model by employing the Linear
Descent module (Figure 5). The input embeddings consist of the gradient of the
loss with respect to the value vectors ($\partial_{{\bm{v}}_{t}}$) computed in
Figure 7. Additionally, we incorporate a residual connection to copy the input
from the Self-attention Forward module (Figure 6) into ${\bm{x}}_{t}$. Before
invoking the Linear Descent module, we represent the value parameters
(${\bm{W}}_{V}$) into the prefix embeddings. TinT simulates the backward pass
of a self-attention layer of the auxiliary model using a Linear Descent module
(Figure 5).

In the upcoming theorem, we formally show that if on a given sequence
$\\{{\bm{x}}_{t}\\}_{t\leq T_{\text{aux}}}$, for all token positions all the
attention heads in a self-attention layer primarily attend to a single token,
then the approximate gradient $\hat{\partial_{{\bm{x}}_{t}}}$ is close to the
true gradient $\partial_{{\bm{x}}_{t}}$ at each position $t$.

######  Definition E.4 ($\varepsilon$-hard attention head).

For the Self-Attention layer of $H_{\text{aux}}$ heads in Definition E.1, on a
given input sequence $\\{{\bm{x}}_{t}\\}_{t=1}^{T_{\text{aux}}}$, an attention
head $h\leq H_{\text{aux}}$ is defined to be $\varepsilon$-hard on the input
sequence, if for all positions $t\leq T_{\text{aux}}$, there exists a position
$t_{0}\leq T_{\text{aux}}$ such that $a_{t,t_{0}}^{h}\geq 1-\varepsilon$.

######  Theorem E.5.

With the notations in Definitions E.1, E.2 and E.3, if on a given input
sequence $\\{{\bm{x}}_{t}\\}_{t=1}^{T_{\text{aux}}}$, with its query, key, and
value vectors
$\\{{\bm{q}}_{t},{\bm{k}}_{t},{\bm{v}}_{t}\\}_{t=1}^{T_{\text{aux}}}$, all the
$H_{\text{aux}}$ attention heads are $\varepsilon$-hard for some
$\varepsilon>0$, then for a given sequence of gradients
$\\{\partial_{{\bm{y}}_{t}}\\}_{t=1}^{T_{\text{aux}}}$,

$\displaystyle\left\|\partial_{{\bm{q}}_{t}}\right\|_{2},\left\|\partial_{{\bm{k}}_{t}}\right\|_{2}\leq\mathcal{O}(\varepsilon
B_{x}^{2}B^{2}_{w}B_{y}),\quad\text{ for all }t\leq T_{\text{aux}},$

where $B_{x}=\max_{t\leq T_{\text{aux}}}\left\|{\bm{x}}_{t}\right\|_{2}$,
$B_{y}=\max_{t\leq T_{\text{aux}}}\left\|\partial_{{\bm{y}}_{t}}\right\|_{2}$,
and
$B_{w}=\max\\{\left\|{\bm{W}}_{K}\right\|_{2},\left\|{\bm{W}}_{Q}\right\|_{2},\left\|{\bm{W}}_{V}\right\|_{2},\left\|{\bm{b}}_{V}\right\|_{2},\left\|{\bm{b}}_{K}\right\|_{2},\left\|{\bm{b}}_{V}\right\|_{2}\\}$.

This implies, for each position $t$,
$\left\|\hat{\partial_{{\bm{x}}_{t}}}-\partial_{{\bm{x}}_{t}}\right\|_{2}\leq\mathcal{O}(\varepsilon
B_{x}^{2}B_{w}^{3}B_{y}).$

####  TinT Self-attention backpropagation module

The input embeddings ${\bm{e}}_{t}$ contain $\partial_{{\bm{y}}_{t}}$ in the
first $D_{\text{aux}}$ coordinates. Since we require to re-compute the
attention scores $\\{a^{h}_{t,j}\\}_{j\leq T_{\text{aux}},h\leq
H_{\text{aux}}}$, we need to copy the query, key, and value vectors
${\bm{q}}_{t}$, ${\bm{k}}_{t}$, and ${\bm{v}}_{t}$ from the TinT self-
attention Forward module at each position $t$. Furthermore, we use the
residual connection to copy the prefix embeddings $\\{{\bm{v}}_{j}\\}$, which
contain the rows of ${\bm{W}}_{V}$, from the TinT self-attention Forward
module.

The operation can be divided into three sub-operations: Computing (a)
attention scores $\\{a_{t,j}^{h}\\}_{h\leq H_{\text{aux}}}$ for all $j\leq
T_{\text{aux}}$, at each position $t$, (b) $\partial_{{\bm{v}}_{t}}$ from
$\\{a_{t,j}^{h}\\}_{h\leq H_{\text{aux}}}$ and $\partial_{{\bm{y}}_{t}}$, and
(c) $\hat{\partial_{{\bm{x}}_{t}}}$ from $\partial_{{\bm{v}}_{t}}$.

  * •

Sub-operation (a): Since, the current embeddings ${\bm{e}}_{t}$ contain
${\bm{q}}_{t},{\bm{k}}_{t}$, we can simply call a self-attention attention
module to compute the attention scores $\\{a_{t,j}^{h}\\}_{h\leq
H_{\text{aux}}}$ for all $j\leq T$ and store them in the current embeddings.
We further retain $\partial_{{\bm{y}}_{t}}$ and ${\bm{v}}_{t}$ for further
operations using residual connections.

  * •

Sub-operation (b): With the current embeddings ${\bm{e}}_{t}$ containing the
attention scores $\\{a_{t,j}^{h}\\}_{h\leq H_{\text{aux}}}$ for all $j\leq T$,
and the gradient $\partial_{{\bm{y}}_{t}}$, we can compute
$\partial_{{\bm{v}}_{t}}$ using a TinT linear self-attention module with
atleast $H_{\text{aux}}$ attention heads, that represent the attention scores
between tokens ${\bm{e}}_{t}$ and ${\bm{e}}_{j}$ for any $j$ as
$\\{a_{j,t}^{h}\\}_{h\leq H_{\text{aux}}}$ and use
$\textsc{Split}_{H_{\text{aux}}}(\partial_{{\bm{y}}_{t}})$ as their value
vectors.

  * •

Sub-operation (c): And finally, the computation of
$\hat{\partial_{{\bm{x}}_{t}}}$ is identical to the backpropagation through a
linear layer, with parameters ${\bm{W}}_{V}$ and ${\bm{b}}_{V}$. Hence, we
call a Linear backpropagation module on the current embeddings, that contain
$\partial_{{\bm{y}}_{t}}$ and the prefix embeddings that contain
${\bm{W}}_{V}$ and ${\bm{b}}_{V}$.

#### Separating sub-operations (a) and (b)

The operation for computing $\partial_{{\bm{v}}_{t}}$ in Definition E.3 looks
very similar to the computation of ${\bm{y}}_{t}$ in Equation 9. However, the
major difference is that instead of the attention scores being
$\\{a_{t,j}^{h}\\}_{h\leq H_{\text{aux}}}$ between token $t$ and any token
$j$, we need the attention scores to be $\\{a_{j,t}^{h}\\}_{h\leq
H_{\text{aux}}}$. Thus, unless our model allows a transpose operation on the
attention scores, we need to first store them in our embeddings and then use
an additional self-attention module that can pick the right attention scores
between tokens using position embeddings. Please see Figure 8.

#### Auxiliary’s value descent update

Similar to the complexity of true backpropagation, the descent updates for
${\bm{W}}_{Q},{\bm{b}}_{Q},{\bm{W}}_{K},{\bm{b}}_{K}$ are quite expensive to
express with the transformer layers. Hence, we focus simply on updating on
${\bm{W}}_{V},{\bm{b}}_{V}$, while keeping the others fixed.

######  Definition E.6 (Auxiliary self-attention value descent).

For query, key, and value weights
${\bm{W}}_{Q},{\bm{W}}_{K},{\bm{W}}_{V}\in\mathbb{R}^{D_{\text{aux}}\times
D_{\text{aux}}}$ and bias
${\bm{b}}_{Q},{\bm{b}}_{K},{\bm{b}}_{V}\in\mathbb{R}^{D_{\text{aux}}}$, the
value descent layer corresponding to a self-attention layer with
$H_{\text{aux}}$ attention heads and any function
$f_{\mathrm{attn}}:\mathbb{R}^{T_{\text{aux}}}\to\mathbb{R}^{T_{\text{aux}}}$
takes in a batch of gradients
$\\{\partial_{{\bm{y}}_{t}}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq
T_{\text{aux}}}$ and inputs
$\\{{\bm{x}}_{t}\in\mathbb{R}^{D_{\text{aux}}}\\}_{t\leq T_{\text{aux}}}$ and
updates ${\bm{W}}_{V},{\bm{b}}_{V}$ as follows:

$\displaystyle{\bm{W}}_{V}\leftarrow{\bm{W}}_{V}-\eta\sum_{t\leq
T_{\text{aux}}}\partial_{{\bm{v}}_{t}}{\bm{x}}_{t}^{\top},\quad{\bm{b}}_{V}\leftarrow{\bm{b}}_{V}-\eta\sum_{t\leq
T_{\text{aux}}}\partial_{{\bm{v}}_{t}},$ $\displaystyle\text{ where
}\partial_{{\bm{v}}_{t}}=\textsc{Vectorize}(\\{\sum_{j}a^{h}_{j,t}\partial_{{\bm{y}}^{h}_{j}}\\}_{h\leq
H_{\text{aux}}})$

Here, ${\bm{v}}_{t}$ refers to value vectors at each position $t$, as defined
in Definition E.1.

####  TinT Self-attention descent module

The input embeddings contain $\partial_{{\bm{v}}_{t}}$ in the first
$D_{\text{aux}}$ coordinates, from the TinT self-attention backpropagation
module. Furthermore, the prefix embeddings $\\{{\bm{v}}_{j}\\}$ contain the
stacked rows of ${\bm{W}}_{V}$ and ${\bm{b}}_{V}$, continuing from the TinT
self-attention backpropagation module.

Since we further need the input ${\bm{x}}_{t}$ to the auxiliary self-attention
layer under consideration, we use residual connections to copy ${\bm{x}}_{t}$
from the TinT self-attention Forward module at each position $t$.

The updates of ${\bm{W}}_{V}$ and ${\bm{b}}_{V}$ are equivalent to the
parameter update in a linear layer, involving gradients
$\\{\partial_{{\bm{v}}_{t}}\\}$ and input $\\{{\bm{x}}_{t}\\}$. Thus, we call
a Linear descent module on the current embeddings and the prefix embeddings to
get the updated value parameters. Please see Figure 9.

###  E.1 Proofs of theorems and gradient definitions

We restate the theorems and definitions, before presenting their proofs for
easy referencing.

See E.2

###### Derivation of gradient in Definition E.2.

Recalling the definition of ${\bm{y}}_{t}$ from Definition E.1,

$\displaystyle{\bm{y}}_{t}=\textsc{Vectorize}(\\{\sum_{j\leq
T_{\text{aux}}}a^{h}_{t,j}{\bm{v}}^{h}_{j}\\}_{h\leq H_{\text{aux}}});\quad
a^{h}_{t,j}=\mathrm{softmax}({\bm{K}}^{h}{\bm{q}}^{h}_{t})_{j},$
$\displaystyle{\bm{q}}_{t}={\bm{W}}_{Q}{\bm{x}}_{t}+{\bm{b}}_{Q}\quad{\bm{k}}_{t}={\bm{W}}_{K}{\bm{x}}_{t}+{\bm{b}}_{K},\quad{\bm{v}}_{t}={\bm{W}}_{V}{\bm{x}}_{t}+{\bm{b}}_{V}.$

${\bm{q}}^{h}_{t},{\bm{k}}^{h}_{t},{\bm{v}}^{h}_{t}$ denote
$\textsc{Split}_{H_{\text{aux}}}({\bm{q}}_{t})_{h}$,
$\textsc{Split}_{H_{\text{aux}}}({\bm{k}}_{t})_{h}$, and
$\textsc{Split}_{H_{\text{aux}}}({\bm{v}}_{t})_{h}$ respectively for all
$t\leq T_{\text{aux}}$, and $h\leq H_{\text{aux}}$.
${\bm{K}}^{h}\in\mathbb{R}^{T_{\text{aux}}\times D_{\text{aux}}}$ is defined
with its rows as $\\{{\bm{k}}^{h}_{t}\\}_{t\leq T_{\text{aux}}}$ for all
$h\leq H_{\text{aux}}$.

We explain the proof for an arbitrary token position $t$. With the application
of the chain rule, we have

$\displaystyle\partial_{{\bm{x}}_{t}}$
$\displaystyle=(\frac{\partial{\bm{q}}_{t}}{\partial{\bm{x}}_{t}})^{\top}\partial_{{\bm{q}}_{t}}+(\frac{\partial{\bm{k}}_{t}}{\partial{\bm{x}}_{t}})^{\top}\partial_{{\bm{k}}_{t}}+(\frac{\partial{\bm{v}}_{t}}{\partial{\bm{x}}_{t}})^{\top}\partial_{{\bm{v}}_{t}}$
$\displaystyle={\bm{W}}_{Q}^{\top}\partial_{{\bm{q}}_{t}}+{\bm{W}}_{K}^{\top}\partial_{{\bm{k}}_{t}}+{\bm{W}}_{V}^{\top}\partial_{{\bm{v}}_{t}},$

where the second step follows from the definitions of
${\bm{q}}_{t},{\bm{k}}_{t},$ and ${\bm{v}}_{t}$ respectively.

#### Computation of $\partial_{{\bm{q}}_{t}}$:

With the Split operation of ${\bm{q}}_{t}$ across $H_{\text{aux}}$ heads for
the computation of ${\bm{y}}_{t}$, the computation of the backpropagated
gradient $\partial_{{\bm{q}}_{t}}$ itself needs to be split across
$H_{\text{aux}}$ heads. Furthermore, query vector ${\bm{q}}_{t}$ only affects
${\bm{y}}_{t}$, implying
$\frac{\partial{\bm{y}}_{t^{\prime}}}{\partial{\bm{q}}_{t}}=0$ for any
$t^{\prime}\neq t$. Thus, we have for any head $h\leq H_{\text{aux}}$, if
${\bm{y}}_{t}^{h}$ represents the output of attention head $h$, given by
$\sum_{j\leq T_{\text{aux}}}a^{h}_{t,j}{\bm{v}}_{j}^{h}$,

$\displaystyle\partial_{{\bm{q}}_{t}^{h}}$
$\displaystyle=(\frac{\partial{\bm{y}}_{t}^{h}}{\partial{\bm{q}}_{t}^{h}})^{\top}\partial_{{\bm{y}}_{t}^{h}}$
$\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{j}^{h},\partial_{{\bm{y}}_{t}^{h}}\rangle\frac{\partial
a^{h}_{t,j}}{\partial{\bm{q}}_{t}^{h}}$ $\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{j}^{h},\partial_{{\bm{y}}_{t}^{h}}\rangle\frac{\partial}{\partial{\bm{q}}_{t}^{h}}\left(\frac{e^{\langle{\bm{k}}_{j}^{h},{\bm{q}}_{t}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}\right)$
(11) $\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{j}^{h},\partial_{{\bm{y}}_{t}^{h}}\rangle\left[\frac{1}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}\frac{\partial
e^{\langle{\bm{k}}_{j}^{h},{\bm{q}}_{t}^{h}\rangle}}{\partial{\bm{q}}_{t}^{h}}-\left(\frac{e^{\langle{\bm{k}}_{j}^{h},{\bm{q}}_{t}^{h}\rangle}}{(\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle})^{2}}\right)\sum_{j^{\prime}\leq
T_{\text{aux}}}\frac{\partial
e^{\langle{\bm{k}}_{j^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}{\partial{\bm{q}}_{t}^{h}}\right]$
(12) $\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{j}^{h},\partial_{{\bm{y}}_{t}^{h}}\rangle\left[\left(\frac{e^{\langle{\bm{k}}_{j}^{h},{\bm{q}}_{t}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}\right){\bm{k}}_{j}^{h}-\left(\frac{e^{\langle{\bm{k}}_{j}^{h},{\bm{q}}_{t}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}\right)\sum_{j^{\prime}\leq
T_{\text{aux}}}\left(\frac{e^{\langle{\bm{k}}_{j^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{t}^{h}\rangle}}\right){\bm{k}}^{h}_{j^{\prime}}\right]$
(13) $\displaystyle=\sum_{j\leq
T_{\text{aux}}}a^{h}_{t,j}\langle{\bm{v}}_{j}^{h},\partial_{{\bm{y}}_{t}^{h}}\rangle\left({\bm{k}}_{j}^{h}-\sum_{j^{\prime}\leq
T_{\text{aux}}}a^{h}_{t,j^{\prime}}{\bm{k}}^{h}_{j^{\prime}}\right).$

In Equation 11, we have expanded the definition of $\mathrm{softmax}$ in
$a^{h}_{t,j}:=\mathrm{softmax}({\bm{K}}^{h}{\bm{q}}_{t}^{h})_{j}$ in order to
better motivate the derivative of $a^{h}_{t,j}$ w.r.t. ${\bm{q}}_{t}^{h}$.
Finally, $\partial_{{\bm{q}}_{t}}$ is given by
$\textsc{Vectorize}(\\{\partial_{{\bm{q}}_{t}^{h}}\\}_{h\leq
H_{\text{aux}}}).$

#### Computation of $\partial_{{\bm{k}}_{t}}$:

Continuing as the computation of $\partial_{{\bm{q}}_{t}}$, we split the
computation of $\partial_{{\bm{k}}_{t}}$ across the $H_{\text{aux}}$ attention
heads. However, unlike ${\bm{q}}_{t}$, ${\bm{k}}_{t}$ affects ${\bm{y}}_{j}$
for all $j\leq T_{\text{aux}}$. For any head $h\leq H_{\text{aux}}$, we follow
the chain-rule step by step to get

$\displaystyle\partial_{{\bm{k}}_{t}^{h}}$ $\displaystyle=\sum_{j\leq
T_{\text{aux}}}(\frac{\partial{\bm{y}}_{j}^{h}}{\partial{\bm{k}}_{t}^{h}})^{\top}\partial_{{\bm{y}}_{j}^{h}}=\sum_{j\leq
T_{\text{aux}}}\left(\frac{\partial\sum_{j^{\prime}\leq
T_{\text{aux}}}a_{j,j^{\prime}}{\bm{v}}_{j^{\prime}}^{h}}{\partial{\bm{k}}_{t}^{h}}\right)^{\top}\partial_{{\bm{y}}_{j}^{h}}$
$\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{t}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle\frac{\partial
a^{h}_{j,t}}{\partial{\bm{k}}_{t}^{h}}+\sum_{j\leq
T_{\text{aux}}}\sum_{j^{\prime}\leq T_{\text{aux}};j^{\prime}\neq
t}\langle{\bm{v}}_{j^{\prime}}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle\frac{\partial
a^{h}_{j,j^{\prime}}}{\partial{\bm{k}}_{t}^{h}}$ (14)
$\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{t}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle\frac{\partial}{\partial{\bm{k}}_{t}^{h}}\left(\frac{e^{\langle{\bm{k}}_{t}^{h},{\bm{q}}_{j}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}\right)$
(15) $\displaystyle+\sum_{j\leq T_{\text{aux}}}\sum_{j^{\prime}\leq
T_{\text{aux}};j^{\prime}\neq
t}\langle{\bm{v}}_{j^{\prime}}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle\frac{\partial}{\partial{\bm{k}}_{t}^{h}}\left(\frac{e^{\langle{\bm{k}}_{j^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}\right)$
(16) $\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{t}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle\left[\left(\frac{e^{\langle{\bm{k}}_{t}^{h},{\bm{q}}_{j}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}\right){\bm{q}}_{j}^{h}-\left(\frac{e^{\langle{\bm{k}}_{t}^{h},{\bm{q}}_{j}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}\right)^{2}{\bm{q}}_{j}^{h}\right]$
(17) $\displaystyle-\sum_{j\leq T_{\text{aux}}}\sum_{j^{\prime}\leq
T_{\text{aux}};j^{\prime}\neq
t}\langle{\bm{v}}_{j^{\prime}}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle\left(\frac{e^{\langle{\bm{k}}_{j^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}\right)\left(\frac{e^{\langle{\bm{k}}_{t}^{h},{\bm{q}}_{j}^{h}\rangle}}{\sum_{t^{\prime}\leq
T_{\text{aux}}}e^{\langle{\bm{k}}_{t^{\prime}}^{h},{\bm{q}}_{j}^{h}\rangle}}\right){\bm{q}}_{j}^{h}$
(18) $\displaystyle=\sum_{j\leq
T_{\text{aux}}}\langle{\bm{v}}_{t}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle(a^{h}_{j,t}-(a^{h}_{j,t})^{2}){\bm{q}}_{j}^{h}-\sum_{j\leq
T_{\text{aux}}}\sum_{j^{\prime}\leq T_{\text{aux}};j^{\prime}\neq
t}\langle{\bm{v}}_{j^{\prime}}^{h},\partial_{{\bm{y}}_{j}^{h}}\rangle
a^{h}_{j,j^{\prime}}a^{h}_{j,t}{\bm{q}}_{j}^{h}$ $\displaystyle=\sum_{j\leq
T_{\text{aux}}}a^{h}_{j,t}\langle\partial_{{\bm{y}}_{j}^{h}},{\bm{v}}^{h}_{t}-\sum_{j^{\prime}}a^{h}_{j,j^{\prime}}{\bm{v}}^{h}_{j^{\prime}}\rangle{\bm{q}}^{h}_{j}$