Efficient and accurate steering of Large Language Models through attention-guided feature learning

TL;DR

This work addresses how to reliably steer large language models by learning concept vectors from attention to prompt prefixes. It introduces an attention-guided framework that (i) dynamically selects token embeddings based on attention to the activating prefix, (ii) uses soft labels derived from token-attention for per-block concept-vector learning, and (iii) employs permutation testing to identify concept-enriched layers for steering. The method yields substantial improvements over prior work, achieving about $95\%$ steerability on a $512$-concept benchmark with Llama-3.1-8b and generalizing to models up to $70$B parameters, while revealing heterogeneous distribution of concept features across layers. This approach offers a scalable path toward efficient fine-tuning and deeper understanding of how semantic concepts are represented in industry-scale LLMs.

Abstract

Steering, or direct manipulation of internal activations to guide LLM responses toward specific semantic concepts, is emerging as a promising avenue for both understanding how semantic concepts are stored within LLMs and advancing LLM capabilities. Yet, existing steering methods are remarkably brittle, with seemingly non-steerable concepts becoming completely steerable based on subtle algorithmic choices in how concept-related features are extracted. In this work, we introduce an attention-guided steering framework that overcomes three core challenges associated with steering: (1) automatic selection of relevant token embeddings for extracting concept-related features; (2) accounting for heterogeneity of concept-related features across LLM activations; and (3) identification of layers most relevant for steering. Across a steering benchmark of 512 semantic concepts, our framework substantially improved steering over previous state-of-the-art (nearly doubling the number of successfully steered concepts) across model architectures and sizes (up to 70 billion parameter models). Furthermore, we use our framework to shed light on the distribution of concept-specific features across LLM layers. Overall, our framework opens further avenues for developing efficient, highly-scalable fine-tuning algorithms for industry-scale LLMs.

Figures (22)

• Figure 1: Overview of our concept extraction and LLM steering approach. All examples are generated on the 8 billion parameter Llama 3.1 model. (A) For concept vector extraction, given a set of prefixed prompts, we dynamically select token embeddings at each block by selecting the token with highest attention to the prefix. (B) Given token embeddings for a token $t$ and a set of prefixed and non-prefixed prompts from (A), we extract concept vectors by training a model to predict the attention from token $t$ to the prefix tokens. (C) To steer LLMs, we perturb the activations from blocks with tokens that had significant attention-to-prefix (as determined by permutation testing).
• Figure 2: Attention-guided token selection improves steering over fixed token selection for concept extraction. (A) Overview of the dataset from Aditpaper for concept extraction. There are two sets of prompts: $\mathcal{P}_c$ with a prefix to activate a concept (e.g., fear of marriage) and $\mathcal{P}_0$ without a prefix to keep the concept inactive. Assistant tokens in common to all prompts are colored. (B) Comparison of steering performance when using fixed token embeddings and using our dynamic, attention-guided token embedding selection strategy. We use RFM rfm_science for concept extraction and the $512$-concept benchmark and evaluation strategy from Aditpaper. Samples of steered responses are presented in Fig. \ref{['fig:steering_examples_hardlabels_llama_3.1_8b']}. Comparison across four other concept extraction methods (linear regression, logistic regression, differences in means, and PCA) is shown in Fig. \ref{['fig:llama-3.1-8b-allmethods']}. (C) Visualization of the distribution of tokens that were selected across the $512$ concepts (stratified by concept class).
• Figure 3: Attention-based soft labels improve over hard (binary) labels for concept vector extraction. (A) PCA visualization of heterogeneity in concept activity as measured by attention-to-prefix. Red points indicate token embeddings of prefixed prompts and gray dots indicate token embeddings of non-prefixed prompts. Dashed gray lines pair each token embedding for a non-prefixed prompt to the corresponding embedding after adding a prefix. Hue of red points indicates the magnitude of attention to the prefix. Solid and dashed arrows indicate the concept vectors learned by RFM with hard labels and soft labels, respectively. Additional visualizations are provided in Figs. \ref{['fig:pca_extra_3.1_8b']} and \ref{['fig:pca_extra_3.3_70b']}. (B) Comparison of steering with soft labels vs. hard labels for llama-3.1-8b and llama-3.3-70b. Additional comparisons for Qwen models are presented in Fig. \ref{['fig:qwen_steerability']}. (C) Examples comparing steered outputs for llama-3.1-8b using concept vectors learned via hard and soft labels. Additional generations for Llama and Qwen models are provided in Figs. \ref{['fig:steering_examples_softhard_llama_3.1_8b']}, \ref{['fig:steering_examples_softhard_llama_3.3_70b']}, \ref{['fig:steering_examples_softhard_qwen-14b']}, and \ref{['fig:steering_examples_softhard_qwen-32b']}.
• Figure 4: Concept enrichment scores (the average number of times a token had significantly high attention to the prefix across attention heads and prompts, as determined by permutation testing) identify effective blocks for steering. (A) Comparison of steering the blocks with highest concept enrichment scores (red) and lowest enrichment scores (gray) for llama-3.1-8b and llama-3.3-70b. Steering scores stratified by concept class are provided in Fig. \ref{['fig:topk_by_cc']}. (B) Heatmap of concept enrichment scores per block for both LLMs. Results for Qwen models are presented in Fig. \ref{['fig:qwen_steerability']}.
• Figure S1: Steering coefficients and blocks selected used for all models and experiment.