PaperlayerView exact revision on arXiv

Source-linked AI summary

In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering

Sheng Liu, Haotian Ye, Lei Xing, James Zou

arXiv:2311.06668v3cs.LGcs.AIcs.CL

TL;DR

Standard in-context learning can be uneven, context-intensive, and difficult to control. The paper proposes In-Context Vectors, which summarize demonstrations into latent vectors and apply them to query representations; across diverse tasks, ICV outperforms standard ICL and LoRA fine-tuning while supporting task combination through vector arithmetic.

  • Problem

    In-context learning can have limited and template-sensitive effectiveness, consume context space, and become inefficient with many demonstrations.

  • Method

    ICV computes a vector from demonstration-example latent states, then adds it to latent states across model layers and token positions while processing only the query.

  • Results

    ICV outperforms standard ICL and LoRA fine-tuning across diverse tasks including safety, style transfer, role-playing, and formatting.

  • Takeaways & Limitations

    Vector magnitude provides direct control, while vector addition and subtraction enable combining multiple task instructions.

Abstract

from arXiv · show

Large language models (LLMs) demonstrate emergent in-context learning capabilities, where they adapt to new tasks based on example demonstrations. However, in-context learning has seen limited effectiveness in many settings, is difficult to quantitatively control and takes up context window space. To overcome these limitations, we propose an alternative approach that recasts in-context learning as in-context vectors (ICV). Using ICV has two steps. We first use a forward pass on demonstration examples to create the in-context vector from the latent embedding of the LLM. This vector captures essential information about the intended task. On a new query, instead of adding demonstrations to the prompt, we shift the latent states of the LLM using the ICV. The ICV approach has several benefits: 1) it enables the LLM to more effectively follow the demonstration examples; 2) it's easy to control by adjusting the magnitude of the ICV; 3) it reduces the length of the prompt by removing the in-context demonstrations; 4) ICV is computationally much more efficient than fine-tuning. We demonstrate that ICV achieves better performance compared to standard in-context learning and fine-tuning on diverse tasks including safety, style transfer, role-playing and formatting. Moreover, we show that we can flexibly teach LLM to simultaneously follow different types of instructions by simple vector arithmetics on the corresponding ICVs.

1 Introduction

The paper introduces In-Context Vectors (ICV) as a controllable, efficient alternative to standard in-context learning, addressing uneven effectiveness, context-length limits, and computational cost. ICV summarizes demonstrations into a latent vector and uses it to steer queries, outperforming standard ICL and LoRA fine-tuning across diverse tasks.

  • Motivation: ICL effectiveness is sensitive to templates, verbalizers, and demonstrations, while long contexts create computational inefficiency.Popular open-source LLMs may have a maximum context length of 4096, making large numbers of demonstrations expensive.
  • ICV approach: ICV computes a task vector from demonstration-example latent states and applies it to shift the LLM’s latent states during query generation.This decomposes adaptation into task summary and feature shifting, without prepending demonstrations to the query.
  • Efficiency and control: ICV supports demonstrations beyond the context-length limit by representing them in a single vector rather than directly adding them to the prompt.The method incurs negligible overhead for computing the vector and introduces no new parameters.
  • Efficiency and control: ICV directly controls latent-state shifts through a specified magnitude, unlike ICL’s indirect self-attention-based shift.It can use demonstration examples without requiring a template.
  • Results: ICV significantly outperforms standard ICL and LoRA fine-tuning across detoxification, style transformation, role-playing, and formatting tasks.The evaluations use models including Falcon and Llama.

2 Backgrounds

Standard in-context learning prepends demonstrations to a query, allowing the self-attention mechanism to modify the query’s latent representation. The resulting modification can be understood as a position-wise latent-feature shift whose direction and distance are automatically determined by attention.

  • In-context learning: Standard ICL concatenates demonstration pairs with a query and asks the model to predict the query’s output from those examples.The demonstrations can specify transformations such as changing sentiment or rewriting text in a particular style.
  • Attention mechanism: Prepending demonstrations changes the self-attention computation by adding their keys and values to the query’s attention context.The query attends over the concatenated demonstrations and query sequence.
  • Latent feature shifting: ICL applies a position-wise modification to the original attention output, with the shift direction and distance controlled by self-attention.The attention-weight scalar α represents the summed normalized attention between demonstrations and query examples.

3 Method

ICV extracts task information from demonstration examples and directly shifts a model’s latent states to guide new queries. Its construction supports paired or unpaired demonstrations, controllable task strength, and vector arithmetic across tasks without fine-tuning.

  • Efficiency and control: ICV summarizes demonstrations in a single vector, allowing examples beyond the context limit and avoiding templates, while adding negligible computation overhead.The method uses inference-only latent-state shifting and does not fine-tune the model or train additional components.
  • ICV construction: ICV obtains latent representations from demonstration inputs and targets, then combines them into a vector encoding task information.Inputs and targets are separately passed through the LLM; final-token latent states across layers form the representations used to construct the ICV.
  • Paired demonstrations: For paired demonstrations, the ICV is the first principal direction of the differences h(y_i) − h(x_i).This follows from the constrained optimization objective that encourages similarity to target representations and dissimilarity from input representations.
  • Unpaired demonstrations: ICV construction extends beyond paired examples using contrastive loss, with positive y examples and negative x examples softly associated through learned weights.Because a closed-form solution is unavailable in this setting, the ICV is obtained from the gradient of the contrastive objective.
  • Feature shifting: At inference, the ICV is added to every token position at every transformer layer, with λ controlling the strength of the task shift.The updated latent states are normalized to match the pre-update ℓ2 norm, preserving their typical magnitude for subsequent modules.
  • Task arithmetic: ICV supports task reversal and composition by negating or adding vectors, such as combining “safe” and “polite” vectors to produce safer but rude text.The method applies vector arithmetic to transfer aligned or opposite task directions and to combine related task behaviors.

4 Experiments

The experiments evaluate ICV across safety, style transfer, role-playing, and related generation settings, comparing it with conventional ICL and LoRA fine-tuning. They use paired and unpaired demonstrations, varied demonstration counts, multiple models, and automatic safety, similarity, style, and role-playing evaluations.

  • Experimental scope: ICV is evaluated on language detoxification, dialogue safety, jailbreak safety, style transfer, role-playing, and formatting tasks.The experiments include models from the LLaMA, Falcon, and Vicuna families and compare ICV with conventional ICL and, where paired demonstrations are available, LoRA fine-tuning.
  • Safety tasks: For language detoxification, demonstrations pair offensive sentences with inoffensive paraphrases, and ICV uses the paired examples directly without prompts or instructions.ParaDetox evaluation uses five randomly selected demonstrations and 670 other queries.
  • Safety tasks: Jailbreak experiments use five malicious-query demonstrations and measure attack success rate over 100 harmful behaviors.Conventional ICL uses the jail-broken column as demonstrations, while ICV uses the listed demonstrations.
  • Evaluation: Automatic evaluation uses a safety classifier, style-classifier accuracy, ROUGE-1, BERT scores, and GPT-3.5-Turbo comparative assessments for role-playing.ROUGE-1 measures raw-text similarity and BERT scores measure feature-domain similarity.
  • Style and role-playing tasks: Style-transfer experiments use five demonstrations for sentiment and formality, ten for Shakespeare role-playing, and three to four for additional style and formatting tasks.The additional tasks include reserved-to-emotive style, rudeness-to-politeness, and capitalization formatting.

5 Results

Across safety, style, role-playing, and task-combination experiments, ICV generally outperforms conventional ICL and LoRA fine-tuning while exposing controllable tradeoffs through its scaling factor. The results also show that demonstration count and vector arithmetic affect behavior, whereas applying vectors across all layers is most effective.

  • Language detoxification and dialogue safety: ICV reduces toxicity by 49.81% on Falcon-7b and 45.31% on Llama-7b compared with the evaluated baselines.It also maintains high semantic similarity to reference sentences, with robust ROUGE-1 and BERT scores.
  • Language detoxification and dialogue safety: Increasing λ strengthens detoxification but lowers ROUGE-1, reflecting a tradeoff between safety and preservation of the original meaning.The paper also reports reduced fluency at larger λ values.
  • Language detoxification and dialogue safety: More demonstrations correlate with fewer toxic generations, and ICV can use additional examples without the context-length constraint of conventional ICL.Figure 3 reports the mean across three random seeds on Falcon-7b.
  • Language detoxification and dialogue safety: Applying ICV across all Transformer layers is more effective than applying it only to the first, middle, or last layer.The layer-specific variants perform similarly to using only the query example.
  • Jailbreak: With stronger ICV vectors, jailbreak attack success rate rises to 99%, matching optimization-based methods that take around 30 minutes per instance while ICV takes seconds.The jailbreak behavior appears after five malicious-query instances.
  • Speaking style: ICV improves formality by 36.8% and positivity by 65.2% and also transfers formatting and emotional expressiveness.These results are reported for style-transfer experiments using Llama-7b.
  • Role-playing: For Shakespeare role-playing, ICV has higher GPT-3.5-Turbo win rates than ICL and LoRA fine-tuning, with performance increasing for larger models.The win rate is the frequency with which a method ranks first among the three methods.
  • Task arithmetic: Adding or subtracting corresponding ICVs combines multiple task directions, such as making text safer while subtracting a politeness vector to make it rude.The arithmetic demonstrations use Falcon-7B without additional training.

6 Related works

Related work situates ICV among methods for improving and understanding ICL, activation editing, task arithmetic, and LLM safety. The paper distinguishes ICV by extracting task information from demonstrations and directly shifting latent states for open-ended generation without templates or prompts.

  • Improving and understanding ICL: Prior ICL work improves performance through template selection, example selection, consistency, diversity, frequency, flipped learning, noisy-channel prompting, and iterative context updates.Other work trains decoder networks as alternatives to few-shot ICL or uses nearest-neighbor labeling for multiple-choice tasks.
  • Activation editing: A concurrent method obtains a task vector from one latent layer and replaces that layer’s query states, whereas ICV targets open-ended generation.The concurrent method selects its layer using development-set prediction accuracy.
  • Activation editing: Activation-editing research steers sentiment, topics, concepts, agent behavior, counterfactual behavior, and TruthfulQA performance through latent representations.The paper positions ICV as an in-context-learning method that uses demonstrations rather than specific templates or prompts.
  • Task arithmetic: Prior task-arithmetic work derives task vectors from fine-tuning weight differences, while this paper reports analogous addition and negation properties for ICV without fine-tuning.Adding vectors can improve task performance, whereas negating them can induce task forgetting in the prior formulation.
  • LLM safety: LLM safety research includes further training, negative-prompt activation updates, and standard ICL for detoxification and dialogue safety.The paper addresses the same safety setting with ICV as an alternative that requires no additional training.

7 Conclusions

The paper presents ICV as a two-stage framework that extracts task information from demonstrations and applies it to query generation through latent-state shifts. Experiments show efficient adaptation across downstream tasks and combining multiple tasks, while requiring model access.

  • ICV extracts an in-context vector from demonstration examples and adds it to latent states across layers and token positions during query inference.Only the query example is then placed into the model.
  • ICV enables efficient adaptation to downstream tasks and combination of multiple tasks.
  • ICV outperforms conventional ICL and LoRA finetuning in performance and efficiency across experiments from detoxification and safety to role-playing.
  • ICV requires access to the model to add the in-context vector, which is easy to implement for open-source models.Open-source models are the main focus of the experiments and applications.

A More experimental details

The experiments use task-specific prompts, demonstrations, evaluation procedures, and implementation settings for ICL, ICV, and LoRA finetuning. ICV uses latent states from post-attention MLP features, PCA-based vectors, adaptive scaling, and jailbreak-specific decay.

  • ICL uses a simple template for style-transfer and role-playing tasks, while ICV uses the same instruction without demonstrations.Formality transfer uses a separate paraphrasing instruction.
  • LoRA finetuning uses 20 epochs, a 3e-4 learning rate, 512-token cutoff, rank 8, dropout 0.05, and alpha 16.Training uses one A100 GPU and gradient accumulation with batch size equal to the number of demonstrations.
  • ICV uses post-attention MLP features as latent states, one PCA component, and a default scaling factor λ of 0.1.
  • ICV adaptively increases modification for latent states less aligned with the task direction and decays λ to 0.1λ within five jailbreak-generation tokens.
  • Role-playing quality is automatically evaluated with GPT-3.5-Turbo using role characteristics and content relevance.
  • Jailbreak attack success is detected by checking whether generated sequences contain specified refusal-related tokens.

B Datasets and demonstration examples

The evaluation datasets use selected demonstrations for ICL, ICV, and LoRA finetuning, with additional examples for tasks lacking large matched datasets. The appendix also derives ICL’s attention-based latent-state interpretation.

  • Table 7 overviews datasets used for quantitative evaluation, with demonstrations randomly selected from training sets when available.
  • For dialogue safety and style transfer, the authors create demonstration examples using GPT-4 because large-scale matched datasets may be difficult to find.
  • The appendix represents ICL inputs as concatenated demonstrations and queries processed through attention using learnable key, query, and value matrices.

D Latent states for text-classification datasets

The paper examines whether latent states support traditional text-classification and multiple-choice in-context tasks, using AGNews as a text-classification example. A PCA-reduced 1-NN classifier on latent states substantially outperforms vanilla ICL.

  • AGNews is used to demonstrate latent-state effectiveness for text classification, with one demonstration provided for each class.
  • The method obtains latent states for class demonstrations and a new query, reduces them with PCA to the number of classes, then assigns the nearest demonstration’s class using 1-NN.
  • The latent-state method substantially outperforms vanilla ICL by a noticeable margin on the reported benchmark.
  • Prior style-transfer methods use prompts or trained style extractors, whereas ICV does not require modifying model architectures.

E Proof of lemma 1

The proof characterizes the relevant direction through the sample covariance matrix and its spectral properties, then derives a gradient expression. The section also includes captions for demonstration-example and AGNews result tables.

  • The proof begins by expressing the objective in terms of the sample covariance matrix ΣD of D.
  • Because ΣD is symmetric, the spectral theorem is used to analyze its eigenvectors.
  • The first eigenvector of ΣD is identified as the first principal direction maximizing the sample variance of D.
  • Tables 11 and 12 list demonstration examples used for quantitative jailbreak evaluations, with results reported in Table 6.
  • Table 13 reports results for conventional ICL and latent-states kNN on the AGNews dataset.
  • The objective in Eq. (3) is rewritten before deriving its gradient with respect to h.
Loading 2311.06668v3…