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

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.

## introduction large language models (llms) are increasingly deployed in high-stakes settings, necessitating reliable mechanisms to guide their behavior at inference time. while alignment methods such as reinforcement learning from human feedback (rlhf) @xcite have improved overall model helpfulness and harmlessness, they are expensive and inflexible to deploy @xcite , may inadvertently make models less safe @xcite and cause model forgetting @xcite . in contrast, activation-based steering methods offer lightweight alternatives that modify internal representations at inference time to influence attributes * equal contribution † corresponding authors: narmeen@withmartian.com, amir.abdullah@thoughtworks.com ‡ advising author such as tone, factuality, or toxicity without retraining the full model @xcite @xcite . most existing approaches focus on singleattribute control, typically using linear interventions. these include difference-in-classes methods @xcite @xcite , vector-based edits @xcite , and gradientguided control using small auxiliary classifiers @xcite . extending such techniques to multi-attribute control is non-trivial; for instance, a naive extension by simply combining steering vectors often leads to interference or degraded fluency @xcite @xcite , as well as separate hyperparameter sweeps required for computing each individual vector. more recent frameworks like @xcite and @xcite attempt to address this through orthogonality constraints and confident vector selection. however, they still assume an additive linear structure of the steering vectors and require per-attribute gating functions to be learnt on the activations at each token. in this paper, we introduce k-steering, a method that extends non-linear gradient-based steering approaches @xcite from single-attribute control to multi-attribute control settings. k-steering fills an important gap in current methods by enabling simultaneous control over multiple behaviors using a single, unified classifier. unlike prior approaches that steer each attribute independently, k-steering learns a joint representation of multiple behaviors by training a classifier on hidden activations. at inference time, it uses the classifier's gradients to adjust the model's behavior toward (or away from) a specified set of target attributes. this approach avoids the assumption that behaviors compose linearly, handles inter-attribute interactions more robustly, and scales naturally to larger attribute sets. empirically, we find that an activation classifier trained at the final layer detects a stronger presence of target class behaviors in k-steered outputs compared to baselines that average independent steering vectors. these results are supported by traditional steering evaluation methods, such as llm judges, which indicate that k-steering better resists the dilution effect of averaging multiple vectors. our contributions are three-fold: ## methodology train an mlp. we use a multi-layer perceptron (mlp) for multi-label classification with input dimension d model , two hidden layers (256 units, relu), and an output layer of size k (softmax). while we train the classifier using only the activations from the final position index at a given layer, when steering we apply it across all positions in the sequence. we train using the cross-entropy loss with the default settings of the pytorch adam optimizer, a batch size of 32, and train for 30 epochs, which is sufficient for the classifier to achieve high accuracy on our datasets. ## datasets and models model selection. for our experiments, we use the llama-3.2-3b-instruct @xcite , mistral-7b-instruct-v0.3 @xcite and olmo-2-1124-7b-instruct (olmo et al., 2024) models. these models provide a good balance between size and performance, allowing for efficient experimentation while still producing high-quality outputs. for our comparisons, we required datasets where we can (1) steer between one attribute to another, (2) steer away from an attribute, and (3) steer towards composable combinations of attributes that compose together. the datasets used by prior work such as @xcite and @xcite are effectively combinations of binary datasets, and not "truly" k attributes that are composable in a single response. to fill this gap, we use gpt-4o-mini to compile two k attribute datasets to demonstrate the versatility of k-steering. in the first case, we create tonebank, a dataset of questions that can be responded to in six conversational tones, and debatemix which can be responded to in ten rhetorical styles. for a list of the labels for tonebank, see appendix f.1. for a list of the labels for debatemix, see appendix f.2. ## single layer, one step k-steering evaluation. we first experiment with a minimal implementation of k-steering on a single layer. we intervene on the residual stream activations at one layer by taking a single gradient step in the direction of k target labels according to the loss function described in equation 1, where k is 1, 2 or 3. we compare k-steering to contrastive activation addition (caa) @xcite and deep causal transcoding (dct) (mack and turner, 2024) (for details on our baselines, see appendix m). for caa, we again intervene on the residual stream activations at a single point, computing a mean vector from 100 contrastive pairs. the positive element of a pair is the final token activations from a generation caused by a prompt instructing the model to respond according to a description of the target label. the negative element is the final token activations caused by prompting for any other label. for the labels and prompts used in each of our datasets, see appendix g. to create multi-label caa vectors, we take the mean of single-label caa vectors. if a label is to be avoided, we subtract the vector for that label before averaging. the dct baseline learns 256 steering vectors over the last 5 layers of the model. all three methods use a novel technique to calibrate the constants by which the steering vectors are scaled, which we refer to as αs. a unique α is calibrated for every label combination through a 12iteration binary search over a large range. a typical range is [0.1, 1024]. although 1024 is much larger than most optimal α values, our binary search is computationally cheap and quickly rules out excessively large candidates. this ensures robustness against the edge case where performant αs are very large. the scoring function for our search is gpt-4o-mini, which outputs an integer score between 0 and 100, representing the coherence of a generation steered by a vector scaled by some α (see appendix j for details). if the judge assigns a score above a threshold we search the upper half of the current range for a higher α; otherwise, we search the lower half. the motivating heuristic that α should be as large as possible without affecting the model's coherence, ensuring the target behaviors are maximally visible. after 12 iterations, we select the largest α deemed coherent. to find an appropriate layer to steer using caa or k-steering, we train a classifier on the activations of the final layer of a model. this classifier is identical in architecture and training to the classifiers used to compute k-steering vectors, however it is trained on a heldout set of labelled activations and is never used to compute k-steering gradient updates. the best layer is decided by the layer that has the largest positive mean difference from the unsteered model across all label combinations over 100 inputs. we depict the full evaluation in figure 2 . the use of a classifier on final layer activations addresses previous criticisms of evaluating steering methods based solely on outputs @xcite . to use the dct vectors in our evaluation, we iterate through the 256 vectors, mapping each target label to a dct vector according to which has the strongest positive effect on the classifier probabilities over 100 samples. when steering in multiple directions with dct vectors, we use the mean of the vector for the individual target labels. results. we find that k-steering outperforms caa and dct across all models and values of k, tying with dct in the @xmath0. all three evaluated methods perform better than the unsteered model, except caa in the @xmath1.2-3b on tonebank. we additionally evaluate the performance of k-steering, caa and dct at single attribute steering on truthfulqa questions in appendix d. to validate that our final layer classifier is giving reasonable scores, we select the label combinations with the largest difference in activation classifier score between k-steering and caa for each model, k, and dataset, finding whether an llm judge agrees with the classifier judgement. we sample 100 generations from an unsteered model, a model steered with caa, and a model steered with k-steering, and prompt the judge to grade from 0 to 10 the extent to which the steered generations exhibit the target labels more so than the unsteered generations. the average of these scores gives a judgement for which method caused generations that exhibit the target labels more strongly. for more details on our judge, see appendix k. as seen in ## in-depth analysis of k-steering given that single-step, single-layer k-steering already outperforms baselines, we now examine whether additional gains can be achieved by modifying key intervention parameters. in particular, we investigate the effects of steering across multiple layers at once and taking multiple gradient steps. ## projection removal in this section, we evaluate the effectiveness of the projection removal technique described in algorithm 2 on both debatemix and tonebank tasks. this method aims to steer the model away from a specific debate style or conversational tone by ablating gradient directions associated with the corresponding class logits. we compare it against the "difference-in-means" directional vector ablation baseline introduced by @xcite , referred to as caa. as shown in ## conclusion we introduced k-steering, a unified method for multi-attribute control in language models using gradients from a non-linear classifier trained on hidden activations. unlike prior approaches, k-steering enables simultaneous control across multiple behaviors without relying on linear vector composition or per-attribute tuning. our main evaluation focuses on a single-step, single-layer variant, showing that k-steering outperforms strong baselines on two new benchmarks: tonebank and debatemix. we also explore multi-step and multi-layer. while multi-step steering can enhance effectiveness, it introduces additional computational overhead. although our results highlight the flexibility and strength of classifier-guided steering, this work also raises challenges related to scalability, dataset construction, and baseline coverage, covered in the limitations section. we hope this work provides a foundation for more robust compositional behavior control in llms.