Where to Steer: Input-Dependent Layer Selection for Steering Improves LLM Alignment

arXiv:2604.03867 · cs.LG · Submitted 2026-04-04 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Where to Steer".

Jane: The gist: Where to Steer (W2S) is a framework that adaptively selects the intervention layer conditioned on the input,

Tom: First, who's behind it and why it matters.

Title and authors: Jane: So Tom, let's talk about who is behind this work and what they’re calling it. The paper is titled Where to Steer: Input-Dependent Layer Selection for Steering Improves LLM Alignment <ref:2604.03867#pg1>.

Tom: Yeah, so the authors are Soham Gadgil from the University of Washington, Chris Lin also from the University of Washington, and Su-In Lee also from the University of Washington <ref:2604.03867#pg1>. They are leading this research effort.

Lu: As someone working on representation learning here at Tsinghua, I’m really interested in this because it challenges the conventional wisdom about how we should be manipulating LLMs for alignment <ref:2604.03867#pg2>.

Meng: Because conventionally, we just pick a fixed layer and assume that's sufficient for most cases, right? This paper seems to be directly questioning that assumption.

Lalam: It’s about recognizing that the representation needed for alignment isn't static; it changes based on the input sequence <ref:2604.03867#pg2>. That means the 'right' place in the model isn't fixed.

Jane: Exactly. The authors are motivated by observing that representations relevant to a target behavior can be encoded at different layers depending on what you feed into the model <ref:2604.03867#pg2>. That’s the starting point for their whole argument.

Tom: So they are proposing a method, Where to Steer, to address that input dependency directly by learning a way to adapt our steering layer based on the input <ref:2604.03867#pg1>.

Lu: It’s an interesting direction because it shifts the focus from adapting *how* we steer to adapting *where* we steer, which they call this complementary axis <ref:2604.03867#pg3>.

Jane: That’s a big conceptual shift, moving from just changing the steering vector to intelligently choosing the layer where that vector has the most impact <ref:2604.03867#pg3>.

Meng: I’m curious about how they defined their input representation z and what kind of embeddings they are using to train that predictor <ref:two thousand six hundred four point zero three eight six seven#pg3. Is this just the raw token embedding, or something richer?

Tom: They use the input sequence as a sequence of tokens denoted as x, and the input to an LLM is generally a sequence of tokens denoted as x <ref:two thousand six hundred four point zero three eight six seven#pg3. So they're starting with those token embeddings.

Lalam: From my side, it means the model needs to be able to understand that different inputs require different internal representations for the same goal <ref:2604.03867#pg2>. It’s about capturing that nuance in the representation space.

Jane: So when we think about alignment, this paper suggests that our understanding of what a target behavior means needs to be input-aware, not just universally fixed <ref:2604.03867#pg1>.

Tom: And it gives us a concrete proposal for how to build that awareness into the steering process itself by learning the layer mapping function <ref:two thousand six hundred four point zero three eight six seven#pg3. That’s what they're delivering here.

Lu: It’s a solid foundation because it validates the intuition that representation layers are context-dependent, which is a key insight for any future work in this area <ref:2604.03867#pg2>.

Jane: So, moving on from the problem statement, let's look at what exactly this paper actually claims it achieved.

The paper's summary: Tom: Okay, so summarizing the W2S framework itself. They introduce a mapping function f where f goes from input embeddings to optimal steering layers, and they train this layer predictor using cross-entropy loss with L2 regularization <ref:two thousand six hundred four point zero three eight six seven#pg3.

Jane: So they are essentially training a shallow multi-layer perceptron parameterized by theta to predict the optimal layer L* given the prompt embedding z <ref:two thousand six hundred four point zero three eight six seven#pg3. That's the mechanism for their input-dependent control.

Lu: They show that this approach is theoretically sound, demonstrating that different inputs can require steering at different layers to achieve alignment with a desirable model behavior <ref:two thousand six hundred four point zero three eight six seven#pg2.

Meng: But the empirical evidence is what really sells it. They showed that this optimal layer index varies substantially across inputs in practice, and they saw this variation even in an example where different inputs required steering at layers one, two, one, and two <ref:two thousand six hundred four point zero three eight six seven#pg2.

Tom: And the performance numbers are strong. W2S consistently improves steering performance over standard fixed-layer steering across thirteen datasets with diverse target behaviors <ref:2604.03867#pg1>. This is an improvement over just tweaking the fixed layer <ref:2604.03867#pg1>.

Jane: It's not just a marginal improvement; they show improvements on nearly all behaviors for both Llama-two-7B-Chat and Qwen-one point five-14B-Chat in the in-distribution setting <ref:2604.03867#pg1>. That’s pretty broad coverage for different models.

Lalam: I can see why that matters for deployment; it means the alignment works reliably across a wider range of prompts and model architectures <ref:two thousand six hundred four point zero three eight six seven#pg1.

Lu: And they didn't just focus on the in-distribution stuff. They showed robustness under out-of-distribution shifts induced by controlled prompt perturbations, where W2S consistently outperforms fixed-layer baselines across all OOD settings <ref:two thousand six hundred four point zero three eight six seven#pg1.

Meng: Out-of-distribution testing is where things usually break down. If it works there, it’s a much more reliable system for real traffic than just relying on in-distribution performance <ref:2604.03867#pg1>.

Tom: And they even showed that W2S recovers cases that exhibit negative steerability, which is something fixed-layer steering often fails to do <ref:2604.03867#pg1>. That’s a key win for handling tricky scenarios.

Jane: So the summary boils down to this: they propose learning a layer predictor, and it empirically works better than just picking one layer, especially when things change unexpectedly <ref:two thousand six hundred four point zero three eight six seven#pg1.

Lu: It really validates that input-dependent control is an important design dimension because it’s what allows the system to handle the variability of real language <ref:two thousand six hundred four point zero three eight six seven#pg3.

Meng: So, if we look at what this means for production deployment, it suggests we need more flexible steering mechanisms, not just one standard setting <ref:two thousand six hundred four point zero three eight six seven#pg1.

The paper's improvements: Tom: Let’s talk about the methodological refinements they made to make W2S even better, because the initial version might have had some training stability issues <ref:2604.03867#pg4>.

Jane: They introduced a frequency-aware label smoothing strategy to help with that. This strategy reassigned the label L i by choosing the layer that maximizes both steerability for that sample and its global frequency across the dataset <ref:2604.03867#pg4>.

Lu: The specific formulation they used is i = arg max in T(k) c, where T(k) is the set of steerable layers and c is the frequency of that layer <ref:2604.03867#pg4>. It’s a clever way to balance per-sample optimality with global label frequency for a target behavior <ref:two thousand six hundred four point zero three eight six seven#pg4.

Meng: Balancing per-sample needs against global frequency sounds like it addresses the risk of the predictor overfitting to just one type of input prompt <ref:two thousand six hundred four point zero three eight six seven#pg4. That’s a practical concern for any production system.

Tom: And they showed this refined approach consistently outperforms all fixed-layer baselines and even standard W2S methods when they look at both LLMs and both evaluation metrics, achieving Top one performance <ref:two thousand six hundred four point zero three eight six seven#pg4. That’s a strong result for stability.

Jane: That's a big deal because it means the method is not only better but also more stable than previous versions of W2S <ref:2604.03867#pg4>. It seems to have solved the training instability issue they were worried about <ref:two thousand six hundred four point zero three eight six seven#pg4.

Lu: From a research angle, this refinement suggests that incorporating knowledge about layer frequency into the selection process is beneficial for generalization <ref:two thousand six hundred four point zero three eight six seven#pg4. It’s adding a layer of contextual awareness to the prediction <ref:two thousand six hundred four point zero three eight six seven#pg3.

Meng: So, in short, they took an initial concept and made it more robust by incorporating data frequency into the label selection process <ref:two thousand six hundred four point zero three eight six seven#pg4. That’s how you build something that actually runs reliably at scale.

Tom: It really solidifies the idea that this layer predictor isn't just a black box guessing a number; it's using statistical information about the entire dataset to make a smarter steering decision <ref:two thousand six hundred four point zero three eight six seven#pg4.

Conclusion: Jane: So we’ve talked about the title and authors, the summary of Where to Steer, and those specific improvements in training stability. What’s the final word on what this paper means for us right now?

Tom: Basically, input-dependent layer selection is a key design axis for improving LLM alignment because it acknowledges that we can’t assume a single best layer for everything <ref:2604.03867#pg1>. The optimal layer varies substantially across inputs, and W2S provides a concrete first step toward implementing this idea <ref:two thousand six hundred four point zero three eight six seven#pg1.

Lu: I think the biggest implication is that we need to shift our design philosophy away from fixed layer steering toward methods that are inherently input-aware, which is what W2S proposes <ref:two thousand six hundred four point zero three eight six seven#pg3.

Meng: For deployment, the caution about access controls and constraints is important because we can't just roll this out blindly without considering the potential ethical implications <ref:two thousand six hundred four point zero three eight six seven#pg1.

Lalam: I just want to say that I'm sure that input-dependent layer selection for steering is a major direction for making LLMs better <ref:two thousand six hundred four point zero three eight six seven#pg1. It’s a promising path forward for the whole field.

Tom: Agreed. So, Where to Steer provides the first concrete step toward making this input-dependent layer selection a standard practice in LLM alignment <ref:two thousand six hundred four point zero three eight six seven#pg1. That's what we’ve been hearing about this research.

Jane: We have covered a lot of ground today, from the authors and the core mechanism to the stability improvements and the big picture implications of this paper on LLM alignment <ref:two thousand six hundred four point zero three eight six seven#pg1. Thanks for tuning in.

University of Washington

cs.LG

Submitted: 2026-04-04

Updated: 2026-10-15

Comments: EMNLP 2026 (Main Conference) Camera Ready

Code: https://github.com/dtch1997/steering-bench4https:

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 92/100

The gist: The gist: Where to Steer (W2S) is a framework that adaptively selects the intervention layer conditioned on the input, which consistently outperforms fixed-layer steering baselines across multiple

Key concepts

Fixed-Layer Steering: This is the traditional method where researchers apply a steering vector (a set of instructions) to a single, predetermined layer inside the LLM. The assumption is that this specific layer contains the most important information for achieving a desired outcome, regardless of what the input prompt is. This approach often fails because different prompts need different parts of the model to be adjusted.
Where to Steer (W2S) Framework: W2S introduces a system that learns which layer is best for any given input. It creates a function that maps the numerical representation of an input prompt directly to the optimal layer index within the LLM. This means if one prompt needs layer 1 and another needs layer 5, W2S identifies this difference automatically, making steering more context-aware.
Frequency-Aware Label Smoothing: This is a refinement technique used during training to make W2S more stable. Instead of just picking the absolute best layer for one sample, this method considers both whether that layer is actually steerable for that specific input and how often that layer appears across the entire dataset. This balances finding the perfect spot with ensuring the chosen spot is generally reliable.

Terminology

Steering — 7× in this paper · explained in 2 episodes
Alignment — 4× in this paper
Embeddings — 1× in this paper · explained in 2 episodes
Large Language Models (LLMs) — 1× in this paper · explained in 8 episodes

Summary

The gist: Where to Steer (W2S) is a framework that adaptively selects the intervention layer conditioned on the input, which consistently outperforms fixed-layer steering baselines across multiple LLMs and alignment behaviors.

Motivation and Problem

Existing methods typically apply steering vectors at a globally fixed layer, implicitly assuming that the optimal intervention layer is invariant across inputs The assumption that the optimal intervention layer is uniform across contexts and prompts is fundamentally limited because representations relevant to a target behavior can be encoded at different layers depending on the input For example, when steering an LLM towards positive sentiment, the relevant concept representations may differ across inputs

Theoretical and Empirical Findings

The paper theoretically demonstrates that different inputs can require steering at different layers to achieve alignment with a target behavior Empirically, the optimal steering layer varies substantially across inputs in practice This variability is shown through an example where different inputs required steering at layers 1, 2, 1, and 2 respectively Empirical analysis shows that the optimal layer index exhibits substantial variation across inputs and often deviates from the fixed layer

Where to Steer (W2S) Framework

The W2S framework is introduced to address this by learning a mapping from input embeddings to optimal steering layers The problem formulation involves learning a function f: V∗ → L, where L is the total number of intermediate LLM layers, and this function is composed of a prompt encoder and a layer predictor The W2S layer predictor f predθ is instantiated as a shallow multi-layer perceptron parameterized by θ It is trained using the cross entropy loss with L2 regularization to predict the optimal layer L∗ given the prompt embedding z

Evaluation and Results

W2S consistently improves steering performance over standard fixed-layer steering across 13 datasets with diverse target behaviors In the in-distribution setting, W2S provides improvements on nearly all behaviors for both Llama-2-7B-Chat and Qwen-1.5-14B-Chat The results are robust under out-of-distribution (OOD) shifts induced by controlled prompt perturbations, where W2S consistently outperforms fixed-layer baselines across all OOD settings Notably, W2S mitigates the failure mode of fixed-layer steering by recovering cases that exhibit negative steerability

Methodological Refinements

To improve training stability and performance, a frequency-aware label smoothing strategy is explored, where the reassigned label L˜i is given by arg max l∈T (k) c(l) (7) This formulation ensures that each sample is assigned to a layer that is both steerable for that sample and is sufficiently frequent across the dataset Across both LLMs and both evaluation metrics, frequency-aware label smoothing consistently outperforms all the fixed-layer baselines and standard W2S (Top 1)

Conclusion

Input-dependent layer selection for steering is a key design axis for improving LLM alignment The paper suggests that the optimal layer at which a target behavior is encoded varies substantially across inputs W2S is proposed as a concrete first step towards input-dependent layer selection. Future work can include extending the idea and methodology of input-dependent layer selection to multi-layer steering. The paper also cautions that real-world deployment should be accompanied by appropriate access controls and constraints due to potential ethical considerations.

--- Page 1 ---

Preprint. Under review.