Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression

Introduction

Transformer neural networks (Vaswani et al., 2017) have become the default architecture for natural language processing (Devlin et al., 2019; Brown et al., 2020; OpenAI, 2023), and have even been adopted by other areas like computer vision (Dosovitskiy et al., 2021). As first demonstrated by GPT-3 (Brown et al., 2020), Transformers excel at in-context learning (ICL)—learning from input-output pairs provided as inputs to the model, without updating their model parameters. Through in-context learning, Transformer-based Large Language Models (LLMs) can achieve state-of-the-art few-shot performance across a wide variety of downstream tasks (Rae et al., 2022; Smith et al., 2022; Thoppilan et al., 2022; Chowdhery et al., 2022).

Given the importance of Transformers and ICL, many prior efforts have attempted to understand how Transformers perform in-context learning. Prior work suggests Transformers learn classification similar to support vector machines (Tarzanagh et al., 2023b, a) and can approximate linear functions well in-context (Garg et al., 2022). Specifically to linear regression tasks, prior work has tried to understand the ICL mechanism and the dominant hypothesis is that Transformers learn in-context by running optimizations internally through gradient-based algorithms (von Oswald et al., 2022, 2023; Ahn et al., 2023; Dai et al., 2023).

This paper presents strong evidence for a competing hypothesis: Transformers trained to perform in-context linear regression learn to implement a higher-order optimization method rather than a first-order method like Gradient Descent. In particular, Transformers implement a method very similar to Newton-Schulz’s Method, also known as the Iterative Newton’s Method, which iteratively improves an estimate of the inverse of the design matrix to compute the optimal weight vector. Across many layers of the Transformer, subsequent layers approximately compute more and more iterations of Newton’s Method, with increasingly better predictions; both eventually converge to the optimal minimum-norm solution found by ordinary least squares (OLS). Interestingly, this mechanism is specific to Transformers: LSTMs do not learn these same higher-order methods, as their predictions do not even improve across layers.

We present both empirical and theoretical evidence for our claims. Empirically, Transformer induced weights and residuals are similar to Iterative Newton and deeper layers match Newton with more iterations (see Figures 1 and 9). Transformers can also handle ill-conditioned problems without requiring significantly more layers, where GD would suffer from slow convergence but Iterative Newton would not. Crucially, Transformers share the same rate of convergence as Iterative Newton and are exponentially faster than Gradient Descent. Theoretically, we show that Transformer circuits can efficiently implement Iterative Newton, with the number of layers depending linearly on the number of iterations and the dimensionality of the hidden states depending linearly on the dimensionality of the data. Overall, our work provides a mechanistic account of how Transformers perform in-context learning that not only explains model behavior better than previous hypotheses, but also hints at what makes Transformers so well-suited for ICL compared with other neural architectures.

Related Work

GPT-3 (Brown et al., 2020) first showed that Transformer-based large language models can “learn” to perform new tasks from in-context demonstrations (i.e., input-output pairs). Since then, a large body of work in NLP has studied in-context learning, for instance by understanding how the choice and order of demonstrations affects results (Lu et al., 2022; Liu et al., 2022; Rubin et al., 2022; Su et al., 2023; Chang and Jia, 2023; Nguyen and Wong, 2023), studying the effect of label noise (Min et al., 2022c; Yoo et al., 2022; Wei et al., 2023), and proposing methods to improve ICL accuracy (Zhao et al., 2021; Min et al., 2022a, b).

Inspired by the phenomenon of ICL by large language models, subsequent work has studied how Transformers learn in-context beyond NLP tasks. Garg et al. (2022) first investigated Transformers’ ICL abilities for various classical machine learning problems, including linear regression. We largely adopt their linear regression setup in this work. Li et al. (2023) formalize in-context learning as an algorithm learning problem where a Transformer model implicitly constructs a hypothesis function at inference-time and obtain generalization bounds for ICL. Han et al. (2023) suggests that Transformers learn in-context by performing Bayesian inference on prompts, which can be asymptotically interpreted as kernel regression. Tarzanagh et al. (2023a) and Tarzanagh et al. (2023b) show that Transformers can find max-margin solutions for classification tasks and act as support vector machines. Zhang et al. (2023) prove that a linear attention Transformer trained by gradient flow can indeed in-context learn class of linear models. Raventós et al. (2023) explore how diverse pretraining data can enable models to perform ICL on new tasks.

A growing body of work has suggested that Transformers learn in-context by implementing gradient descent within their internal representations. Akyürek et al. (2022) summarize operations that Transformers can implement, such as multiplication and affine transformations, and show that Transformers can implement gradient descent for linear regression using these operations. Concurrently, von Oswald et al. (2022) argue that Transformers learn in-context via gradient descent, where one layer performs one gradient update. In subsequent work, von Oswald et al. (2023) further argue that Transformers are strongly biased towards learning to implement gradient-based optimization routines. Ahn et al. (2023) extend the work of von Oswald et al. (2022) by showing Transformers can learn to implement preconditioned Gradient Descent, where the pre-conditioner can adapt to the data. Bai et al. (2023) provides detailed constructions for how Transformers can implement a range of learning algorithms via gradient descent. Finally, Dai et al. (2023) conduct experiments on NLP tasks and conclude that Transformer-based language models performing ICL behave similarly to models fine-tuned via gradient descent ; however, concurrent work (Shen et al., 2023b) argues that real-world LLMs do not perform ICL via gradient descent. In this paper, we argue that Transformers actually learn to perform in-context learning by implementing a higher-order optimization method, not gradient descent. Predictions made by different Transformer layers match iterations of higher-order optimization methods better than they match iterations of gradient descent; moreover, Transformers can handle ill-conditioned data, unlike Gradient Descent.

Our work attempts to understand the mechanism through which Transformers perform in-context learning. Prior work has studied other aspects of Transformers’ internal mechanisms, including reverse-engineering language models (Wang et al., 2022), the grokking phenomenon (Power et al., 2022; Nanda et al., 2023), manipulating attention maps (Hassid et al., 2022), and automated circuit finding (Conmy et al., 2023).

Problem Setup

Does the algorithm Transformers learn for linear regression resemble any known algorithm?

To investigate this question, we first discuss various known algorithms for linear regression. We then compare them with Transformers empirically in §4 and theoretically in §5. For iterative algorithms, we care particularly about their rates of convergence (the number of steps required to reach $\epsilon$ error).

This method finds the minimum-norm solution to the objective:

The Ordinary Least Squares (OLS) solution has a closed form given by the Normal Equations:

where $\bm{S}:=\bm{X}^{\top}\bm{X}$ and $\bm{S}^{\dagger}$ is the pseudo-inverse (Moore, 1920) of $\bm{S}$.

It is known that Gradient Descent requires $\mathcal{O}\left(\kappa(\bm{S})\log(1/\epsilon)\right)$ steps to converge to an $\epsilon$ error where $\kappa(\bm{S})=\frac{\lambda_{\max}(\bm{S})}{\lambda_{\min}(\bm{S})}$ is the condition number. Thus, when $\kappa(\bm{S})$ is large, Gradient Descent converges slowly (Boyd and Vandenberghe, 2004).

While GD computes the gradient with respect to the full data matrix $\bm{X}$ at each iteration, Online Gradient Descent (OGD) is an online algorithm that only computes gradients on the newly received data point $\{\bm{x}_{k},y_{k}\}$ at step $k$:

This computes an approximation of the psuedo inverse using the moments of $\bm{S}$. In contrast to GD, the Iterative Newton’s method only requires $\mathcal{O}(\log\kappa(\bm{S})+\log\log(1/\epsilon))$ steps to converge to an $\epsilon$ error (Soderstrom and Stewart, 1974; Pan and Schreiber, 1991). Note that this is exponentially faster than the convergence rate of Gradient Descent.

2 Solving Linear Regression with Transformers

We will use neural network models such as Transformers to solve this linear regression task. As shown in Figure 2, at time step $t+1$, the model sees the first $t$ in-context examples $\{\bm{x}_{i},y_{i}\}_{i=1}^{t}$, and then makes predictions for $\bm{x}_{t+1}$, whose label $y_{t+1}$ is not observed by the Transformers model.

We randomly initialize our models and then train them on the linear regression task to make predictions for every number of in-context examples $t$, where $t\in[n]$. Training and test data are both drawn from $P_{\mathcal{D}}$. To make the input prompts contain both $\bm{x}_{i}$ and $y_{i}$, we follow same the setup as Garg et al. (2022)’s to zero-pad $y_{i}$’s, and use the same decoder-only GPT-2 model (Radford et al., 2019) with softmax activation and causal attention mask (discussed later in eq. 6).

We now present the key mathematical details for the Transformer architecture, and how they can be used for in-context learning. First, the causal attention mask enforces that attention heads can only attend to hidden states of previous time steps, and is defined as follows.

In particular for the linear regression task, Transformers perform in-context learning as follows

3 LSTM

4 Measuring Algorithmic Similarity

We propose two metrics to measure the similarity between linear regression algorithms.

For a linear regression algorithm ${\mathcal{A}}$, let ${\mathcal{A}}(\bm{x}_{t+1}\mid\{\bm{x}_{i},y_{i}\}_{i=1}^{t})$ denote its prediction on the $(t+1)$-th in-context example $\bm{x}_{t+1}$ after observing the first $t$ examples (see Figure 2). We write ${\mathcal{A}}(\bm{x}_{t+1}):={\mathcal{A}}(\bm{x}_{t+1}\mid\{\bm{x}_{i},y_{i}\}_{i=1}^{t})$ for brevity. The errors (i.e., residuals) on the data sequence are:the indices start from 2 to $n+1$ because we evaluate all cases where $t$ can choose from $1,\cdots,n$.

For any two algorithms ${\mathcal{A}}_{a}$ and ${\mathcal{A}}_{b}$, their similarity of errors, corresponding to the metric ${\mathcal{C}}(\cdot,\cdot)$, is

where we use the cosine similarity as our correlation metric ${\mathcal{C}}(\bm{u},\bm{v})=\dfrac{\left\langle{\bm{u},\bm{v}}\right\rangle}{\|\bm{u}\|_{2}\|\bm{v}\|_{2}}$. Here $n$ is the total number of in-context examples and $P_{\mathcal{D}}$ is the data generation process discussed previously.

Then, the induced weight vector for ${\mathcal{A}}$ and these $t$ in-context examples is:

4.2 Best Matching Hyper-parameters Between Algorithms

Each algorithm we consider has its own hyper-parameter(s), for example the number of iterations for Iterative Newton and Gradient Descent, and the number of layers for Transformers (see Section 4.1). When comparing two algorithms, given a choice of hyper-parameters for the first algorithm, we compare with the hyper-parameters for the second algorithm that maximize algorithmic similarity. In other words, we measure whether there exists hyperparameters for the second algorithm that make the two algorithms are similar.

Let $\mathcal{M}$ be the metric for evaluating similarities between two algorithms ${\mathcal{A}}_{a}$ and ${\mathcal{A}}_{b}$, which have hyper-parameters $p_{a}\in\mathcal{P}_{a}$ and $p_{b}\in\mathcal{P}_{b}$, respectively. For a given choice of $p_{a}$, We define the best-matching hyper-parameters of algorithm ${\mathcal{A}}_{b}$ for ${\mathcal{A}}_{a}$ as:

The matching processes can be visualized as heatmaps as shown in Figure 3, where best-matching hyperparameters are highlighted. This enables us to better compare the rate of convergence of algorithms and we will discuss these results further in §4.

Experimental Evidence

We mainly work on the Transformers-based GPT-2 model with 12 layers with 8 heads per layer. Alternative configurations with fewer heads per layer also support our findings and we defer them to Appendix A. We initially focus on isotropic cases where $\bm{\Sigma}=\bm{I}$ and later consider ill-conditioned $\bm{\Sigma}$ in §4.3. Our training setup is exactly the same as Garg et al. (2022): models are trained with at most $n=40$ in-context examples for $d=20$ (with the same learning rate, batch size etc.).

We claim that Transformers learn high-order optimization methods in-context. We provide evidence that Transformers improve themselves with more layers in §4.1; Transformers share the same rate of convergence as Iterative Newton, which is exponentially faster than that of Gradient Descent, in §4.2; and they also perform well on ill-conditioned problems in §4.3. Finally, we contrast Transformers with LSTMs in §4.4.

Many known algorithms for linear regression, including GD, OGD, and Iterative Newton, are iterative: their performance progressively improves as they perform more iterations, eventually converging to a final solution. How could a Transformer implement such an iterative algorithm? von Oswald et al. (2022) propose that deeper layers of the Transformer may correspond to more iterations of an iterative method; in particular, they show that there exist Transformer parameters such that each attention layer performs one step of Gradient Descent.

2 Transformers are more similar to Iterative Newton’s Method

In contrast, even though Gradient Descent has a comparable similarity with the Transformers at later layers, their best matching follows an exponential trend. As discussed in the Section 3.1, for well-conditioned problems where $\kappa\approx 1$, to achieve $\epsilon$ error, the rate of convergence of Gradient Descent is $\mathcal{O}(\log(1/\epsilon))$ while the rate of convergence of Iterative Newton is $\mathcal{O}(\log\log(1/\epsilon))$. Therefore the rate of convergence of Iterative Newton is exponentially faster than Gradient Descent. Transformer’s linear correspondence with Iterative Newton and its exponential correspondence with Gradient Descent provides strong evidence that the rate of convergence of Transformers is similar to Iterative Newton, i.e., $\mathcal{O}(\log\log(1/\epsilon))$. We also note that it is not possible to significantly improve Gradient Descent’s rate of convergence without using second-order methods: Nemirovski and Yudin (1983) showed a \Omega\big{(}\log(1/\epsilon)\big{)} lower bound on the rate of convergence of gradient-based methods for smooth and strongly convex problems.

Overall, we conclude that a Transformer trained to perform in-context linear regression learns to implement an algorithm that is very similar to Iterative Newton’s method, not Gradient Descent. Starting at layer 3, subsequent layers of the Transformer compute more and more iterations of Iterative Newton’s method. This algorithm successfully solves the linear regression problem, as it converges to the optimal OLS solution in the final layers.

3 Transformers perform well on ill-conditioned data

We repeat the same experiments with data $\bm{x}_{i}\overset{\text{i.i.d.}}{\sim}\mathcal{N}(\bm{0},\bm{\Sigma})$ sampled from an ill-condition covariance matrix $\bm{\Sigma}$ with condition number $\kappa(\bm{\Sigma})=100$, and eigenbasis chosen uniformly at random. The first $d/2$ eigenvalues of $\bm{\Sigma}$ are 100, the last $d/2$ are 1.

As shown in Figure 4, the Transformer model’s performance still closely matches Iterative Newton’s Method with 21 iterations, same as when $\bm{\Sigma}=\bm{I}$ (see layer 10-12 in Figure 3). The convergence of higher-order methods has a mild logarithmic dependence on the condition number since they correct for the curvature. On the other hand, Gradient Descent’s convergence is affected polynomially by conditioning. As $\kappa(\bm{\Sigma})$ increase from 1 to 100, the number steps required for GD’s convergence increases significantly (see Fig. 4 where GD requires 800 steps to converge), making it impossible for a 12-layer Transformers to implement these many gradient updates. We also note that preconditioning the data by $(\bm{X}^{\top}\bm{X})^{\dagger}$ can make the data well-conditioned, but since the eigenbasis is chosen uniformly at random, with high probability there is no sparse pre-conditioner or any fixed pre-conditioner which works across the data distribution. Computing $(\bm{X}^{\top}\bm{X})^{\dagger}$ appears to be as hard as computing the OLS solution (Eq. 1)—in fact Sharan et al. (2019) conjecture that first-order methods such as gradient descent and its variants cannot avoid polynomial dependencies in condition number in the ill-conditioned case.

See Section A.2 for detailed experiments on ill-conditioned problems. These experiments on ill-conditioned data further strengthen our hypothesis that Transformers are learning to perform higher-order optimization methods in-context, not Gradient Descent.

4 LSTM is more similar to OGD than Transformers

As discussed in Section 3, LSTM is an alternative auto-regressive model widely used before the introduction of Transformers. Thus, a natural research question is: If Transformers can learn in-context, can LSTMs do so as well? If so, do they learn the same algorithms? To answer this question, we train a 10-layer LSTM model, with 5.3M parameters, in an identical manner to the Transformers (with 9.5M parameters) studied in the previous sections.While the LSTM has fewer parameters than the Transformer, we found in preliminary experiments that increasing the hidden dimension or number of layers in the LSTM would not substantively change our results.

Figure 5 plots the mean squared error of Transformers, LSTMs, and other standard methods as a function of the number of in-context (i.e., training) examples provided. While LSTMs can also learn linear regression in-context, they have much higher mean-squared error than Transformers. Their error rate is similar to Iterative Newton’s Method after only 5 iterations, a point where it is far from converging to the OLS solution.

LSTMs’ inferior performance to Transformers can be explained by the inability of LSTMs to use deeper layers to improve their predictions. Figure 1 shows that LSTM performance does not improve across layers—a readout head fine-tuned for the first layer makes equally good predictions as the full 10-layer model. Thus, LSTMs seem poorly equipped to fully implement iterative algorithms.

Finally, we show that LSTMs behave more like an online learning algorithm than Transformers. In particular, its predictions are biased towards getting more recent training examples correct, as opposed to earlier examples, as shown in Figure 5. This property makes LSTMs similar to online Gradient Descent. In contrast, five steps of Newton’s method has the same error on average for recent and early examples, showing that the LSTM implements a very different algorithm from a few iterations of Newton. Similarly, Figure 5 shows that LSTMs are more similar to OGD than Transformers are, whereas Transformers are more similar to Newton and GD than LSTMs. We hypothesize that since LSTMs have limited memory, they must learn in a roughly online fashion; in contrast, Transformers’ attention heads can access the entire sequence of past examples, enabling it to learn more complex algorithms.

Mechanistic Evidence

Our empirical evidence demonstrates that Transformers behave much more similarly to Iterative Newton’s than to Gradient Descent. Iterative Newton is a higher-order optimization method, and is algorithmically more involved than Gradient Descent. We begin by first examining this difference in complexity. As discussed in Section 3, the updates for Iterative Newton are of the form,

Theorem 1 shows that this is indeed possible.

for some $\alpha>0$ and $\bm{S}=\bm{X}^{\top}\bm{X}$. The number of layers of the transformer is $\mathcal{O}(k)$ and the dimensionality of the hidden layers is $\mathcal{O}(d)$.

Here we provide a skecth of the proof. We note that our proof uses full attention instead of causal attention and ReLU activations for the self-attention layers. The definitions of these and the full proof appear in Appendix B.

Conclusion and Discussion

In this work, we studied how Transformers perform in-context learning for linear regression. In contrast with the hypothesis that Transformers learn in-context by implementing gradient descent, our experimental results show that different Transformer layers match iterations of Iterative Newton’s method linearly and Gradient Descent exponentially. This suggests that Transformers share a similar rate of convergence to Iterative Newton’s method but not Gradient Descent. Moreover, Transformers can perform well empirically on ill-conditioned linear regression, whereas first-order methods such as Gradient Descent struggle. Theoretically, we provide exact Transformer circuits that can implement $k$-steps of Iterative Newton’s method with $\mathcal{O}(k)$ layers to make predictions from in-context examples.

An interesting direction is to explore a wider range of higher-order methods that Transformers can implement. It also seems promising to extend our analysis to classification problems, especially given recent work showing that Transformers resemble SVMs in classification tasks (Li et al., 2023; Tarzanagh et al., 2023a). Finally, a natural question is to understand the differences in the model architecture that make Transformers better in-context learners than LSTMs. Based on our investigations with LSTMs, we hypothesize that Transformers can implement more powerful algorithms because of having access to a longer history of examples. Investigating the role of this additional memory in learning appears to be an intriguing direction.

Acknowledgement

We would like to thank the USC NLP Group and Center for AI Safety for providing compute resources. DF would like to thank Oliver Liu and Ameya Godbole for extensive discussions. RJ was supported by an Open Philanthropy research grant and a Google Research Scholar Award. VS was supported by NSF CAREER Award CCF-2239265 and an Amazon Research Award.

References

Appendix

Appendix A Additional Experimental Results

We present heatmaps with all values of similarities.

A.1.2 Additional Results on Comparison over Transformer Layers

A.1.3 Additional Results on Similarity of Induced Weights

A.2 Experiments on Ill-Conditioned Problems

In this section, we repeat the same experiments as we did on isotropic data in the main text and in Section A.1, and we change the covariance matrix to be ill-conditioned such that $\kappa(\bm{\Sigma})=100$.

We also present the heatmaps to find the best-matching hyper-parameters and conclude that Transformers are similar to Newton’s method than GD in ill-conditioned data.

A.3 Experiments on alternative configurations

In this section, we present experimental results from an alternative model configurations than the main text. We show in the main text that Transformers learn higher-order optimization methods in-context where the experiments are using a GPT-2 model with 12 layers and 8 heads per layer. In this section, we present experiments with a GPT-2 model with 12 layers but only 1 head per layer.

We conclude that our experimental results are not restricted to a specific model configurations, smaller models such as GPT-2 with 12 layers and 1 head each layer also suffice in implementing the Iterative Newton’s method, and more similar than gradient descents, in terms of rate of convergence.

A.4 Definitions for Evaluating Forgetting

Note: the further away $i$ is from $n$, the more possible algorithm ${\mathcal{A}}$ forgets.

Appendix B Detailed Proofs for Section 5

This is slightly different from the causal attention layer (see eq. 6) in that at each time stamp $t$, the attention layer in definition 5 has full information of all hidden states $j\in[n]$, unlike causal attention layer which requires $j\in[t]$.

We begin by constructing a useful component for our proof, and state some existing constructions from Akyürek et al. .

Given hidden states $\{\bm{h}_{1},\cdots,\bm{h}_{n}\}$, there exists query, key and value matrices $\bm{Q},\bm{K},\bm{V}$ respectively such that one attention layer can compute $\sum_{j=1}^{n}\bm{h}_{j}$.

Let $\bm{V}_{1}=\bm{V}_{2}=\begin{bmatrix}\bm{O}_{(d+1)\times d}&\bm{O}_{(d+1)\times(d+1)}\\ n\bm{I}_{d\times d}&\bm{O}_{d\times(d+1)}\end{bmatrix}$ so that $\bm{V}_{m}\begin{bmatrix}\bm{h}_{j}\\ 1\\ \bm{0}_{d}\end{bmatrix}=\begin{bmatrix}\bm{0}_{d+1}\\ n\bm{h}_{j}\end{bmatrix}$.

We apply one attention layer to these 1-padded hidden states and we have

mov$(\bm{H};s,t,i,j,i^{\prime},j^{\prime})$: selects the entries of the $s$-th column of $\bm{H}$ between rows $i$ and $j$, and copies them into the $t$-th column ($t\geq s$) of $\bm{H}$ between rows $i^{\prime}$ and $j^{\prime}$.

mul$(\bm{H};a,b,c,(i,j),(i^{\prime},j^{\prime}),(i^{\prime\prime},j^{\prime\prime}))$: in each column $\bm{h}$ of $\bm{H}$, interprets the entries between $i$ and $j$ as an $a\times b$ matrix $\bm{A}_{1}$, and the entries between $i^{\prime}$ and $j^{\prime}$ as a $b\times c$ matrix $\bm{A}_{2}$, multiplies these matrices together, and stores the result between rows $i^{\prime\prime}$ and $j^{\prime\prime}$, yielding a matrix in which each column has the form $\begin{bmatrix}\bm{h}_{:i^{\prime\prime}-1},\bm{A}_{1}\bm{A}_{2},\bm{h}_{j^{\prime\prime}:}\end{bmatrix}^{\top}$. This allows the layer to implement inner products.

div$(\bm{H};(i,j),i^{\prime},(i^{\prime\prime},j^{\prime\prime}))$: in each column $\bm{h}$ of $\bm{H}$, divides the entries between $i$ and $j$ by the absolute value of the entry at $i^{\prime}$, and stores the result between rows $i^{\prime\prime}$ and $j^{\prime\prime}$, yielding a matrix in which every column has the form $\begin{bmatrix}\bm{h}_{:i^{\prime\prime}-1},\bm{h}_{i:j}/|\bm{h}_{i^{\prime}}|,\bm{h}_{j^{\prime\prime}:}\end{bmatrix}^{\top}$.

aff$(\bm{H};(i,j),(i^{\prime},j^{\prime}),(i^{\prime\prime},j^{\prime\prime}),\bm{W}_{1},\bm{W}_{2},\mathbf{b})$: in each column $\bm{h}$ of $\bm{H}$, applies an affine transformation to the entries between $i$ and $j$ and $i^{\prime}$ and $j^{\prime}$, then stores the result between rows $i^{\prime\prime}$ and $j^{\prime\prime}$, yielding a matrix in which every column has the form $\begin{bmatrix}\bm{h}_{:i^{\prime\prime}-1},\bm{W}_{1}\bm{h}_{i:j}+\bm{W}_{2}\bm{h}_{i^{\prime}:j^{\prime}}+\mathbf{b},\bm{h}_{j^{\prime\prime}:}\end{bmatrix}^{\top}$. This allows the layer to implement subtraction by setting $\bm{W}_{1}=\bm{I}$ and $\bm{W}_{2}=-\bm{I}$.

B.2 Proof of Theorem 1

We call each column after mov as $\bm{h}_{j}$. With an full attention layer, one can construct two heads with query and value matrices of the form $\bm{Q}_{1}^{\top}\bm{K}_{1}=-\bm{Q}_{2}^{\top}\bm{K}_{2}=\begin{bmatrix}\bm{I}_{d\times d}&\bm{O}_{d\times d}\\ \bm{O}_{d\times d}&\bm{O}_{d\times d}\end{bmatrix}$ such that for any $t\in[n]$, we have

Now we can use the aff operator to make subtractions and then

We call this transformed hidden states as $\bm{H}^{(0)}$ and denote $\bm{T}^{(0)}=\alpha\bm{S}$:

Notice that $\bm{S}$ is symmetric and thereafter $\bm{T}^{(0)}$ is also symmetric.

Let the input prompt be the same as Equation 27,

This is exactly the same as the Newton iteration.

Now we apply Lemma 1 over all even columns in Equation 36 and we have

Initialization: mov needs $\mathcal{O}(1)$ layer; gathering $\alpha\bm{S}$ needs $\mathcal{O}(1)$ layer; and aff needs $\mathcal{O}(1)$ layer. In total, Transformers need $\mathcal{O}(1)$ layers for initialization.

Newton Iteration: each exact Newton’s iteration requires $\mathcal{O}(1)$ layer. Implementing $k$ iterations requires $\mathcal{O}(k)$ layers.

Making prediction on test query: We need one operation of mov and mul each, requiring $\mathcal{O}(1)$ layer each.

Hence, in total, Transformers can implement $k$-step Iterative Newton and make predictions accordingly using $\mathcal{O}(k)$ layers.

B.3 Iterative Newton as a sum of moments method

Recall that Iterative Newton’s method finds $\bm{S}^{\dagger}$ as follows

We can expand the iterative equation to moments of $\bm{S}$ as follows.

Let’s do this one more time for $\bm{M}_{2}$.

We can see that $\bm{M}_{k}$ are summations of moments of $\bm{S}$, with respect to some pre-defined coefficients from the Newton’s algorithm. Hence Iterative Newton is a special of an algorithm which computes an approximation of the inverse using higher-order moments of the matrix,

We note that Transformer circuits can represent other sum of moments other than Newton’s method. We can introduce different coefficients $\beta_{i}$ than in the proof of Theorem 1 by scaling the value matrices or through the MLP layers.

B.4 Estimated weight vectors lie in the span of previous examples

What properties can we infer and verify for the weight vectors which arise from Newton’s method? A straightforward one arises from interpreting any sum of moments method as a kernel method.

We test this hypothesis. Given a sequence of in-context examples $\{\bm{x}_{i},y_{i}\}_{i=1}^{t}$, we fit coefficients $\{a_{i}\}_{i=1}^{t}$ in Equation 45 to minimize MSE loss:

We then measure the quality of this fit across different number of in-context examples $t$, and visualize the residual error in Figure 16. We find that even when $t<d$, Transformers’ induced weights still lie close to the span of the observed examples $\bm{x}_{i}$’s. This provides an additional validation of our proposed mechanism.