Table of Contents
Fetching ...

Beyond Linear Steering: Unified Multi-Attribute Control for Language Models

Narmeen Oozeer, Luke Marks, Shreyans Jain, Fazl Barez, Amirali Abdullah

Abstract

Controlling multiple behavioral attributes in large language models (LLMs) at inference time is a challenging problem due to interference between attributes and the limitations of linear steering methods, which assume additive behavior in activation space and require per-attribute tuning. We introduce K-Steering, a unified and flexible approach that trains a single non-linear multi-label classifier on hidden activations and computes intervention directions via gradients at inference time. This avoids linearity assumptions, removes the need for storing and tuning separate attribute vectors, and allows dynamic composition of behaviors without retraining. To evaluate our method, we propose two new benchmarks, ToneBank and DebateMix, targeting compositional behavioral control. Empirical results across 3 model families, validated by both activation-based classifiers and LLM-based judges, demonstrate that K-Steering outperforms strong baselines in accurately steering multiple behaviors.

Beyond Linear Steering: Unified Multi-Attribute Control for Language Models

Abstract

Controlling multiple behavioral attributes in large language models (LLMs) at inference time is a challenging problem due to interference between attributes and the limitations of linear steering methods, which assume additive behavior in activation space and require per-attribute tuning. We introduce K-Steering, a unified and flexible approach that trains a single non-linear multi-label classifier on hidden activations and computes intervention directions via gradients at inference time. This avoids linearity assumptions, removes the need for storing and tuning separate attribute vectors, and allows dynamic composition of behaviors without retraining. To evaluate our method, we propose two new benchmarks, ToneBank and DebateMix, targeting compositional behavioral control. Empirical results across 3 model families, validated by both activation-based classifiers and LLM-based judges, demonstrate that K-Steering outperforms strong baselines in accurately steering multiple behaviors.

Paper Structure

This paper contains 39 sections, 8 equations, 3 figures, 39 tables, 2 algorithms.

Figures (3)

  • Figure 1: An illustration of gradient-based K-Steering. For an activation vector $A$, we calculate a steering loss that penalizes higher logits from a classifier on $A$ for undesired labels and rewards higher logits for desired labels. By backpropagating this loss through the classifier, we obtain the steered activations $A'=A-\alpha\Delta L$.
  • Figure 2: Illustration of our evaluation setup for comparing CAA, DCT and K-Steering. In Step 1, we perform a binary search over $\alpha$ on the layers we want to steer. If we detect that an $\alpha$ causes the steered model to behave out-of-distribution, we decrease the alpha according to the binary search, otherwise we increase it. In Step 2, we use the alphas found through our binary search to select the best layer to steer on according to a classifier on the final layer activations. This is done by finding which layer when steered gives the highest mean probability across the target labels.
  • Figure 3: Steering scores across steps for 3 groups of $\alpha$ values when steering towards the "Empirical Grounding" style and away from the "Straw Man Reframing" style. We perform the experiment on layer 14 in Llama-3.2-3B. We divide $\alpha$ into 3 ranges: early ($\alpha \leq 1.0$), middle ($1.0 < \alpha \leq 3.0$), and late ($\alpha > 3.0$). Each subplot shows how multi-step steering influences intervention effectiveness, with darker shades representing larger $\alpha$ values within each range. Highest performance is observed at $\alpha = 3.0$ and $\alpha = 2.6$ using two steps.