SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kwangjun Ahn*
MIT EECS/LIDS
MIT LIDS
Hadi Daneshmand*
MIT LIDS/FODSI
hdanesh@mit.edu
Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity, we ask: Can transformers learn to implement such algorithms by training over random problem instances? To our knowledge, we make the first theoretical progress on this question via an analysis of the loss landscape for linear transformers trained over random instances of linear regression. For a single attention layer, we prove the global minimum of the training objective implements a single iteration of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the input distribution but also to the variance induced by data inadequacy. For a transformer with  $L$  attention layers, we prove certain critical points of the training objective implement  $L$  iterations of preconditioned gradient descent. Our results call for future theoretical studies on learning algorithms by training transformers.
In-context learning (ICL) is the striking capability of large language models: Given a prompt containing examples and a query, the transformer produces the correct output based on the context provided by the examples, without adapting its parameters (Brown et al., 2020; Lieber et al., 2021; Rae et al., 2021; Black et al., 2022). This property has become the focus of body of recent research that aims to shed light on the underlying mechanism of large language models (Garg et al., 2022; Akyurek et al., 2022; von Oswald et al., 2023; Li and Malik, 2017; Min et al., 2021; Xie et al., 2021; Elhage et al., 2021; Olsson et al., 2022).
A line of research studies ICL via the expressive power of transformers. Transformer architectures are powerful Turing machines, capable of implementing various algorithms (Pérez et al., 2021; Wei et al., 2022). Given an in-context prompt, Edelman et al. (2022); Olsson et al. (2022) argue that transformers are able to implement algorithms through the recurrence of multi-head attentions to extract coarse information from raw input prompts. Akyurek et al. (2022); von Oswald et al. (2023) assert that transformers can implement gradient descent on linear regression encoded in a given input prompt. It is thought provoking that transformers can implement such algorithms.
Although transformers are universal machines to implement algorithms, they need specific parameter configurations for achieving these implementations. In practice, their parameters are adjusted via training using non-convex optimization over random problem instances. Hence, it remains unclear whether this non-convex optimization can be used to learn algorithms. The present paper investigates the possibility of learning algorithms via training over random problem instances.
More specifically, we investigate the learning of gradient-based methods. It is hard to mathematically formulate what it means to learn gradient descent for general functions with transformers. Yet, Garg et al. (2022) elegantly examine it in the specific setting of ICL for learning functions. Empirical evidence suggests that transformers indeed learn to implement gradient descent, after training on random instances of linear regression (Garg et al., 2022; Akyurek et al., 2022; von Oswald et al., 2023). Motivated by these observations, we theoretically investigate the loss landscape of a simple transformer architecture based on attention without softmax (Schlag et al., 2021; von Oswald et al., 2023) (see Section 2 for details).
Summary of our main results. Our main contributions are the following:
- We provide a complete characterization of the global optimum of a single-layer linear transformer. In particular, we observe that, with the optimal parameters, the transformer implements a single step of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the distribution of input data but also to the variance caused by data inadequacy. We present this result in Theorem 1 in Section 3.
$\triangleright$  Next, we focus on a subset of the transformer parameter space, defined by a special sparsity condition (8). Such a parameter configuration allows us to formulate training transformers as a search over  $k$ -step adaptive gradient-based algorithms. Theorem 2 characterizes the global minimizers of the training objective of a two-layer linear transformer over isotropic regression instances, and shows that the optima correspond to gradient descent with adaptive step sizes. For multilayer transformers, Theorem 3 demonstrates that gradient descent, with a data-dependent preconditioning, can be derived from a critical point of the training objective.
Finally, we study the loss landscape in the absence of the sparsity condition (8), which goes beyond searching over conventional gradient-based optimization methods. In this case, we prove and interpret the structure of a critical point of the training objective. We show that a certain critical point in parameter space leads to an intriguing gradient-based algorithm that simultaneously takes gradient steps preconditioned by data covariance, and applies a linear transformation to further improve the conditioning. In the specific case when data covariance is isotropic, this algorithm corresponds to the  $\mathrm{GD}++$  algorithm of von Oswald et al. (2023) which is experimentally observed to be the outcome of training.
We empirically validate the critical points analyzed in Theorem 3 and Theorem 4. For a transformer with three layers, our experimental results confirm the structural of critical points. Furthermore, we observed the objective value associated with these critical points is close to 0, suggesting that the critical points might be global optima. These experiments substantiate our theoretical analysis and suggests that our theory indeed aligns with practice. Code for our experiments is available at https://github.com/chengxiang/LinearTransformer.
The ability of neural network architectures to implement algorithms has been investigated in various context. The seminal work by Siegelmann and Sontag (1992) investigate the Turing completeness of recurrent neural networks. Despite this computational power, training recurrent networks remains a challenge. Graves et al. (2014) design an alternative neural architecture known as the neural Turing machine, building on attention layers introduced by Hochreiter and Schmidhuber (1997). Leveraging attention, Vaswani et al. (2017) propose transformers as powerful neural architectures, capable of solving various tasks in natural language processing (Devlin et al., 2019). This capability inspired a line of research that examines the algorithmic power of transformers (Pérez et al., 2021; Wei et al., 2022; Giannou et al., 2023; Akyurek et al., 2022; Olsson et al., 2022). What sets transformers apart from conventional neural networks is their impressive performance after training. In this work, we focus on understanding how transformers learn to implement algorithms by training over problem instances.
A line of research investigates how deep neural networks process data across their layers. The seminal work by Jastrzebski et al. (2018) observes that hidden representations across the layers of deep neural networks approximately implement gradient descent. Recent observations provide novel insights into the working mechanism of ICL for large language models, showing they can implement optimization algorithms across their layers (Garg et al., 2022; Akyurek et al., 2022; von Oswald et al., 2023). Moreover, Zhao et al. (2023); Allen-Zhu and Li (2023) observe transformer perform
dynamic programming to generate text. In this work, we theoretically study how transformer learns gradient-based algorithms for ICL.
We discuss here two related works (Zhang et al., 2023; Mahankali et al., 2023) that appeared shortly after publication of our original draft. Both of these studies focus on a single layer attention network (see Section 3). Zhang et al. (2023) prove the global convergence of gradient descent to the global optimum whose structure is analyzed independently from this study and it the same as that in Theorem 1. Mahankali et al. (2023) also characterize the global minimizer of a single layer attention without softmax for a different data distribution. In addition to results for a single-layer attention, we analyze the landscape of two and multi-layer transformers.
2 Setting: training linear transformers over random linear regression
In order to understand the mechanism of ICL, we consider the setting of training transformers over the random instances of linear regression, following (Garg et al., 2022; Akyurek et al., 2022; von Oswald et al., 2023). In particular, the random instances of linear regression are formalized as follows.
Data distribution: random linear regression instances. Let  $x^{(i)} \in \mathbb{R}^d$  be the covariates drawn i.i.d. from a distribution  $D_{\mathcal{X}}$ , and  $w_{\star} \in \mathbb{R}^{d}$  be drawn from  $D_{\mathcal{W}}$ . Let  $X \in \mathbb{R}^{(n + 1)\times d}$  be the matrix of covariates whose row  $i$  contains tokens  $x^{(i)}$ . Given  $x^{(i)}$ 's and  $w_{\star}$ , the responses are defined as  $y = [\langle x^{(1)},w_{\star}\rangle ,\dots ,\langle x^{(n)},w_{\star}\rangle ] \in \mathbb{R}^n$ . Define the input matrix  $Z_{0}$  as
$$
Z _ {0} = \left[ z ^ {(1)} z ^ {(2)} \dots z ^ {(n)} z ^ {(n + 1)} \right] = \left[ \begin{array}{c c c c c} x ^ {(1)} & x ^ {(2)} & \dots & x ^ {(n)} & x ^ {(n + 1)} \\ y ^ {(1)} & y ^ {(2)} & \dots & y ^ {(n)} & 0 \end{array} \right] \in \mathbb {R} ^ {(d + 1) \times (n + 1)}, \tag {1}
$$
where zero in the above matrix is used to replace the unknown response variable corresponding to  $x^{(n + 1)}$ . Then, our goal is to predict  $w_{\star}^{\top}x^{(n + 1)}$  given  $Z_{0}$ . In other words, the training data consists of pairs  $(Z_0,w_\star^\top x^{(n + 1)})$  for  $x^{(i)}\sim D_{\mathcal{X}}$  and  $w_{\star}\sim D_{\mathcal{W}}$ . We then consider training transformers over this data distribution.
Self-attention layer without softmax. Following (Schlag et al., 2021; von Oswald et al., 2023), we consider the linear self-attention layer. To motivate, we first briefly review the standard self-attention layer (Vaswani et al., 2017). Letting  $Z \in \mathbb{R}^{(d + 1) \times (n + 1)}$  be the input matrix with  $n + 1$  tokens in  $\mathbb{R}^{d + 1}$ , a single-head self-attention layer denoted by  $\text{Attn}^{\text{max}}$  is a parametric map defined as
$$
\operatorname {A t t n} _ {W _ {k, q, v}} ^ {\mathsf {s m a x}} (Z) = W _ {v} Z M \cdot \mathsf {s m a x} \left(Z ^ {\top} W _ {k} ^ {\top} W _ {q} Z\right), \quad M := \left[ \begin{array}{c c} I _ {n} & 0 \\ 0 & 0 \end{array} \right] \in \mathbb {R} ^ {(n + 1) \times (n + 1)}, \tag {2}
$$
where  $W_{v}, W_{k}, W_{q} \in \mathbb{R}^{(d + 1) \times (d + 1)}$  are the (value, key and query) weight matrices, and  $\mathrm{smax}(\cdot)$  is the softmax operator which applies softmax operation to each column of the input matrix. Note that the prompt is asymmetric since the label for  $x^{(n + 1)}$  is excluded from the input. To reflect this asymmetric structure, the mask matrix  $M$  is included in the attention. In our setting, we consider the self-attention layer that omits the softmax operation in (2). In particular, we reparameterize weights as  $P := W_{v} \in \mathbb{R}^{(d + 1) \times (d + 1)}$  and  $Q := W_{k}^{\top} W_{q} \in \mathbb{R}^{(d + 1) \times (d + 1)}$  and consider
$$
\operatorname {A t t n} _ {P, Q} (Z) = P Z M \left(Z ^ {\top} Q Z\right). \tag {3}
$$
At first glance, the omission of the softmax operation (3) might seem over-simplified. But, (von Oswald et al., 2023) proves such attention can implement gradient descent, and we will prove in Lemma 1 that it can also implement various algorithms to solve linear regression in-context.
Architecture for prediction. We now present the neural network architecture that will be used throughout this paper. For the number of layers  $L$ , we define an  $L$ -layer transformer as a stack of  $L$  linear self-attention blocks. Formally, denoting by  $Z_{\ell}$  the output of the  $\ell^{\mathrm{th}}$  layer attention, we define
$$
Z _ {\ell + 1} = Z _ {\ell} + \frac {1}{n} \operatorname {A t t n} _ {P _ {\ell}, Q _ {\ell}} \left(Z _ {\ell}\right) \quad \text {f o r} \ell = 0, 1, \dots , L - 1, \tag {4}
$$
The scaling factor  $\frac{1}{n}$  is used only for ease of notation and does not influence the expressive power of the transformer. Given  $Z_{L}$ , we define  $\mathsf{TF}_L(Z_0; \{P_\ell, Q_\ell\}_{\ell = 0,1,\dots L - 1}) = -[Z_L]_{(d + 1),(n + 1)}$ , i.e., the  $(d + 1,n + 1)$ -th entry of  $Z_{L}$ . The reason for the minus sign is to be consistent with (von Oswald
et al., 2023), and we will clarify such a choice in Lemma 1. For training, the parameters are optimized to minimize in-context loss as
$$
f \left(\left\{P _ {\ell}, Q _ {\ell} \right\} _ {\ell = 0} ^ {L}\right) = \mathbb {E} _ {\left(Z _ {0}, w _ {\star}\right)} \left[ \left(\mathsf {T F} _ {L} \left(Z _ {0}, \left\{P _ {\ell}, Q _ {\ell} \right\} _ {\ell = 0} ^ {L}\right) + w _ {\star} ^ {\top} x ^ {(n + 1)}\right) ^ {2} \right]. \tag {5}
$$
Goal: the landscape analysis of the training objective functions. We are interested in understanding how the optimization of  $f$  leads to in-context learning. We investigate this question by analyzing its loss landscape. Such analysis is challenging due to two major reasons: (i)  $f$  is non-convex in parameters  $\{P_i, Q_i\}$  even for a single layer transformer. (ii) The cross-product structures in attention makes  $f$  a highly nonlinear function in its parameters. Hence, we analyze a spectrum of settings from single-layer transformers to multi-layer transformers. For simpler settings such as single-layer transformers, we prove stronger results such as the full characterization of the global minimizers. For networks with more layers, we characterize the structure of critical points. Furthermore, we provide algorithmic interpretations of the critical points. Table 1 summarizes our results for various parametric models.
Results	x(i)	w*	Setting	Guarantees
Theorem 1	N(0,Σ)	N(0,I)	single-layer	global minimizers
Theorem 2	N(0,I)	N(0,I)	two-layer + symmetric (8)	global minimizers
Theorem 3	N(0,Σ)	N(0,Σ-1)	multi-layer + (8)	critical points
Theorem 4	N(0,Σ)	N(0,Σ-1)	multi-layer + (11)	critical points
Theorem 5	N(0,I)	N(0,I)	single-layer + ReLU activation	global minimizers
Table 1: Summary of our analyses for various models and input distributions. The additional conditions (8) and (11) are about the sparsity structure of parameters. In addition, "symmetric (8)" means we additionally impose the weights to be symmetric.
Remark 1 (Optimizing (5) vs. practical transformer optimization). Interestingly, a recent work by Ahn et al. (2023) reports that common optimization algorithms such as SGD/ADAM behave remarkably similarly on the (linear Transformers + linear regression) problem as they do on (practical transformers + real language modeling tasks). In particular, they reproduce several distinctive features of transformer optimization under a simple shallow linear transformer. This work suggests that (linear transformer + linear regression) may serve as a good proxy for understanding practical transformer optimization.
3 The global optimum for a single-layer transformer
For the single layer case of  $L = 1$ , the following result characterizes the optimal parameters  $P_0$  and  $Q_0$  for the in-context loss (5).
Theorem 1 (Single-layer; non-isotropic data). Assume that vector  $x^{(i)}$  is sampled from  $\mathcal{N}(0, \Sigma)$ , i.e., a Gaussian with covariance  $\Sigma = U\Lambda U^{\top}$  where  $\Lambda = \mathrm{diag}(\lambda_1, \dots, \lambda_d)$ . Moreover, assume that  $w_{\star}$  is sampled from  $\mathcal{N}(0, I_d)$ . Then, the following choice of parameters
$$
P _ {0} = \left[ \begin{array}{c c} 0 _ {d \times d} & 0 \\ 0 & 1 \end{array} \right], \quad Q _ {0} = - \left[ \begin{array}{c c} U \operatorname {d i a g} \left(\left\{\frac {1}{\frac {n + 1}{n} \lambda_ {i} + \frac {1}{n} \cdot \left(\sum_ {k} \lambda_ {k}\right)} \right\} _ {i = 1, \dots , d}\right) U ^ {\top} & 0 \\ 0 & 0 \end{array} \right]. \tag {6}
$$
is a global minimizer of  $f(P, Q)$  up to re-scaling, i.e.,  $P_0 \gets \gamma P_0$  and  $Q_0 \gets \gamma^{-1} Q_0$  for a scalar  $\gamma$ .
See Appendix A for the proof of Theorem 1. In the specific case when the Gaussian is isotropic, i.e.,  $\Sigma = I_d$ , the optimal  $Q_0$  has the following simple form
$$
Q _ {0} = - \frac {1}{\left(\frac {n - 1}{n} + (d + 2) \frac {1}{n}\right)} \left[ \begin{array}{l l} I _ {d} & 0 \\ 0 & 0 \end{array} \right]. \tag {7}
$$
Up to scaling, the above parameter configuration is equivalent to the parameters used by von Oswald et al. (2023) to perform one step of gradient descent. Thus, in the single-layer setting, the in-context loss is indeed minimized by a transformer that implements the gradient descent algorithm.
More generally, when the in-context samples are non-isotropic, the transformer learns to implement one step of a preconditioned gradient descent as we shall detail in Lemma 1. Here the "preconditioning matrix" given in (6) has interesting properties:
- When the number of samples  $n$  is large, the first  $d \times d$  submatrix of  $Q_0$  approximates  $\Sigma^{-1}$ , the inverse of the data covariance matrix, which is also close to the Gram matrix formed from  $x^{(1)}, \ldots, x^{(n)}$ . Hence the preconditioning can lead to considerably faster convergence rate when  $\Sigma$  is ill-conditioned.
- Moreover,  $\frac{1}{n}\sum_{k}\lambda_{k}$  in (6) acts as a regularizer. It becomes more significant when  $n$  is small and variance of the  $x^{(i)}$ 's is high. Such an adjustment resembles structural risk minimization (Vapnik, 1999) where the regularization strength is adapted to the sample size.
4 Multi-layer transformers with sparse parameters
Theorem 1 proves a single layer of linear attention can implement a single step of preconditioned gradient descent. Inspired by this result, we investigate the algorithmic power of the linear transformer architecture. We show that the model can implement various optimization methods even under sparsity constraints. In particular, we impose the following restrictions on the parameters:
$$
P _ {i} = \left[ \begin{array}{c c} 0 _ {d \times d} & 0 \\ 0 & 1 \end{array} \right], \quad Q _ {i} = - \left[ \begin{array}{c c} A _ {i} & 0 \\ 0 & 0 \end{array} \right] \quad \text {w h e r e} A _ {i} \in \mathbb {R} ^ {d \times d}. \tag {8}
$$
The next lemma proves that a forward-pass of a  $L$ -layer transformer, with the parameter configuration (8) is the same as taking  $L$  steps of gradient descent, preconditioned by  $A_{\ell}$ .
Lemma 1 (Forward pass as a preconditioned gradient descent). Consider the  $L$ -layer linear transformer parameterized by  $A_0, \ldots, A_{L-1}$  as in (8). Let  $y_\ell^{(n+1)}$  be the  $(d+1, n+1)$ -th entry of the  $\ell$ -th layer output, i.e.,  $y_\ell^{(n+1)} = [Z_\ell]_{(d+1), (n+1)}$  for  $\ell = 1, \ldots, L$ . Then, it holds that  $y_\ell^{(n+1)} = -\langle x^{(n+1)}, w_\ell^{\mathrm{gd}} \rangle$  where  $\{w_\ell^{\mathrm{gd}}\}$  is defined as  $w_0^{\mathrm{gd}} = 0$  and as follows for  $\ell = 1, \ldots, L-1$ :
$$
w _ {\ell + 1} ^ {\mathbf {g d}} = w _ {\ell} ^ {\mathbf {g d}} - A _ {\ell} \nabla R _ {w _ {\star}} \left(w _ {\ell} ^ {\mathbf {g d}}\right) \quad \text {w h e r e} \quad R _ {w _ {\star}} (w) := \frac {1}{2 n} \sum_ {i = 1} ^ {n} \left(w ^ {\top} x _ {i} - w _ {\star} ^ {\top} x _ {i}\right) ^ {2}. \tag {9}
$$
See Subsection C.1 for a proof. The iterative scheme (9) includes various optimization methods including gradient descent with  $A_{\ell} = \gamma_{\ell}I_{d}$ , and (adaptive) preconditioned gradient descent, where the preconditioner  $A_{\ell}$  depends on the time step. In the upcoming sections, we characterize how the optimal  $\{A_{\ell}\}$  are linked to the input distribution.
4.1 Warm-up: optimal two-layer transformer with symmetric weights
For the rest of this section, we will study the optimal parameters for the in-context loss under the constraint of Eq. (8). Later in Section 5, we analyze the optimal model for a more general parameters. For a two-layer transformer, the next Theorem proves the optimal in-context loss obtains the simple gradient descent with adaptive coordinate-wise stepsizes.
Theorem 2 (Global optimality for the two-layer (symmetric) transformer). Consider the optimization of in-context loss for a two-layer transformer with the parameter configuration in Eq. (8), and additionally assume that  $A_{1}, A_{2}$  are symmetric matrices. More formally, consider
$$
\min  _ {A _ {1}, A _ {2} \text {a r e s y m m e t r i c}} f \left\{P _ {\ell} = \left[ \begin{array}{c c} 0 _ {d \times d} & 0 \\ 0 & 1 \end{array} \right], Q _ {\ell} = \left[ \begin{array}{c c} - A _ {\ell} & 0 \\ 0 & 0 \end{array} \right] \right\} _ {\ell = 1, 2}.
$$
Assume  $x^{(i)} \stackrel{i.i.d.}{\sim} N(0, I_d)$  and  $w_\star \sim N(0, I_d)$ ; then, there are diagonal matrices  $A_1$  and  $A_2$  that are a global minimizer of  $f$ .
Combining the above result with Lemma 1 concludes that the two iterations of gradient descent with coordinate-wise adaptive step sizes achieve the minimal in-context loss for isotropic Gaussian inputs. Gradient descent with adaptive step sizes such as Adagrad (Duchi et al., 2011) are widely used in machine learning. While Adagrad adjusts its stepsize based on the individual problem instance, the algorithm learned adjusts its stepsize to the underlying data distribution.
4.2 Multi-layer transformers
We now turn to the setting of general  $L$ -layer transformers, for any positive integer  $L$ . The next theorem proves that certain critical points of the in-context loss effectively implement a specific preconditioned gradient algorithm, where the preconditioning matrix is the inverse covariance of the input distribution. Before stating this result, let us first consider a motivating scenario in which the data-covariance matrix is non-identity:
Linear regression with distorted view of the data: Suppose that  $\overline{w}_{\star} \sim \mathcal{N}(0, I)$  and the latent covariates are  $\overline{x}^{(1)}, \ldots, \overline{x}^{(n+1)}$ , drawn i.i.d from  $\mathcal{N}(0, I)$ . We are given  $y^{(1)}, \ldots, y^{(n)}$ , with  $y^{(i)} = \langle \overline{x}^{(i)}, \overline{w}_{\star} \rangle$ . However, we do not observe the latent covariates  $\overline{x}^{(i)}$ . Instead, we observe the distorted covariates  $x^{(i)} = W \overline{x}^{(i)}$ , where  $W \in \mathbb{R}^{d \times d}$  is a distortion matrix. Thus the prompt consists of  $(x^{(1)}, y^{(1)}), \ldots, (x^{(n)}, y^{(n)})$ , as well as  $x^{(n+1)}$ . The goal is still to predict  $y^{(n+1)}$ . Note that this setting is quite common in practice, when covariates are often represented in an arbitrary basis.
Assume that  $\Sigma := WW^{\top} \succ 0$ . We verify from our definitions that for  $w_{\star} := \Sigma^{-1/2}\overline{w}_{\star}, y^{(i)} = \langle x^{(i)}, w_{\star} \rangle$ . Furthermore,  $x^{(i)} \sim \mathcal{N}(0, \Sigma)$  and  $w_{\star} \sim \mathcal{N}(0, \Sigma^{-1})$ . From Lemma 1, the transformer with weight matrices  $\{A_0, \ldots, A_{L-1}\}$  implements preconditioned gradient descent with respect to  $R_{w_{\star}}(w) = \frac{1}{2n} (w - w_{\star})^T XX^{\top} (w - w_{\star})$ , with  $X = [x^{(1)}, \ldots, x^{(n)}]$ . Under this loss, the Hessian matrix  $\nabla^2 R_{w_{\star}}(w) = \frac{1}{2n} XX^{\top}$  (at least in the case of large  $n$ ). For any fixed prompt, Newton's method corresponds to  $A_i \propto (XX^{\top})^{-1}$ , which makes the problem well-conditioned even if  $\Sigma$  is very degenerate. As we will see in Theorem 3 below, the choice of  $A_i \propto \Sigma^{-1} = \mathbb{E}[XX^{\top}]^{-1}$  appears to be a stationary point of the loss landscape, in expectation over prompts.