Compositional Steering of Large Language Models with Steering Tokens

TL;DR

This work tackles the challenge of compositional steering for large language models by introducing steering tokens that reside in the input embedding space. Behavior tokens $\mathbf{e}_b$ and a dedicated composition token $\mathbf{e}_{\texttt{<and>}}$ are learned via a two-stage self-distillation process on frozen LLMs, enabling zero-shot generalization to unseen behavior combinations and varying numbers of behaviors. Empirical results across multiple architectures show that the proposed approach outperforms instruction-based baselines on unseen compositions, with cross-model robustness and favorable scaling; combining steering tokens with natural-language instructions yields the best performance. The method provides a parameter-efficient, modular alternative to fine-tuning, while complementing prompts, and demonstrates strong generalization to new behaviors and compositions across model families and sizes.

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 control 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.

Figures (2)

• Figure 1: Illustration of our compositional self-distillation. Left:Training individual behavior steering tokens; Right:Training the composition token <and>. The LLM remains fully frozen (including its subword embeddings). We self-distil from the instruction-prompted LLM to train the respective steering token: single behavior tokens in , and the <and> token in .
• Figure 2: Average relative performance gains (%) for Qwen14B (top row) and Llama (bottom row). Top 20 behaviour combinations (2- and 3-behaviour) yielding the largest differences between methods. Green bars indicate improvements over the text baseline; red bars indicate degradations.