The Closeness of In-Context Learning and Weight Shifting for Softmax Regression

Shuai Li, Zhao Song, Yu Xia, Tong Yu, Tianyi Zhou

Introduction

In recent years, there has been a significant increase in research and development in the field of Artificial Intelligence (AI), with large language models (LLMs) emerging as an effective way to tackle complex tasks. Transformers have achieved state-of-the-art results in various natural language processing tasks, such as machine translation , language modeling, question answering, and text generation . As a result, they have become the preferred architecture for NLP. Based on that architecture, BERT , GPT-3 , PaLM , and OPT were proposed. They have demonstrated remarkable learning and reasoning capabilities and have proven to be more efficient than smaller models and traditional techniques when processing natural language.

Additionally, LLMs can be fine-tuned for multiple purposes without requiring a new build from scratch, making them a versatile tool for AI applications. A prime example of this is ChatGPT, a chat software developed by OpenAI utilizing GPT-3’s potential to its fullest. Moreover, GPT-4 has the capability to handle intricate tasks that its predecessors were unable to accomplish. It has demonstrated a remarkable level of proficiency comparable to human performance in various professional and academic benchmarks.

Transformers have a specific type of sequence-to-sequence neural network architecture. They utilize the attention mechanism that allows them to capture long-range dependencies and context from input data effectively. The core of the attention mechanism is the attention matrix which is comprised of rows and columns, corresponding to individual words or “tokens”. The attention matrix represents the relationships within the given text. It measures the importance of each token in a sequence as it relates to the desired output. During the training process, the attention matrix is learned and optimized to improve the accuracy of the model’s predictions. Through the attention mechanism, each input token is evaluated based on its relevance to the desired output by assigning a token score. This score is determined by a similarity function that compares the current output state with input states.

Mathematically, the attention computation problem can be formulated as a regression problem in the following sense

Several theoretical transformer work have studied either exponential regression or softmax regression problem . In this work, to take a more step forward to understand the softmax unit in the attention scheme in LLMs. We consider the following softmax regression and study the in-context learning phenomena based on it.

We remark that the Definition 1.3 is a formulation is between Definition 1.2 and Definition 1.1.

We consider the softmax regression (Definition 1.3) problem

Part 1. If we move the $x_{t}$ to $x_{t+1}$ , then we’re solving a new softmax regression problem with

Part 2. If we move the $A_{t}$ to $A_{t+1}$ , then we’re solving a new softmax regression with

In Section 2, we introduce some related work. In Section 3, we give some preliminaries. In Section 4, we compute the gradient of the loss function with softmax function with respect to $x$ . Those functions include $\alpha(x)^{-1}$ , $\alpha(x)$ and $f(x)$ . In Section 5, we show the Lipschitz for the self-attention function with respect to $x$ . In Section 6, we compute the gradient of the loss function with softmax function with respect to $A$ . In Section 7, we show the Lipschitz for the self-attention function with respect to $A$ . In Section 8, we give the main result of our paper.

Related Work

indicated that Transformer-based in-context learners are able to perform traditional learning algorithms implicitly. This is achieved by encoding smaller models within their internal activations. These smaller models are updated by the given context. They theoretically investigate the learning algorithms that Transformer decoders can implement. They demonstrate that Transformers need only a limited number of layers and hidden units to implement various linear regression algorithms. For $d$ -dimensional regression problems, a $O(d)$ -hidden-size Transformer can perform a single step of gradient descent. They also demonstrate that the Transformer with $O(d^{2})$ hidden size is able to update a ridge regression problem. The study reveals that Transformers theoretically have the ability to perform multiple linear regression algorithms.

concentrate on training Transformer to learn certain functions, under in-context conditions. The goal is to have a more comprehensive understanding of in-context learning and determine if Transformers can learn the majority of functions within a given class after training. They found that in-context learning is possible even when there is a distribution shift between the training and inference data or between in-context examples and query inputs. In addition, they find out that Transformers can learn more complex function classes such as sparse linear functions, two-layer neural networks, and decision trees. These trained Transformers have comparable performance to task-specific learning algorithms.

demonstrate and provide an explanation of the similarity between the training process of the Transformers in in-context tasks and some meta-learning formulations based on gradient descent. During the process of training Transformers for auto-regressive tasks, the implementation of in-context learning in the Transformer forward pass is carried out through gradient-based optimization of an implicit auto-regressive inner loss that is constructed from the in-context data.

Formally speaking, they consider the following problem $\min_{x}\|Ax-b\|_{2}$ defined in Definition 1.2. They first show that doing one step of gradient descent carries out data transformation as follows:

where $\widehat{b}$ denotes the updated $b$ and $P$ denotes the projection matrix such that a Transformer step $\widehat{b}_{j}$ on every $j$ is identical to the gradient-induced dynamics $\widetilde{b}_{j}$ . This equivalence implies that when training linear-self-attention-only Transformers for fundamental regression tasks, the models learned by GD and Transformers show great similarity.

explores the occurrence of in-context learning during pre-training when documents exhibit long-range coherence. The Language Model (LLM) develops the ability to generate coherent next tokens by deducing a latent document-level concept. During testing, in-context learning is observed when the LLM deduces a shared latent concept between examples in a prompt. Through the research conducted, it has been demonstrated that in-context learning happens even when there is a distribution mismatch between prompts and pretraining data, especially in scenarios where the pretraining distribution is a mixture of Hidden Markov Models . Theoretically, they show that the error of the in-context predictor is optimal when a distinguishability condition holds. In cases where this condition does not hold, the expected error still reduces as the length of each example increases. This finding highlights the importance of both input and input-output mapping contributes to in-context learning.

2 Transformer Theory

The advancements of Transformers have been noteworthy, however, their learning mechanisms are not completely comprehensible yet. Although these models have performed remarkably well in structured and reasoning activities, our comprehension of their mathematical foundations lags significantly behind. Past research has indicated that the outstanding performance of Transformer-based models can be attributed to the information within their components, such as the multi-head attention. Various studies have presented empirical proof that these components carry a substantial amount of information, which can help in resolving different probing tasks.

Recent research has investigated the potential of Transformers through both theoretical and experimental methods, including Turing completeness , function approximation , formal language representation , and abstract algebraic operation learning . Some of these studies have indicated that Transformers may act as universal approximators for sequence-to-sequence operations and emulate Turing machines . demonstrate the existence of contextual sparsity in LLM, which can be accurately predicted. They exploit the sparsity to speed up LLM inference without degrading the performance from both a theoretical perspective and an empirical perspective. proposed the Pixelated Butterfly model that uses a simple fixed sparsity pattern to speed up the training of Transformer. Other studies have focused on the expressiveness of attention within Transformers .

Furthermore, has demonstrated that moderately sized masked language models may effectively parse and recognize syntactic information that helps in the partial reconstruction of a parse tree. Inspired by the language grammar model studied by , consider the tensor cycle rank approximation problem. consider the exponential regression in neural tangent kernel over-parameterization setting. studied the computation of regularized version of the exponential regression problem but they ignore the normalization factor. consider the softmax regression which considers the normalization factor compared to exponential regression problems . The majority of LLMs can perform attention computations in an approximate manner during the inference process, as long as there are sufficient guarantees of precision. This perspective has been studied by various research, including . With this in mind, have conducted a study on the computation of the attention matrix from the hardness perspective and developed faster algorithms.

Preliminary

In Section 3.1, we introduce the notations used in this paper. In Section 3.2, we give some facts about the basic algebra. In Section 3.3, we propose the lower bound on $\langle\exp(Ax),{\bf 1}_{n}\rangle$ .

For a positive integer $n$ , we use $[n]$ to denote $\{1,2,\cdots,n\}$ , for any positive integer $n$ .

We use ${\bf 1}_{n}$ to denote the vector where all entries are one. We use ${\bf 0}_{0}$ to denote the vector where all entries are zero. The identity matrix of size $n\times n$ is represented by $I_{n}$ for a positive integer $n$ .

The notation $B_{i}$ is used to indicate the $i$ -th row of matrix $B$ .

For a square and full rank matrix $B$ , we use $B^{-1}$ to denote the true inverse of $B$ .

2 Basic Algebras

$\|x\circ y\|_{2}\leq\|x\|_{\infty}\cdot\|y\|_{2}$

$\|x\|_{\infty}\leq\|x\|_{2}\leq\sqrt{n}\|x\|_{\infty}$

$\|\exp(x)\|_{\infty}\leq\exp(\|x\|_{2})$

For any $\|x-y\|_{\infty}\leq 0.01$ , we have $\|\exp(x)-\exp(y)\|_{2}\leq\|\exp(x)\|_{2}\cdot 2\|x-y\|_{\infty}$

If $X\preceq\alpha\cdot Y$ , then $\|X\|\leq\alpha\cdot\|Y\|$

3 Lower bound on β𝛽\beta

Let $\beta$ be lower bound on $\langle\exp(Ax),{\bf 1}_{n}\rangle$

Softmax Function with respect to x𝑥x

In Section 4.1, we give the definitions used in the computation. In Section 4.2, we compute the gradient of the loss function with softmax function with respect to $x$ . Those functions includes $\alpha(x)^{-1}$ , $\alpha(x)$ and $f(x)$ .

We define function softmax $f$ as follows

For convenient, we define two helpful notations $\alpha$ and $c$

Then, we can rewrite $f(x)$ (see Definition 4.1) and $L_{\exp}(x)$ (see Definition 4.2) as follows

$L_{\exp}(x)=0.5\cdot\|\alpha(x)^{-1}\cdot\exp(Ax)-b\|_{2}^{2}$ .

$L_{\exp}(x)=0.5\cdot\|f(x)-b\|_{2}^{2}$ .

Then we can rewrite $L_{\exp}(x)$ (see Definition 4.2) as follows

2 Gradient Computations

Let $\alpha(x)$ be defined in Definition 4.3.

Let $L_{\exp}(x)$ be defined in Definition 4.2.

Lipschitz with respect to x𝑥x

In Section 5.1, we give the preliminary to compute the Lipschitz. In Section 5.2, we show the upper bound of $\delta_{b}$ . In Section 5.3, we compute the Lipschitiz of function $\exp(Ax)$ with respect to $x$ . In Section 5.4, we compute the Lipschitiz of the function $\alpha$ with respect to $x$ . In Section 5.5, we compute the Lipschitiz of function $\alpha^{-1}$ with respect to $x$ .

Let $\{-1,+1\}^{n}$ denote a vector that each entry can be either $-1$ or $+1$ . In the worst case, there are $2^{n}$ possible solutions, e.g.,

The norm of all the choices are the same. Thus, it is sufficient to only consider one solution as follows.

The proof directly follows from Definition 5.1. ∎

For convenience, we split $\delta_{b}$ into two terms, and provide the following definitions

where the 1st step follows from the definitions of $\delta_{b}$ , the 2nd step follows from the definitions of $\delta_{b,1}$ and $\delta_{b,2}$ , the 3rd step follows from simple algebra, the 4th step comes from the definition of $\alpha$ . ∎

Let $\delta_{b}=\delta_{b,1}+\delta_{b,2}$ .

where the first step follows from definition, the second step follows from assumption on $A$ and $x$ , the third step follows Lemma 5.8, the forth step follows from Lemma 5.7, the fifth step follows from Lemma 5.6.

where the first step follows from definition, the 2nd step comes from Lemma 5.6.

where the 1st step follows from the definition of $\delta_{b}$ , the 2nd step follows from triangle inequality, the 3rd step follows from the results in Part 1 and Part 2, the 4th step follows from the fact that $n\geq 1$ and $\beta^{-1}\geq 1$ , the 5th step follows from simple algebra. ∎

3 Lipschitz for function exp⁡(Ax)𝐴𝑥\exp(Ax) with respect to x𝑥x

Let $x,y$ satisfy that $\|x\|_{2}\leq R$ and $\|y\|_{2}\leq R$

where the 1st step follows from $\|A(y-x)\|_{\infty}<0.01$ and Fact 3.1, the 2nd step comes from Fact 3.1, the 3rd step follows from Fact 3.2, the 4th step follows from Fact 3.2, the last step follows from $\|A\|\leq R$ . ∎

4 Lipschitz for function α(x)𝛼𝑥\alpha(x) with respect to x𝑥x

Let $\alpha(x)$ be defined as Definition 4.3

Let $\langle\exp(Ax),{\bf 1}_{n}\rangle\geq\beta$

Let $\langle\exp(Ay),{\bf 1}_{n}\rangle\geq\beta$

Softmax Function with respect to A𝐴A

In this section, we consider the function with respect to $A$ . We define function softmax $f$ as follows

Similarly, we reparameterized $x$ by $A$ for our loss function $L$ . We define loss function $L$ as follows

For convenience, we define two helpful notations $\alpha$ and $c$ with respect to $A$ as follows:

Then, we can rewrite $f(A)$ (see Definition 6.1) and $L_{\exp}(A)$ (see Definition 6.2) as follows

$L_{\exp}(A)=0.5\cdot\|\alpha(A)^{-1}\cdot\exp(Ax)-b\|_{2}^{2}$ .

$L_{\exp}(A)=0.5\cdot\|f(A)-b\|_{2}^{2}$ .

Then we can rewrite $L_{\exp}(A)$ (see Definition 6.2) as follows

Lipschitz with respect to A𝐴A

In Section 7.1, we give the preliminary to compute the Lipschitz. In Section 7.2, we show the upper bound of $\delta_{b}$ with respect to $A$ . In Section 7.3, we compute the Lipschitiz of function $\exp(Ax)$ with respect to $A$ . In Section 7.4, we compute the Lipschitiz of the function $\alpha$ with respect to $A$ . In Section 7.5, we compute the Lipschitiz of function $\alpha^{-1}$ with respect to $A$ .

The proof directly follows from Definition 7.1. ∎

For convenient, we split $\delta_{b}$ into two terms, and provide the following definitions

where the 1st step follows from the definitions of $\delta_{b}$ , the 2nd step follows from the definitions of $\delta_{b,1}$ and $\delta_{b,2}$ , the 3rd step comes from simple algebra, the 4th step comes from the definition of $\alpha$ . ∎

Let $\delta_{b}=\delta_{b,1}+\delta_{b,2}$ .

where the first step follows from definition, the second step follows from assumption on $A$ and $x$ , the third step follows Lemma 7.8, the forth step follows from Lemma 7.7, the fifth step follows from Lemma 7.6.

where the 1st step follows from the definition of $\delta_{b}$ , the 2nd step comes from triangle inequality, the 3rd step comes from the results in Part 1 and Part 2, the 4th step follows from the fact that $n\geq 1$ and $\beta^{-1}\geq 1$ , the 5th step follows from simple algebra. ∎

3 Lipschitz for function exp⁡(Ax)𝐴𝑥\exp(Ax) with respect to A𝐴A

where the 1st step follows from $\|A(y-x)\|_{\infty}<0.01$ and Fact 3.1, the 2nd step follows from Fact 3.1, the 3rd step follows from Fact 3.2, the 4th step comes from Fact 3.2, the last step follows from $\|A\|\leq R$ . ∎

4 Lipschitz for function α(A)𝛼𝐴\alpha(A) with respect to A𝐴A

Let $\alpha(A)$ be defined as Definition 6.3

where the 1st step comes from the definition of $\alpha(x)$ , the 2nd step follows from Cauchy-Schwarz inequality (Fact 3.1). ∎

Let $\langle\exp(Ax),{\bf 1}_{n}\rangle\geq\beta$

Let $\langle\exp(Bx),{\bf 1}_{n}\rangle\geq\beta$

where the 1st step follows from simple algebra, the 2nd step follows from $\alpha(A)\geq\beta,\alpha(B)\geq\beta$ . ∎

Main Results

In Section 8.1, we show our upper bound result of $\delta_{b}$ with respect to $x$ . In Section 8.2, we show our upper bound result of $\delta_{b}$ with respect to $A$ .

We consider the softmax regression problem

If we move the $x_{t}$ to $x_{t+1}$ , then we’re solving a new softmax regression problem with

where the 1st step follows from Lemma 5.5, the 2nd step comes from Lemma 3.3, the 3rd step comes from simple algebra, the 4th step follows from simple algebra, the 5th step follows from simple algebra and the 6th step follows from the definition of $M$ . ∎

2 Shifting Sentence Data A𝐴A