Compositional Steering of Large Language Models with Steering Tokens

AI summary

Overview

Research area: Natural Language Processing — controllable generation and behavior steering of large language models.

Technical level: Intermediate. The paper assumes familiarity with self-distillation, KL-divergence training objectives, embeddings, LoRA adapters, and activation steering, though the central idea can be grasped without deep math.

Scope: The paper proposes and evaluates "steering tokens" — trainable input-space embeddings, plus a learned composition token, that steer an LLM toward multiple verifiable output constraints simultaneously without modifying the model's parameters.

What This Paper Is About

Deploying LLMs in practice usually requires controlling several output properties at once (for example, answering in Spanish, in 10–50 words, and in title case), but existing steering research mostly targets one behavior at a time. The authors introduce compositional steering tokens: individual behaviors are compressed into dedicated input-side tokens via self-distillation, and a separate, explicitly trained composition token (<and>) is learned to represent the act of combining behaviors. The goal is zero-shot compositional generalization — correctly combining behaviors, including behaviors and numbers of behaviors never seen during composition training.

Key Contributions

1. Compositional steering tokens in the input space. Behaviors expressed as natural language instructions are distilled into dedicated trainable token embeddings that are appended to the prompt while the base LLM (including its subword embeddings) stays fully frozen. Contrasting with most prior work in the activation space, the steers live in the space of input tokens.

2. A dedicated learned composition token (<and>). Instead of interpolating or concatenating individual steering vectors, the authors train a single token to encode the concept of composition itself, with behavior tokens frozen during this second training stage so the operator learns a behavior-independent composition function.

3. Demonstrated zero-shot compositional generalization. The composition token generalizes to unseen compositions of seen behaviors, to compositions involving unseen behaviors, and to compositions with an unseen number of behaviors (training only on 2-behavior pairs, then testing on 3-behavior compositions).

4. Rigorous evaluation on verifiable constraints and hybrid steering. The authors evaluate on automatically checkable behaviors, compare against instruction steering, activation steering (CAA), LM-Steer, and LoRA DARE merging, and show that combining steering tokens with natural language instructions gives further gains.

Main Findings

• Steering tokens outperform instruction steering on unseen compositions. With Qwen3-8B, the <and> composition reaches 76.9% on unseen 2-behavior compositions versus 71.8% for instruction steering (+5.1%), and 59.5% versus 54.0% on unseen 3-behavior compositions (+5.5%). On seen compositions the two are comparable (90.9% vs. 90.7% for 2-behavior; 83.1% vs. 83.7% for 3-behavior).

• An explicit composition operator is essential. Simple concatenation of behavior tokens collapses on unseen 3-behavior compositions (33.2%), and using the LLM's native "and" token reaches only 47.1%, compared with 59.6% for the learned <and> token; concatenation also shows much higher order variance (55.8% versus 25.5% on 3-behavior compositions).

• Competing steering approaches fail at composition. LoRA DARE performs well on seen compositions (81.5% for 2-behavior) but generalizes poorly to unseen ones (44.8%), while CAA (1.6% seen, 0.5% unseen) and LM-Steer (18.1% seen, 13.4% unseen) essentially collapse; response quality for these two is also degraded (1.1 and 1.3 on a 1–5 scale).

• Hybrid steering (tokens + instructions) is best overall. On the hardest generalization case — unseen 3-behavior compositions — hybrid steering achieves 62.9% accuracy with 15.2% order variance, better than either component alone. Response quality stays high (4.9) across the token-based and hybrid methods.

• Gains hold across architectures. Across the seven evaluated models (Qwen3-4B/8B, Llama3-3B/8B, SmolLM3-3B, OLMo-7B, Gemma3-4B), hybrid steering consistently gives the best results, with large gains on weaker compositional models: +12.7% for Llama3-8B on 3-behavior compositions. Qwen models are more steerable generally, while Llama models lag regardless of method (Llama3-8B: 39.5% token steering and 40.2% instruction steering on unseen 2-behavior compositions).

• Performance scales with model size. On 3-behavior compositions, pure steering improves from 59.5% at 8B to 68.0% at 14B (+8.5%), and instruction steering from 52.1% to 61.4% (+9.3%). Hybrid steering reaches 69.2% at 14B with only 6.2% order variance, compared with 18.6% at 4B.

• Training on 2-behavior pairs is sufficient for larger models. Explicitly adding 3-behavior supervision (2+3) helps variance at 8B (15.5% vs. 21.2%) but degrades accuracy at 14B (63.9% vs. 68.0%) and raises order variance (17.3% vs. 13.9%), suggesting larger models already learn composition from simpler examples.

• Orthogonality regularization is critical for zero-shot fusion. Ablating it, the "and" embedding initialization drops from 70.8% to 55.2% unseen accuracy; zero initialization with orthogonality gives the best accuracy-variance tradeoff (69.1% unseen, 6.2% order variance). Without the <and> token at all, seen accuracy falls to 73.6%, unseen to 49.7%, and order variance rises to 27.0%.

• The method is not highly sensitive to the orthogonality weight. λ = 0.5 performs best (69.1% on 2-behavior, 60.7% on 3-behavior unseen), λ = 1.0 worst (63.9%, 51.3%), with no clear trend in order variance.

• Unseen cross-category combinations drive the gains. Per-composition analysis for Qwen3-14B and Llama3-8B shows title_case paired with language or length constraints benefits most from token steering, though Llama3-8B has a failure mode on the unseen length behavior (words_70_90) where instructions are better; hybrid steering removes these failure modes.

• Individual behavior tokens are competitive on their own. Appendix D reports that steering tokens for single behaviors are, alone, as effective as instruction steering — unlike contrastive activation steers, which are dramatically worse than prompt steering when used alone.

• Beyond three behaviors, hybrid becomes necessary. In Appendix C, Qwen3-14B applied to 4-behavior compositions (with a json wrap as the fourth behavior) shows pure steering becoming order-unstable, while hybrid is the regime that scales.

Methodology in Plain English

The approach keeps the LLM completely frozen and learns only a small set of vectors that are inserted into the model's input, where the model normally sees word embeddings.

Step one — behavior tokens. For each behavior (e.g., "Answer in Spanish"), a trainable token <b> is added. A teacher copy of the same frozen LLM is prompted with the behavior instruction; the student copy sees the prompt with the steering token instead. Training minimizes the KL-divergence between the two output distributions (temperature-scaled with T = 10.0 and multiplied by T²), so the token learns to reproduce the instruction's effect. Ten paraphrases are defined per behavior (e.g., "Answer / Respond / Reply in Spanish") and one is sampled per training example to avoid overfitting to a specific wording. Behavior tokens are initialized as the mean of the frozen embeddings of the instruction's words.

Step two — the composition token. A single trainable token <and> is initialized to a zero vector and trained on pairs of behaviors, with both the LLM and all behavior tokens held frozen. The teacher sees the prompt plus both instruction texts; the student sees the prompt plus <b_i> <and> <b_j>. Freezing the behavior tokens is what forces <and> to learn a behavior-independent notion of "compose," rather than editing individual behaviors. An orthogonality loss pushes the <and> embedding to be orthogonal to all behavior embeddings seen in training, preventing it from collapsing into an existing behavior; the total loss is the distillation loss plus λ times the orthogonality loss, with λ = 0.5.

Evaluation setup. Experiments use 15 behaviors from four categories — languages (Spanish, French, Italian, Portuguese, German), length (10–50, 50–70, 70–90, 90–120 words), formatting (lowercase, uppercase, title case), and structure (1–5 sentences) — split into 11 seen behaviors for <and> training and 4 held out (German, title case, 70–90 words, 3 sentences), with half of each category held out. Prompts come from the Smoltalk dataset; Qwen3-30B-A3B-Instruct separates the core question from constraints in each prompt, and 50k prompts per behavior are used to generate training answers. Testing uses 1,000 held-out prompts per 2- and 3-behavior combination, totaling over 1M evaluations per model.

Metrics. Mean accuracy across all token orderings (the percentage of generations satisfying all behaviors), order variance (the largest accuracy gap between any two orderings of a composition), and response quality judged by Qwen3-30B-A3B-Instruct on a 1–5 Likert scale.

Why This Matters

Impact on research. The paper reframes compositional steering as learning an explicit, reusable composition operator rather than interpolating individual steering vectors — a framing that sidesteps the known destructiveness of naive vector addition and the impracticality of retraining all vectors when a new behavior is added (a limitation the authors attribute to orthogonality-constrained prior work such as Nguyen et al.). Its controlled comparison against instruction steering and cross-architecture evaluation also raises the evaluation bar for a subfield that previously relied on anecdotal composition results.

Real-world applications.

• Customer-facing assistants that must simultaneously respect language, length, and formatting constraints (e.g., localized product answers under a strict word limit).
• Structured output generation, where a model must produce a specific language, casing convention, and paragraph or sentence count for downstream parsing.
• Deploying models on constrained hardware or under latency budgets: only a handful of d-dimensional vectors need to be trained and stored per behavior, and the base model is never updated.
• Multi-tenant or multi-policy deployments, where a service swaps behavior tokens on the fly instead of maintaining separate fine-tuned models per constraint combination.

Industry relevance. Because the method avoids parameter updates, it sidesteps the compute cost of fine-tuning, the risk of catastrophic forgetting, and the 2^N fine-tuning procedures implied by arbitrary behavior combinations. The finding that a hybrid of tokens and ordinary prompts is the most robust configuration (and rescues weaker models) maps directly onto production systems that already rely on prompt-level control and want to add reliability without discarding what they have.

Future Directions

• Extending beyond verifiable constraints. The evaluation covers only automatically checkable properties (length, format, language, structure). Whether steering tokens work for subjective behaviors such as tone, style, or persona — which require human or LLM-judge evaluation — remains untested.

• Scaling to higher-order compositions. Compositions of up to three properties are evaluated, with a 4-behavior probe in the appendix. Whether accuracy degrades gracefully or rapidly as the number of simultaneous constraints grows is stated as an open question.

• Behavior toward frontier-scale models. Experiments span 3B to 14B parameters. Whether steering tokens keep improving at larger scales or hit diminishing returns — and whether hybrid methods remain necessary at those scales — is untested by this work.

• Understanding what the <and> token encodes. The authors include an interpretability analysis in Appendix E, but the deeper question of how a single learned vector represents composition, and why orthogonality regularization is so decisive for zero-shot fusion, invites further mechanistic study.

Target Audience

Researchers and engineers working on LLM controllability, behavior steering, prompt compression, and parameter-efficient adaptation will get the most from this paper. It is also relevant to practitioners who need reliable multi-constraint generation in deployed systems and are weighing instruction prompting against fine-tuning or activation steering, and to evaluation-focused researchers interested in verifiable benchmarks for compositional generalization.

Authors’ abstract

Deploying LLMs in real-world applications requires controllable output that satisfies multiple desiderata at the same time. While existing work extensively addresses LLM steering for a single behavior, \textit{compositional steering} -- i.e., steering LLMs simultaneously towards multiple behaviors -- remains an underexplored problem. In this work, we propose \emph{compositional steering tokens} for multi-behavior steering. We first embed individual behaviors, expressed as natural language instructions, into dedicated tokens via self-distillation. Contrary to most prior work, which operates in the activation space, our behavior steers live in the space of input tokens, enabling more effective zero-shot composition. We then train a dedicated \textit{composition token} on pairs of behaviors and show that it successfully captures the notion of composition: it generalizes well to \textit{unseen} compositions, including those with unseen behaviors as well as those with an unseen \textit{number} of behaviors. Our experiments across different LLM architectures show that steering tokens lead to superior multi-behavior steering of verifiable constraints (e.g., length, format, structure, language) compared to competing approaches (instructions, activation steering, and LoRA merging). Moreover, we show that steering tokens complement natural language instructions, with their combination resulting in further gains.