Compositional Steering of Large Language Models with Steering Tokens

1. Problem Statement

Deploying Large Language Models (LLMs) in real-world scenarios often requires controlling multiple output behaviors simultaneously (e.g., generating a response that is in French, under 50 words, and in lowercase). This is known as compositional steering.

Existing approaches face significant limitations:

• Fine-tuning: Updating model parameters for every combination of $N$ behaviors is computationally prohibitive ($2^N$ procedures) and risks catastrophic forgetting or negative interference.
• Natural Language Prompting: While flexible, standard instruction prompting is brittle; semantically equivalent prompts often yield inconsistent results, and performance degrades significantly when multiple constraints are concatenated.
• Activation Steering: Methods that inject vectors into the model's internal activation space (e.g., Contrastive Activation Addition) often fail to generalize to unseen combinations of behaviors or require orthogonal constraints that are difficult to scale.
• Gist Tokens: Existing token-compression methods focus on single instructions and lack mechanisms for composing multiple distinct behaviors.

The core challenge is to create a steering mechanism that allows for zero-shot composition of behaviors (including unseen behaviors and unseen numbers of behaviors) without modifying the model's internal parameters.

2. Methodology: Compositional Steering Tokens

The authors propose Steering Tokens, a method that operates entirely in the input embedding space while keeping the base LLM frozen. The approach consists of two distinct training stages using self-distillation.

A. Architecture

• Behavior Tokens ($<b>$): For each behavior $b$ (e.g., "French", "10-50 words"), a dedicated trainable embedding vector $e_b$ is learned.
• Composition Token ($$): A dedicated trainable vector $e_{<and>}$ is learned to represent the concept of "composition" itself.
• Input Format: To steer a model with behaviors $b_i$ and $b_j$, the input sequence is constructed as: [Prompt Tokens] + $e_{b_i}$ + $e_{<and>}$ + $e_{b_j}$.

B. Training Process (Self-Distillation)

The model is trained using a frozen "Teacher" (the base LLM) and a "Student" (the same LLM with steering tokens appended).

1. Stage 1: Training Individual Behavior Tokens

   • Teacher Input: Prompt + Natural Language Instruction (e.g., "Answer in French").
   • Student Input: Prompt + Behavior Token (e.g., <french>).
   • Objective: Minimize the KL-divergence between the Teacher's and Student's output distributions. The behavior token learns to mimic the instruction's effect.
   • Initialization: Behavior tokens are initialized as the mean of the embeddings of the words in their corresponding instruction.

2. Stage 2: Training the Composition Token

   • Teacher Input: Prompt + Instructions for two behaviors (e.g., "Answer in French. Use 10-50 words.").
   • Student Input: Prompt + Behavior Tokens + Composition Token (e.g., <french> <and> <words_10_50>).
   • Objective: Minimize KL-divergence. Crucially, the behavior tokens ($$, $$) are frozen during this stage. This forces the $$ token to learn the function of composition rather than just modifying specific behavior vectors.
   • Orthogonality Regularization: A loss term ($L_{orth}$) is added to force the $$ embedding to remain orthogonal to all frozen behavior embeddings. This prevents the composition token from collapsing into a representation of a single behavior and ensures it learns a distinct "combinatorial" function.

3. Key Contributions

1. Input-Space Composition: Unlike activation steering, this method operates in the input token space, enabling more effective zero-shot composition without interfering with the model's internal state.
2. Dedicated Composition Operator: The introduction of the learnable <and> token allows the model to generalize to:
   • Unseen combinations of seen behaviors.
   • Combinations involving unseen behaviors.
   • Combinations with an unseen number of behaviors (e.g., training on pairs, inferring triples).
3. Complementarity: The method demonstrates that steering tokens and natural language instructions are complementary; combining them yields superior performance.
4. Rigorous Evaluation: The paper provides a large-scale evaluation on verifiable constraints (length, language, format, structure) across multiple model families (Qwen, Llama, OLMo, Gemma) and sizes (3B to 14B), avoiding the subjectivity of human evaluation.

4. Experimental Results

The authors evaluated their method against baselines including standard prompting, activation steering (CAA), LM-Steer, and LoRA merging (DARE).

• Superior Generalization: Steering tokens with the <and> operator significantly outperformed text instructions on unseen compositions.
  • Example: On Qwen3-8B, for unseen 3-behavior compositions, the Hybrid method (Tokens + Instructions) achieved 62.9% accuracy, compared to 54.0% for instructions alone.
• Robustness to Order: While pure token steering showed some sensitivity to token order, the Hybrid approach (Tokens + Instructions) drastically reduced order variance, making the system robust to how constraints are ordered in the input.
• Scaling Laws:
  • Performance improved with model size (3B $\to$ 14B).
  • Training the <and> token only on 2-behavior compositions was sufficient for 14B models to generalize to 3-behavior compositions. Explicitly training on 3-behavior data sometimes degraded performance at larger scales, suggesting larger models learn the compositional logic more efficiently from simpler examples.
• Model Agnosticism: The method worked across diverse architectures (Qwen, Llama, OLMo), though the magnitude of improvement varied (e.g., Llama models benefited more from the Hybrid approach than Qwen models).
• Ablation Studies:
  • Removing the <and> token caused performance to collapse on unseen compositions.
  • Orthogonality regularization was critical; without it, the composition token failed to generalize effectively.
  • Zero-initialization of the <and> token with orthogonality constraints provided the best trade-off between accuracy and variance.

5. Significance and Impact

This work addresses a critical bottleneck in LLM deployment: the need for ad-hoc, multi-constraint control without the cost of fine-tuning.

• Efficiency: It offers a parameter-efficient alternative to fine-tuning, requiring only the learning of a small set of vectors ($|B| + 1$) rather than updating model weights.
• Scalability: The ability to compose unseen behaviors and varying numbers of constraints makes this approach viable for dynamic, real-world applications where constraint sets change frequently.
• Practicality: By focusing on verifiable constraints, the paper establishes a rigorous benchmark for compositional control, moving beyond anecdotal evidence common in prior work.
• Future Direction: The findings suggest that "composition" is a learnable, distinct function within the embedding space, opening new avenues for controlling LLMs via learned operators rather than just natural language or activation vectors.

In conclusion, Compositional Steering Tokens provide a robust, scalable, and efficient mechanism for multi-behavior control, demonstrating that learned input tokens can effectively generalize complex compositional logic across diverse LLM architectures.