PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2406.02847v2cs.LGstat.ML

TL;DR

ICL prompts are interpretable but temporary, while existing approaches to permanently distilling them into weights are limited. The paper adds bias terms to convert ICL prompts exactly for linearized attention and approximately for regular transformers; experiments report exact incorporation in a linear-attention model and retained context in GPT-2.

  • Problem

    ICL demonstrations provide interpretable context without parameter updates, but their information is temporary, and existing methods for permanently encoding it are limited and largely lack theoretical guarantees.

    Source: §1 · p. 1

  • Method

    The method incorporates ICL prompts into attention-module weights through a bias term, with exact conversion for linearized attention and an approximate extension to regular transformers.

    Source: §1 · pp. 1–2

  • Results

    Experiments demonstrate exact incorporation of ICL tokens in a linear-attention model, while GPT-2 with regular attention retains some ICL information after approximate conversion.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The conversion offers a computationally inexpensive way to incorporate interpretable ICL prompts into model weights, with the exact guarantee applying to linearized attention.

    Source: §8 · p. 9

    For approximate softmax conversion, scalability depends on whether the kernel function can approximate softmax in pretrained transformers and scale to larger models.

    Source: §8 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses the gap between interpretable but temporary ICL prompts and costly, empirically distilled weight updates. It introduces exact bias-based conversion for linearized attention and approximate conversion for regular transformers.

1. Introduction

  • ICL demonstrations are interpretable but temporary: they apply only to the prompt in which they appear, while permanent conversion studies remain limited.Prior context-distillation methods are described as largely empirical and without theoretical guarantees.Source: §1 · p. 1
  • The conversion requires a linear bias update, adds about 1% to the model's parameters, and can be combined or removed to incorporate or discard prompts.The authors describe the update as computationally inexpensive compared with gradient descent.Source: §1 · p. 2
  • For regular softmax-attention transformers, the paper proposes approximate conversion and reports that GPT-2 retains some ICL information from the added biases.The introduction also reports exact incorporation of ICL tokens into a linear-attention model.Source: §1 · p. 2
  • The method converts ICL prompts into model weights by adding a bias term to the attention module, exactly for linearized attention models.The paper identifies the Key-Value matrices as where attention captures ICL relationships.Source: §1 · pp. 1–2
  • Unlike existing context-distillation methods, the proposed conversion is theoretically supported and avoids expensive parameter updates.The paper describes existing distillation approaches as largely empirical and lacking theoretical guarantees.Source: §1 · pp. 1–2

2. Related Works

The related-work discussion situates the proposal among transformer architectures, efficient attention variants, and approaches to understanding or permanently incorporating ICL.

Transformers and Large Language Models

  • Transformers use attention to capture long-range dependencies, while their adoption has extended across natural language processing and other fields.The section contrasts this capability with the sequential operations of LSTMs and RNNs.Source: §2 · p. 2

Linear and Kernel Attention

  • Linear and kernel attention aim to address softmax attention's quadratic computation in input length, and RetNet claims state-of-the-art matching with linear attention and added mechanisms.The section notes that linearized transformers remain a minority despite the computational pressures of lengthy prompts.Source: §2 · p. 2

In Context Learning and Context Distillation

  • Prior ICL distillation methods use fine-tuning to match prompt input-output relationships, but the cited work is described as empirical, costly, and lacking theoretical guarantees.The section also reviews studies explaining ICL through Bayesian frameworks and transformer learning algorithms.Source: §2 · pp. 2–3
  • The paper positions its approach as using linear and kernel attention structure for exact ICL-to-weight conversion, avoiding the approximation and expense it attributes to gradient-descent alternatives.The authors characterize current permanent-encoding approaches as gradient-descent fine-tuning.Source: §2 · p. 3

3. Preliminaries

The preliminaries set up token-sequence inputs, describe standard and linearized attention, and define how demonstration tokens are represented for ICL.

3. Preliminaries

  • The paper represents transformer input as a sequence of N tokens, each with dimensionality d_in, passed through an encoder before transformer layers.The input matrix is denoted X and has shape N × d_in.Source: §3 · p. 3

3.1. The Attention Mechanism

  • Attention uses query, key, and value matrices to weigh token influences; the paper focuses on single-headed attention but says its methods translate by treating each head separately.The general attention formulation is extended to linearized attention by approximating similarity with a kernel and associated feature map.Source: §3.1 · p. 3
  • The paper defines a Key-Value matrix in linearized attention and identifies it as the part that captures ICL information.For masked linearized attention, the matrix's sum runs only through the current token rather than the full sequence.Source: §3.1 · p. 4

3.2. In Context Learning

  • For ICL, the paper abstracts away demonstration structure and treats demonstration and input prompts as token sets, concatenating the demonstration tokens before the input tokens.It denotes their lengths M and N, respectively.Source: §3.2 · p. 4

4. Exact Conversion

The section develops exact ICL-prompt conversion for linearized attention by adding bias terms to model weights. Direct weight changes alone fail generally; the conversion is formalized under stated assumptions and implemented by ICLCA.

4.1. Direct Conversion for Linear Attention

  • Directly changing existing linear-attention weights cannot generally convert ICL prompts, except under circumstances the authors consider unlikely in real-world scenarios.The attempted conversion is framed as an affine equation, but the required condition does not hold in general.Source: §4.1 · pp. 4–5

4.2. Addition of Free Parameters

  • Adding a Key-Value bias bKV supplies the free parameter used for conversion, which the theorem shows works for all X without considering causal masks or positional encodings.The paper presents this simpler case before extending the method to architectures with masks and positional encodings.Source: §4.2 · p. 5

4.3. Assumptions Necessary for Conversion

  • Conversion assumes autoregressive attention and a normalizer separation property; relative positional encoding simplifies the process but is not strictly necessary.With absolute positional embeddings, a following linear transformation is needed to correct the positional shift.Source: §4.3 · p. 5

4.4. ICL Conversion for Linearized Attention

  • Theorem 4.5 gives bias terms that make a linearized-attention model with ICL prompts equivalent for all input prompts X.For RoPE, the method adds separate biases to the Key-Value matrix and the normalizing term.Source: §4.4 · p. 6
  • The RoPE update shifts the existing Key-Value bias left by M tokens to account for added ICL prompts, while the Key-Value matrix itself appears to capture contextual information.The authors identify interactions between the Key and Value matrices as a direction for future study.Source: §4.4 · p. 6

4.5. ICL Conversion Algorithm for Linearized Attention

  • ICLCA processes layers in sequence, passing prompts through non-attention layers and updating biases at linearized-attention layers before saving the new biases.The method considers models whose other sub-layers are element-wise operators.Source: §4.5 · p. 6

Conversion in Regular Attention

For regular softmax-attention transformers, the paper approximates the attention similarity with a kernel, applies the linearized conversion, and returns the bias to the original model. Conversion is approximate, and its accuracy depends on the kernel approximation.

Conversion in Regular Attention

  • Regular softmax-attention models cannot receive exact ICL conversion with this method, but can receive a direct approximation.The authors note that many models use vanilla attention without linearization.Source: Conversion in Regular Attention · p. 7
  • The approximate pipeline replaces the similarity function with a kernel, performs the linearized conversion, and reintroduces the bias into the original model.The corresponding ICLAA algorithm is included in Appendix D.1.Source: Conversion in Regular Attention · p. 7
  • Approximation accuracy depends directly on how accurately the kernel represents the similarity function.The paper notes that studies continue to explore kernel approximation.Source: Conversion in Regular Attention · p. 7

6. Experiments

Experiments test exact conversion in linearized attention, prompt capture on an induction task, the effect of adding bias terms, and approximate conversion on GPT-2. The linearized settings support exact prompt capture, while GPT-2’s relative logit error decreases after approximate conversion.

6.1. Effectiveness of Conversion

  • Linear-attention conversion matches the prompted model’s logits up to numerical rounding errors across model sizes.Table 1 reports relative logit error averaged over 100 randomly generated ICL prompts and input prompts.Source: §6.1 · p. 7, Table 1 · p. 7

6.2. Induction Head Task

  • On the induction-head task, ICLCA exactly captures the ICL prompt, as assessed by in-context accuracy under three settings.The evaluation uses a 12-layer linear-attention transformer with embedding dimension 128 and RoPE.Source: §6.2 · p. 7, Table 2 · p. 7

6.3. Effect of the Added Bias Term

  • Adding bias terms makes training slightly faster and improves the model’s induction-head task performance when trained from scratch.Figure 1 plots training loss and in-context accuracy for models with and without the modification.Source: §6.3 · pp. 7–8, Figure 1 · p. 8

6.4. Experiment on Pretrained GPT-2

  • 9.17% relative logit error after conversion, down from 16.56% before conversion on pretrained GPT-2.The converted model also generates text that reflects information in the ICL prompt.Source: §6.4 · p. 8

7. Transformers

The paper recasts linearized transformers as RNNs, where adding bias terms corresponds to changing the RNN’s initial states. This view suggests a possible extension of ICL conversion to autoregressive RNN architectures.

as RNNs and Associated Extensions

  • Table 3 is identified as a comparison of generated texts.The supplied passage identifies the table’s subject but gives no comparison outcome.Source: Table 3 · p. 8
  • In the RNN formulation, adding bKV and bD is equivalent to changing the initial states s0 and z0.The paper notes that future tokens observe only previous hidden states, making capture of the ICL prompt’s hidden states sufficient.Source: as RNNs and Associated Extensions · p. 8
  • If autoregressive RNN architectures such as H3 or MAMBA exhibit ICL, the authors suggest capturing ICL prompts in their initialization.This extension is proposed as a possibility, not demonstrated experimentally.Source: as RNNs and Associated Extensions · p. 8

8. Conclusion and Discussion

The paper presents exact, inexpensive conversion for linearized attention by adding a bias term to the Key-Value matrix, and discusses possible practical uses. Approximate conversion for standard transformers remains constrained by kernel scalability and may be useful as a fine-tuning initialization.

8. Conclusion and Discussion

  • Adding a bias term to the Key-Value matrix permanently incorporates ICL prompts into linearized-attention models through an exact, interpretable, inexpensive conversion.The paper presents this as a theoretical and practical contribution.Source: §8 · p. 9
  • Converted context could support localized model modifications and avoid repeatedly appending long prompts such as system instructions or personal history.The authors describe potential computation savings for inference from converting prompts that would otherwise be repeatedly included.Source: §8 · p. 9
  • Approximate conversion for standard transformers may serve as an initialization for fine-tuning, but may not suffice for real-world applications on its own.The authors suggest it could reduce the number of fine-tuning steps required.Source: §8 · p. 9
  • Approximate-conversion scalability depends primarily on whether the kernel approximates softmax well in pretrained transformers; exact linear-attention conversion scales without the same issue.The authors note that their approximation applies to ICL tokens rather than all tokens.Source: §8 · p. 9
  • The authors identify testing more sophisticated kernels, combining approximate conversion with fine-tuning, and pursuing exact softmax-attention conversion as future directions.Source: §8 · p. 9

Impact Statement

The impact statement describes potential access and cost benefits alongside misuse concerns. In an appendix discussion, the authors reject a direct-conversion setup because it lacks guaranteed uniqueness and requires an input-dependent shift; a proof section introduces RoPE’s rotary matrix.

Impact Statement

  • The authors frame permanent, interpretable context incorporation without expensive parameter updates as potentially improving access and cost, while noting risks of reinforcing bias or misinformation.They call for ethical guidelines and safeguards.Source: Impact Statement · p. 9

A. Naive Attempt at Direct Conversion

  • A direct-conversion setup lacks guaranteed unique W1 and W2 and requires a shift c that depends on X, so the authors choose free parameters instead.The required equality must hold for all X.Source: A. Naive Attempt at Direct Conversion · p. 12

B. Proof of Theorem 4.5

C. Analysis of Other Choices of Free Parameters

The section weighs alternative locations and forms for free parameters, ultimately focusing on attention-module bias terms. It also identifies a possible future FFN-in-key-value-matrix approach whose exact conversion remains unverified.

C. Analysis of Other Choices of Free Parameters

  • An input encoding layer could capture ICL-token information before the transformer, but the authors considered this and compressor-based alternatives fanciful.The compressor option would initialize from compressed ICL inputs.Source: C. Analysis of Other Choices of Free… · p. 13
  • Modifying the FFN could target factual associations located in GPT networks' MLP/FFN, but exact conversion could remain challenging.This alternative would avoid adding bias terms to the numerator and denominator in theorem 4.5.Source: C. Analysis of Other Choices of Free… · p. 13
  • Attention-module bias terms were chosen over scaling, which could be matched by scaling input tokens, and initialization tokens, which would cost more computation for the same result.The authors also considered modifying inputs or the FFN as alternatives.Source: C. Analysis of Other Choices of Free… · p. 13
  • Including the full FFN within the key-value matrix is proposed as future work, and its claims require detailed experiments for verification.The proposal departs from the current transformer architecture and draws on the key-value matrix's special significance in the authors' findings.Source: C. Analysis of Other Choices of Free… · p. 13

D. Extension to Other Architectures

The section proposes an FFN-based route for ICL conversion in GPT-J, while noting architectural and normalization constraints. It also presents an approximation algorithm for regular attention and additional before-and-after GPT-2 text examples.

D. Extension to Other Architectures

  • GPT-J's tandem attention-and-FFN structure offers a conversion route that adds a term to each token's FFN output instead of adding a new attention bias.The authors note GPT-J matches the typical stacked architecture in performance but does not exceed it.Source: D. Extension to Other Architectures · p. 13
  • The GPT-J route requires a linearized-attention study and a linear setting with normalizing term D(·)=1, a potentially undesirable restriction.The authors note that softmax attention has not been studied in this linearized setting and that RetNet's promise does not ensure D(·)=1 will suit future transformers.Source: D. Extension to Other Architectures · p. 13

D.1. ICL Approximation Algorithm for Regular Attention

  • ICLAA approximately incorporates prompts into regular-attention models by processing non-attention sublayers and saving new bias terms.The algorithm is introduced for models with regular attention and other element-wise sublayers.Source: §D.1 · pp. 13–14

E. More Examples of GPT-2 Generated Text

Loading 2406.02847v2…