Narmeen developed, ideated and validated K-steering at Martian. Luke generated the baselines, figures and wrote this blog post. Amir proposed the research direction and supervised the project. The full interactive blog will be available closer to the publication of the complete paper on the Martian website.

TL;DR: We introduce K-steering, a steering method for language models that allows for steering in multiple simultaneous directions. Our preliminary results show it outperforms a contrastive activation addition (CAA) baseline.

Introduction

We introduce K-Steering, a method for steering language models in multiple directions simultaneously by perturbing activations according to the logits of a multilabel classifier. We experiment with steering conversational tone, showing that K-steering can cause a classifier to give high probabilities to multiple labels/tones. Qualitatively, we find that models steered with a K-steering vector respond in multiple simultaneous tones well.

Current steering methods such as ActAdd and CAA are designed to elicit a single behavior based on contrastive pairs of inputs. Unsupervised steering methods such as deep causal transcoding allow for learning many steering vectors, some of which may elicit multiple behaviors. However, there has been limited work on methods for learning steering vectors that elicit multiple specified behaviors, which K-Steering achieves. Although we experiment only with conversational tone, we expect K-steering to work for most tasks for which an accurate classifier can be trained.

As an example, we give the following steered and unsteered generations to the prompt "What is the structure of the human heart?": 

Method

We first create a synthetic dataset of questions using GPT-4o with the following prompt:

Generate {prompts_per_category} diverse, high-quality questions that fall under the category: "{category}".

IMPORTANT: Create questions that could be meaningfully answered in multiple different tones, such as:

• helpful: balanced and informative
• expert: technically precise with domain terminology
• casual: conversational and friendly
• cautious: highlighting limitations and uncertainties
• empathetic: emotionally attuned and validating
• concise: minimalist and efficient

Each generated question should:

1. Be a clear and well-formed question ending with a question mark
2. Be tone-neutral (able to be answered well in any of the tones)
3. Avoid numbering or special formatting

Focus on creating questions where the SAME question can receive meaningfully different responses depending on which tone is used to answer.

We generate 3500 prompts for these 6 tones across 18 categories such as technical questions, personal advice and nutritional advice. The model that will be steered is then given a system prompt describing a conversational tone, and is prompted with the questions generated by GPT-4o. As an example, this is the system prompt we use for the 'casual' tone:

KEEP IT CASUAL AND CONVERSATIONAL! Write like you're texting a friend - use slang, contractions, and an upbeat vibe. Feel free to throw in some humor or pop culture references. Skip the formalities and technical jargon completely. Use short, punchy sentences. Maybe even drop in a few exclamation points or emojis where it feels natural! Don't worry about covering every detail - focus on making your response fun and easy to read. Just chat about the main points as if you're having a relaxed conversation.

The tone the model is instructed to respond in forms the label for the activations collected from the model. Using those labels and the activations from the final token generated by the model, we train a $K$-label classifier

$f\left(x\right)=W^{\left(3\right)}\sigma (W^{\left(2\right)}\sigma (W^{\left(1\right)}x+b^{\left(1\right)})+b^{\left(2\right)})+b^{\left(3\right)}$,

by minimizing the cross-entropy loss.

We then devise a steering loss

$L\left(x\right)=-\frac{1}{|T|}{\sum }_{k\in T}{f}_k\left(x\right)+\frac{1}{|A|}{\sum }_{k\in A}{f}_k\left(x\right),$

where $T$ and $A$ are the target labels and the labels to avoid respectively. To steer in a single direction we give no labels to avoid and a single target label. This straightforwardly aggregates the logits for the target and avoid labels.

For each token we then update the model's activations $x$ by taking a gradient step with respect to $L\left(x\right)$, making the new activations $x-\alpha {\nabla }_{x}L\left(x\right)$, where $\alpha$ is the step size. This is similar to the non-linear steering method in Kirch et al., however they steer only with a binary classifier.

Results

To get a quantitative estimate for steering vector performance we train a classifier on a held out dataset of activations from layer 22^[1] of either Qwen2-1.5B or Llama-3.2-3B-Instruct and tone labels created with the same method used for the K-steering classifier. This classifier is never optimized against by any of the steering methods we test.

We sample 300 activation vectors from both models, and apply each of the steering vectors we learn to all 300 activation vectors. When sampling these activations we use an unseen dataset of questions, and do not prompt the model to respond in a tone. If a given steering vector consistently yields higher probability for its target labels we consider it more performant. We compare against CAA as a baseline, using the same questions answered in different tones to create a contrastive dataset^[2]. When steering toward multiple tones with CAA, we use the mean of the steering vectors for those tones.

Both CAA and K-steering use a constant to scale the steering vector applied. We denote the constant used to scale a steering vector as $\alpha$. We use the largest $\alpha$ that doesn't negatively affect the coherence of the model for both methods^[3].

We first measure the performance of CAA and K-steering at steering toward two simultaneous tones by finding the mean of the classifier probabilities for the two target labels, and consider the steering vector with the higher mean probability to have steered toward the target tones more strongly. While K-steering gives high probabilities across all tones for both models, the CAA baseline probabilities are consistently around ~0.3. Both CAA and K-steering beat the unsteered model in every tone combination.

We found that CAA struggles further when steering toward three tones, but that K-steering remains performant. A problem with both methods is that one of the tones would sometimes 'vanish', being left with very little probability. This occurred with most of the three tone combinations we tested with CAA, and occasionally with K-steering (e.g., in the figure below for Llama-3.2-3B-Instruct).

We hypothesize that the vanishing of one tone could be caused by the steering vectors becoming too noisy when averaging three CAA vectors. Likewise for K-steering, the gradient of the steering loss may not have captured information about one of the tones in the gradient step on the activations. This could be corrected by doing multiple gradient steps on the activations, and we plan to experiment with this in the future.

We find that K-Steering outperforms CAA on all three tone combinations. To measure performance, we calculate the mean classifier probability across the three target labels. A larger mean indicates that the target tones were more successfully steered toward. We observe that the CAA probabilities tend to "wash out", with lower confidences on the target tone combinations as the number of tones increases.

We compare some responses before and after steering:

A. "What can I do if I feel my mental health is deteriorating despite seeking help?"

B. "How do you perceive the role of literature in shaping societal values?"

Caveats

We note the following limitations in our methodology which we are addressing for a paper version of this blog:

1. Training a classifier on activations we are steering, may not be as reliable as classifiers trained on the model outputs or even last layer activations.
2. Currently we compute our gradients for k-steering using loss over batches, while we expect that applying gradients on a per input basis will be more performant.
3. In this blog we choose our layer arbitrarily to apply the various steering methods, rather than sweeping over layers.

Acknowledgements

We thank Martian for supporting Narmeen and Amir, and Nirmalendu Prakesh for helpful discussion and assistance in generating our synthetic datasets. We also thank Thoughtworks for supporting Amir for the second half of this project.

1. ^{^}

   This choice is mostly arbitrary. Mid-later layers seem to respond better to steering, and both Qwen2-1.5B and Llama-3.2-3B-Instruct are 28 layer models.

2. ^{^}

   Specifically, we find the mean difference of thousands of pairs of activations where the positive examples represent text in the target tone (e.g., conversational, formal, technical) and the negative examples represent the same text in alternative tones.

3. ^{^}

   For Llama-3.2-3B-Instruct, $\alpha$ can be ~2 for CAA and ~50 for K-steering without affecting coherence in open-ended questions. For Qwen2-1.5B we found that $\alpha$ could be as high as 10 for CAA and 500 for K-steering. Note that for K-steering $\alpha$ is sensitive to the batch size, so controlling for batch size the values for $\alpha$ tend to be very similar.