Exact Conversion of In-Context Learning to Model Weights in Linearized-Attention Transformers

Abstract: In-Context Learning (ICL) has been a powerful emergent property of LLMs that has attracted increasing attention in recent years. In contrast to regular gradient-based learning, ICL is highly interpretable and does not require parameter updates. In this paper, we show that, for linearized transformer networks, ICL can be made explicit and permanent through the inclusion of bias terms. We mathematically demonstrate the equivalence between a model with ICL demonstration prompts and the same model with the additional bias terms. Our algorithm (ICLCA) allows for exact conversion in an inexpensive manner. Existing methods are not exact and require expensive parameter updates. We demonstrate the efficacy of our approach through experiments that show the exact incorporation of ICL tokens into a linear transformer. We further suggest how our method can be adapted to achieve cheap approximate conversion of ICL tokens, even in regular transformer networks that are not linearized. Our experiments on GPT-2 show that, even though the conversion is only approximate, the model still gains valuable context from the included bias terms.

• The paper introduces ICLCA, a method that precisely converts in-context learning into model weights by adjusting bias terms in linearized-attention transformers.
• It leverages the theory that Key-Value matrices embed ICL information, demonstrating negligible output logits errors on synthetic tasks.
• Experimental results on GPT-2 show that both exact and approximate conversions enhance model interpretability and computational efficiency.

ICL to Model Weights via Exact Conversion

The paper "Exact Conversion of In-Context Learning to Model Weights in Linearized-Attention Transformers" (2406.02847) introduces a method, ICLCA, to explicitly and permanently incorporate ICL into linearized transformer networks by adding bias terms. This approach achieves exact conversion without requiring parameter updates, offering an interpretable and computationally efficient alternative to existing context distillation techniques. The paper further proposes an approximate conversion method for standard transformers, demonstrating its efficacy on GPT-2.

Theoretical Foundation

The paper begins by establishing the theoretical groundwork for ICL conversion. It highlights that ICL relationships are captured within the Key and Value matrices of the attention module. The paper demonstrates the equivalence between a model with ICL demonstration prompts and the same model augmented with additional bias terms. For a linearized attention layer, the output is given by:

$$O^{T}_{i} = \frac{1}{D_{1}(Q_{i})^{T}D_{2}(X)}\phi(Q_{i})^{T}\sum_{j=1}^{i}\phi(K_{j})V^{T}_{j}$$

where $\phi(\cdot)$ is a feature map, $Q$, $K$, and $V$ are the query, key, and value matrices, and $D_1$ and $D_2$ are normalization terms. The key insight is to add bias terms to the Key-Value matrix and the normalizing term. The Key-Value matrix $A$ is defined as:

$A = \sum_{j=1}^{N} \phi(K_{j})V^{T}_{j}$

Figure 1: The training loss curves and in context accuracy curve during training.

Exact Conversion with Linearized Attention

The paper's core contribution lies in its exact conversion method for linearized attention models. This method involves adding bias terms to the Key-Value matrix and the normalizing term within the attention module. Specifically, the updated biases are calculated as:

$b'_{KV} = R^{d_{K}_{\Theta, -M}b_{KV} + \sum_{j=1}^{M}R^{d_{K}_{\Theta,j-M}\phi(K'_{j})V'^{T}_{j}$

$b'_{D} = b_{D} + D^{*}_{2}(X')$

where $R$ is a rotary matrix, $M$ is the length of the ICL prompt, and $X'$ represents the ICL prompt. This approach ensures that the model's behavior with ICL prompts is precisely replicated by the modified model without ICL prompts. Algorithm 1 in the paper details the ICL conversion algorithm (ICLCA) for models with L sub-layers, outlining the steps to update the bias terms in each linearized attention layer.

Approximate Conversion for Standard Transformers

Recognizing the prevalence of standard transformers with softmax attention, the paper extends its methodology to an approximate conversion method. This involves approximating the softmax similarity function with a kernel function and then applying the bias conversion technique. The resulting output takes the form:

$O ^{T}_{i} = \frac{[\sum_{j=1}^{N}\mathrm{sim}(Q_{i},K_{j})V^{T}_{j}] + R_{\Theta,i}^{d_{K}\phi(Q_{i})^{T}b_{KV}}{[\sum_{j=1}^{N}\mathrm{sim}(Q_{i},K_{j})] + R_{\Theta,i}^{d_{K}\phi(Q_{i})^{T}b_{D}}$

where

$b_{KV} = \sum_{j=1}^{M}R^{d_{K}_{\Theta,j-M}\phi(K'_{j})V'^{T}_{j}, \quad b_{D} = \sum_{j=1}^{M}\phi(K'_{j}).$

While this conversion is not exact, experiments on GPT-2 demonstrate that the model retains a degree of in-context information.

Experimental Results

The paper presents empirical evidence to support its claims. The authors demonstrate the numerical exactness of ICLCA for linearized attention models, achieving negligible relative errors in output logits before and after conversion. On a synthetic induction head task, ICLCA accurately captures the ICL prompt in the bias terms, resulting in high in-context learning accuracy. Furthermore, experiments on GPT-2 show that the approximate conversion method reduces the relative error of logits and enables the model to generate context-aware text.

Implications and Future Directions

The paper's findings have significant implications for the field of LLMs. The exact conversion method offers a computationally efficient and interpretable way to incorporate new knowledge into models, potentially reducing the need for extensive fine-tuning. The authors suggest that future research could explore more sophisticated kernels for approximating softmax attention, investigate methods to achieve exact conversion for standard transformers, and apply the approximate conversion method in conjunction with existing fine-tuning techniques. The observation that contextual information is captured within the Key-Value matrix opens up new avenues for research into the interactions between these matrices and alternative approaches for manipulating them.

Conclusion

The paper "Exact Conversion of In-Context Learning to Model Weights in Linearized-Attention Transformers" (2406.02847) presents a novel approach to incorporating ICL into LLMs. By leveraging the linearity of attention modules and introducing bias terms, the authors achieve exact conversion for linearized attention models and provide an effective approximation for standard transformers. This work contributes to a deeper understanding of ICL and offers practical tools for enhancing the capabilities of LLMs.