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

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.

💡 Research Summary

The paper tackles the problem of “steering” large language models (LLMs)—the direct manipulation of internal activations to bias model outputs toward or away from specific semantic concepts. While prior work has shown that many concepts can be represented as linear directions (concept vectors) in activation space, existing steering pipelines are notoriously brittle. Small algorithmic choices—such as which token embeddings to use, how to treat heterogeneous concept signals, and which transformer layers to perturb—can turn a seemingly unst steerable concept into a steerable one, or vice‑versa.

To address these three pain points, the authors propose an attention‑guided steering framework that (1) automatically selects the most informative token embeddings for each layer, (2) uses soft labels derived from attention scores to capture the degree of concept activation, and (3) identifies the most relevant layers via permutation‑based statistical testing. The workflow consists of three stages: (i) dataset creation with paired “prefixed” (concept‑activated) and “non‑prefixed” prompts, (ii) dynamic token selection based on the token’s attention to the prefixed tokens, and (iii) supervised feature learning where the attention score itself serves as a continuous label. The resulting concept vectors are then added (or subtracted) to the hidden states of the selected layers during inference, scaled by a tunable coefficient ε.

Key technical contributions

1. Dynamic token selection – For each transformer block ℓ, the token that exhibits the highest attention to the concept‑activating prefix is chosen, rather than relying on a fixed position (e.g., the last token). This yields embeddings that are empirically richer in concept‑specific information.
2. Soft‑label supervision – Instead of binary labels (concept present/absent), the authors use the actual attention weight as a soft label, allowing the learning algorithm to respect varying degrees of activation across prompts. This mitigates the heterogeneity of concept signals.
3. Layer relevance via permutation testing – By randomly permuting the soft labels and recomputing attention statistics, the method isolates layers where the prefix‑attention signal is statistically significant, thereby focusing steering perturbations where they matter most and reducing unnecessary computation.

The authors evaluate the framework on a benchmark of 512 diverse concepts (including affective states, political stances, safety‑related refusals, stylistic tones, etc.) across multiple model families: Llama‑3.1 (8 B), Llama‑2 (70 B), Mistral‑7 B, Falcon‑40 B, among others. Using the same evaluation protocol as the seminal work of