Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Brian K Chen, Tianyang Hu, Hui Jin, Hwee Kuan Lee, Kenji Kawaguchi

TL;DR

The paper tackles permanent integration of in-context learning into transformer weights by exploiting linearized attention and introducing a bias-based conversion (ICLCA). It proves that exact conversion is possible for linearized attention, and extends the idea to an approximate method for regular softmax attention, validating with synthetic tasks and a GPT-2 experiment. Key contributions include linking ICL to the Key-Value matrix, outlining an efficient algorithm, and demonstrating practical benefits such as faster training and improved context retention. The work has practical implications for efficient, interpretable context integration in large language models and suggests directions for scalable approximate methods in broader architectures.

Abstract

In-Context Learning (ICL) has been a powerful emergent property of large language models 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.

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

TL;DR

The paper tackles permanent integration of in-context learning into transformer weights by exploiting linearized attention and introducing a bias-based conversion (ICLCA). It proves that exact conversion is possible for linearized attention, and extends the idea to an approximate method for regular softmax attention, validating with synthetic tasks and a GPT-2 experiment. Key contributions include linking ICL to the Key-Value matrix, outlining an efficient algorithm, and demonstrating practical benefits such as faster training and improved context retention. The work has practical implications for efficient, interpretable context integration in large language models and suggests directions for scalable approximate methods in broader architectures.

Abstract

In-Context Learning (ICL) has been a powerful emergent property of large language models 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.
Paper Structure (28 sections, 2 theorems, 25 equations, 1 figure, 5 tables, 2 algorithms)

This paper contains 28 sections, 2 theorems, 25 equations, 1 figure, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Linear and Kernel Attention
  4. Preliminaries
  5. The Attention Mechanism
  6. In Context Learning
  7. Exact Conversion of In-Context Learning to Model Weights
  8. Direct Conversion for Linear Attention
  9. Addition of Free Parameters
  10. Assumptions Necessary for Conversion
  11. ICL Conversion for Linearized Attention
  12. ICL Conversion Algorithm for Linearized Attention
  13. Approximate Conversion in Regular Attention
  14. Experiments
  15. Effectiveness of Conversion
  16. ...and 13 more sections

Key Result

Theorem 4.1

For weights $W_{Q},W_{K},W_{V}$, input prompt $X$, ICL prompt $X'$ and initial bias $b_{VK}$, we have that: holds for all $X$ if we set

Figures (1)

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

Theorems & Definitions (8)

  • Definition 3.1: Single-headed Attention Layer
  • Definition 3.2: linearized attention Layer
  • Definition 3.3: Masked linearized attention Layer
  • Definition 3.4: Key-Value Matrix
  • Theorem 4.1: Bias conversion for Linear Attention
  • Theorem 4.5
  • proof
  • Definition 2.1: Rotary matrix