PaperlayerView exact revision on arXiv

Source-linked AI summary

UniSteer: Text-Guided Flow Matching in Activation Space for Versatile LLM Steering

Yingdong Shi, Ruiming Zhang, Changming Li, Zhiyu Yang, Kaixing Zhang, Jingyi Yu, Kan Ren

arXiv:2605.30076v1cs.CL

TL;DR

Existing activation-steering methods often depend on fixed directions or task-specific modules, limiting adaptation to fine-grained and compositional controls. UniSteer learns a natural-language-conditioned activation flow for unified steering and classification, performing strongly across five evaluated capabilities on three target LLMs.

  • Problem

    Existing activation-steering methods rely on fixed directions or task-specific modules, making fine-grained concepts and compositional constraints difficult to adapt to.

  • Method

    UniSteer learns a conditional velocity field over frozen-LLM residual-stream activations from natural-language conditions and edits activations through flow inversion.

  • Results

    Experiments across three target LLMs show strong performance across behavioral control, truthfulness steering, fine-grained concept steering, multi-constraint instruction following, and activation-space classification.

  • Takeaways & Limitations

    A single text-conditioned activation model can provide a unified interface for diverse steering targets and activation-space classification.

  • Takeaways & Limitations

    The evaluations do not fully characterize UniSteer’s effects on broader model capabilities and have not yet evaluated long-form generation.

Abstract

from arXiv · show

Activation-based control steers large language models (LLMs) by intervening on their internal representations during inference, and has emerged as an effective paradigm for controlling behaviors such as persona and style. However, existing methods often rely on fixed steering directions or task-specific intervention modules, making them difficult to adapt to fine-grained concepts and compositional constraints. We propose UniSteer, a text-guided activation flow matching model that learns a conditional distribution over residual-stream activations from natural-language conditions. Instead of fitting a separate intervention for each target behavior, UniSteer learns a universal conditional velocity field in activation space. At inference time, UniSteer performs flow inversion by partially transporting a source activation toward a latent state and regenerating it under a target textual condition before injecting it back into the frozen LLM. The same conditional model supports activation-space classification by selecting the textual label with the lowest reconstruction energy. Experiments on three target LLMs show that UniSteer provides a unified interface across behavioral control, truthfulness steering, fine-grained concept steering, multi-constraint instruction following, and activation-space classification.

1 Introduction

Activation-based control offers lightweight, modular behavior steering for frozen LLMs, but existing approaches rely on behavior-specific directions or modules. UniSteer instead uses natural-language-conditioned activation flows for unified editing and classification.

  • Activation intervention modifies frozen LLM representations during inference without updating parameters, supporting lightweight and modular behavior control.
  • Existing steering methods often require separately fitted directions or modules for each target behavior and can interfere when combining requirements.
  • UniSteer learns a conditional velocity field over activations whose editing dynamics are specified by natural-language conditions.
  • Natural-language conditions provide one interface for behavioral traits, fine-grained concepts, and compositional constraints without post-hoc component combinations.
  • The same conditional activation model supports flow-based steering and activation-space classification using conditional reconstruction energy.

2 Related Work

Prior work shows that LLM activations encode structured behavioral and conceptual information, while activation steering typically uses fixed directions or task-specific modules. Flow-based methods additionally motivate partial-inversion editing and reconstruction-score classification.

  • Probing and unsupervised representation methods identify activation directions or subspaces associated with truthfulness, knowledge, factuality, refusal, concepts, style, sentiment, and task complexity.
  • Most activation-steering methods construct fixed behavior directions from contrastive examples or learn task-specific intervention modules.
  • Prior steering approaches usually require separate fitting for each target behavior and may suffer interference when multiple requirements are combined.
  • Flow matching supports high-dimensional generation, partial-inversion editing, and classification through conditional reconstruction or likelihood scores.

3 Methodology

UNISTEER learns a text-conditioned flow over frozen-LLM residual activations, then uses flow inversion to edit activations for target conditions or classify them by reconstruction energy.

  • Text-Conditioned Activation Modeling: UNISTEER trains a conditional flow model on residual-stream activations from selected layers and token positions paired with natural-language conditions.A frozen condition model encodes each condition, while the flow model learns the corresponding activation dynamics.
  • Text-Conditioned Activation Modeling: The learned flow transports activation states between a latent prior and condition-specific residual-stream activations.Forward integration maps prior samples to activations, while reverse integration maps observed activations toward the latent prior.
  • Training Corpus: Training instances pair frozen-model activations with natural-language conditions derived from labels, metadata, or annotations, including merged strings for compositional requirements.Categorical labels can be verbalized with templates such as “Be [trait]”.
  • Activation Steering via Flow Inversion: At inference, UNISTEER transports a source activation backward under its source condition and forward under a target condition before injecting the edited activation into the frozen model.The edit strength λ controls how closely the result remains to the source activation versus how strongly it regenerates under the target condition.
  • Activation Space Classification: For classification, UNISTEER reconstructs an activation under each candidate textual condition and selects the label with the lowest conditional reconstruction energy.This reuses the same conditional activation model rather than training a separate classifier for each label set.
  • Activation Space Classification: The unified model supports activation editing and activation-space classification through natural-language conditions.The classification procedure evaluates candidate conditions using short flow-inversion reconstruction cycles.

4 Experiments

UNISTEER is evaluated as a shared activation-steering interface across five settings, three target LLMs, and both editing and classification tasks. Results show strong steering, concrete realization of textual conditions, activation-space classification, and localized multi-constraint edits.

  • Experimental settings: UNISTEER uses one shared text-conditioned model across traits, concepts, and constraint types, whereas baselines are fitted separately for each target.This evaluates generalization across natural-language behavior descriptions rather than per-task direction fitting.
  • Experimental settings: UNISTEER is evaluated across behavioral control, truthfulness steering, fine-grained concepts, multi-constraint instruction following, and activation-space classification on three target LLMs.The target models are Llama-3.2-1B-Instruct, Qwen2.5-1.5B-Instruct, and Qwen2.5-7B-Instruct.
  • Unified steering results: UNISTEER achieves the best Persona target-trait score across all three target LLMs and improves Truth*Info over the original model on all three.It is strongest on TruthfulQA for Qwen2.5-7B; on AxBench, it leads on Qwen2.5-1.5B and Qwen2.5-7B, while LoReFT leads on Llama-3.2-1B.
  • Unified steering results: Qualitative generations show UNISTEER satisfying both required RECAST phrases while preserving tone and steering neutral responses toward explicit agreement and intensified endorsement.These examples demonstrate realization of natural-language conditions in concrete generations.
  • Activation-space classification: UNISTEER predicts activation-space labels using the candidate condition with the lowest conditional reconstruction energy.Across three target LLMs, it achieves best or tied-best accuracy and the highest AUC on two models.
  • Multi-constraint editing: For start_with constraints, edits at response-start tokens align substantially more with the corresponding CAA direction than edits at middle or ending positions.The position-sensitive pattern indicates localized, constraint-aligned updates rather than a uniform perturbation across token activations.

5 Conclusion

UNISTEER provides a unified interface for LLM steering and activation-space classification through a text-conditioned activation flow model. Its evaluations span behavioral control, truthfulness steering, fine-grained concepts, multi-constraint instruction following, and classification.

  • UNISTEER provides a unified interface for LLM steering and activation-space classification.
  • Experiments across three target LLMs evaluate behavioral control, TruthfulQA Truth*Info steering, fine-grained concept steering, multi-constraint instruction following, and activation-space classification.
  • Figure 3 shows that edits for the start_with constraint align more strongly with the CAA start_with direction at start-position tokens than at middle or ending positions.

6 Limitations and Safety Discussion

The paper’s evaluations do not fully characterize UNISTEER’s effects on broader capabilities, and several important generation and reasoning settings remain unevaluated. The authors also identify safety risks from steering toward undesirable behaviors.

  • The evaluations do not fully characterize UNISTEER’s effects on broader model capabilities.
  • Long-form generation, multi-turn stability, multi-step mathematics, and planning remain unevaluated.
  • Natural-language steering could in principle amplify undesirable behaviors such as sycophancy, deception, or harmful personas.
  • The paper recommends restricting unsafe target conditions, adding condition-level safety filters, and auditing edited generations with external classifiers or human review.
  • During editing, UNISTEER partially inverts a source activation and then follows the target-conditioned flow forward, with edit strength controlled by λ.

B Training Corpus Construction

UNISTEER constructs a shared activation-condition corpus from heterogeneous supervision sources by pairing frozen-model residual-stream activations with natural-language conditions. The corpus covers behavioral, fine-grained concept, and compositional constraint supervision.

  • Training examples pair residual-stream activations from a frozen target language model with textual conditions describing behaviors, concepts, constraints, or labels.
  • Activations are extracted under teacher forcing from selected layers and response-token positions.
  • Behavioral supervision covers persona traits, truthfulness, helpfulness, harmlessness, refusal, sycophancy, and hallucination using Persona Vectors, HH-RLHF, and HelpSteer.
  • Fine-grained concept supervision uses AxBench Concept500, converting each concept description into a textual condition.
  • Constraint-following supervision uses RECAST-5 and RECAST-10, verbalizing multiple requirements into a single condition.
  • Categorical labels, scalar attributes, and rule annotations are converted into natural-language conditions with short templates.

C Implementation Details

UNISTEER is a DiT-style conditional flow model trained over residual-stream activations, using frozen target models and a frozen condition encoder. Classifier-free guidance is used during training, and optimization runs for 10 epochs on approximately 270K examples.

  • UNISTEER uses a DiT-style text-conditioned flow model whose input and output match the target LLM’s residual-stream activation dimensionality.
  • Qwen3-Embedding-0.6B encodes textual conditions, and its parameters remain frozen during training.
  • Classifier-free guidance trains the model with both textual conditions and null conditions, using the same flow-matching loss.
  • All parameters of each target LLM and the condition encoder are frozen; only the DiT-based activation flow model is trained.
  • Training samples activation states from N(0, I), samples timesteps from U(0, 1), optionally drops conditions, and optimizes the flow-matching objective.
  • 10 epochs on approximately 270K training examples use AdamW, cosine decay with linear warmup, two GPUs, and gradient accumulation of 8 steps.

D Experimental Details

Experiments use three instruction-tuned target language models: Llama-3.2-1B-Instruct, Qwen2.5-1.5B-Instruct, and Qwen2.5-7B-Instruct.

  • The evaluation includes Llama-3.2-1B-Instruct.
  • The evaluation includes Qwen2.5-1.5B-Instruct.
  • The evaluation includes Qwen2.5-7B-Instruct.

D.2 Benchmarks and Metrics

The experiments evaluate steering across behavioral, truthfulness, concept, and constraint-following tasks, alongside activation-space toxicity classification, using shared data and benchmark-specific metrics and settings.

  • Evaluation setup: All activation-steering methods use the same data mixture for training or direction extraction, covering attribute families represented in the evaluation benchmarks.
  • Benchmarks and metrics: Persona evaluation measures target-trait and coherence scores across evil, hallucinating, and sycophantic traits.The final target-trait score averages generations whose coherence score exceeds 40.
  • Benchmarks and metrics: TruthfulQA reports Truth*Info, the percentage of responses judged both truthful and informative.Responses to 817 questions are evaluated by the specified truthfulness judge model.
  • Benchmarks and metrics: AxBench combines concept relevance, instruction relevance, and fluency using their harmonic mean.Each component score ranges from 0 to 2.
  • Benchmarks and metrics: RECAST-5 and RECAST-10 evaluate instructions containing up to 5 or 10 constraints with Rule-based Constraint Satisfaction Rate.RSR requires the response to obey all rule-based constraints.
  • Benchmarks and metrics: ToxiGen evaluates activation-space classification by comparing toxic and non-toxic textual conditions using reconstruction energy.The predicted label is the condition with lower energy, and results report accuracy and ROC-AUC.

E Additional Results

Hyperparameter sweeps show that UNISTEER’s best RECAST edit strength depends on the target model and constraint count, with potential perturbation when the original model already follows constraints.

  • RECAST Hyperparameter Sensitivity: The optimal classifier-free guidance scale varies by target model and number of RECAST constraints.The analysis uses 10 flow-inversion steps and reconstructs the activation after rolling back one step.
  • RECAST Hyperparameter Sensitivity: For Llama-3.2-1B, RECAST-5 favors a relatively large guidance scale, whereas RECAST-10 favors a moderate scale.
  • RECAST Hyperparameter Sensitivity: For Qwen2.5-1.5B, the preferred guidance scale shifts from smaller-to-moderate on RECAST-5 to slightly stronger on RECAST-10.
  • RECAST Hyperparameter Sensitivity: Qwen2.5-7B improves over the original model across several guidance scales and reaches its best RECAST-10 result at a small scale.
  • RECAST Hyperparameter Sensitivity: On Qwen2.5-7B RECAST-5, UNISTEER does not consistently improve over the original model across the tested guidance range.
  • RECAST Hyperparameter Sensitivity: Figure 4 presents hyperparameter sweeps for RECAST and a trait–coherence trade-off for the Persona evil trait.
Loading 2605.30076v1…