DHSA Achieves DPO-Level Alignment at Inference Time Using Only 256-16384 Control Parameters

DHSA: Achieving Inference-Time Preference Alignment Without Modifying Base Model Weights

Contributed by the DHSA Team

QbitAI | Official Account: QbitAI

Enabling large models to act consistently according to human expectations is the core objective of preference alignment.

Methods like RLHF and DPO optimize models through feedback. Meanwhile, existing research has found that while post-training improves the success rate of single generations, base models can generate more diverse correct answers through multiple sampling attempts.

In these scenarios, part of the alignment challenge lies in "how existing behaviors are expressed": The computation supporting the target already exists; what prevents it from stably dominating generation, and can it be directly corrected?

A research team from Chengdu University, North Carolina State University, Fudan University, Australian Catholic University, and the University of Macau investigated this issue and proposed Direct Hidden State Alignment (DHSA), treating inference-time hidden states—which carry preference computations—as a direct adaptation space.

△ RCM quantifies preference-related internal effects, while DHSA directly controls inference-time states. Sampling trajectories and Q&A examples are conceptual illustrations.

Opposing Effects Can Coexist Behind the Same Answer

For every token generated, components such as attention heads and feed-forward modules write updates into the shared residual stream, influencing subsequent calculations and ultimately shaping the probability distribution for the next token. Even if the computation supporting the target appears, it may be overshadowed by other effects; such choices occur sequentially, affecting how the entire response unfolds.

External preferences describe final behavior, but the model internally executes layered computations. Simply identifying components with a "high impact" on generation does not reveal whether they are promoting or counteracting the target. The team therefore proposed Residual Competition Maps (RCM) to map externally specified preferences back to native computations during inference.

Taking RCM-Zero as an example: For the same input, compare the component's normal write operation against a zeroed-out write, allowing the model to continue completing subsequent calculations. If the target behavior is relatively dominant over alternative behaviors during normal writing, this native effect supports the target; conversely, it counteracts the target. The difference between the two provides the direction and intensity of the effect, offering a quantifiable internal description of how preferences form during inference.

Using RCM, the team observed across different preference tasks that when a single prediction is formed, effects supporting the target and those counteracting it can coexist. The paper terms this residual competition, whose prevalence varies by task. Some components average support for the target, others average counteraction; the same component may also change its directional effect depending on the input.

This competition also appears in samples where the model has already provided a target-aligned answer. DPO training significantly reorganizes internal effects and can weaken some counteracting effects, but these effects may still persist. Output aligning better with preferences does not mean internal effects have become consistent.

△ Internal competition in ConFiQA and the reorganization brought by DPO. Threshold units represent logit-margin changes.

A knowledge conflict case is particularly intuitive: A model originally insisted on its own knowledge. By simply zeroing out the output of one attention head, it shifted to answering based on the provided reference material. Neither the input nor the base weights changed; a local state modification altered the generated response.

Integrating Alignment into Ongoing Inference

RCM provides causal evidence for intervention locations, and generation cases demonstrate how local state changes affect final behavior. Based on this, the team proposed DHSA, treating inference-time hidden states carrying preference computations as a direct adaptation space, while keeping base weights frozen.

DHSA allows control at different states, positions, and granularities. Causal Activation State Transfer (CAST) is the specific implementation presented in the paper: RCM selects a small number of preference-related interfaces, and CAST learns local adjustments via constant vectors or low-rank mappings, modifying the state before components write back to the residual stream, allowing the model to continue generating from there.

How states change and what objectives are used for learning can be designed separately. This paper trains CAST using a pairwise preference objective in the form of DPO, updating only the controller. The scope of DHSA can extend to other state transformations and learning objectives, designing control around the process of preference formation during inference.

Intervention timing can also be selected: covering the entire prompt processing and generation process, or intervening only during the prompt processing (prefill) stage, influencing subsequent generation by altering prompt computation and cached states. Controllers are separated from base weights, enabling them to be enabled or removed at inference time, and can also be stacked onto DPO models.

Few Control Parameters Achieve DPO-Level Performance

Experiments were conducted on ConFiQA, IMDb, and TL;DR, simulating RAG knowledge conflicts in actual multi-hop, multi-source fact generation, assessing sentiment control, and evaluating open-ended summarization, respectively. Models included GPT-2-large (0.774B), GPT-J (6B), Qwen2 (7B), Llama-3 and its LION version (8B); RCM localization experiments also covered Qwen2.5 (14B).

Position selection controls supported RCM's guidance role in interventions. Across three types of tasks, CAST achieved alignment levels competitive with DPO using only 256–16,384 control parameters when deployed. Some settings required intervention only during the prefill stage; transferring controllers learned via SFT to DPO models, calibrating only the intensity, further improved performance.

△ IMDb Sentiment Control. Left: Position selection controls; Right: Reward and token KL. SV/LR denote vector/low-rank controllers; black diamonds denote DPO.

△ Q&A and Summarization Results. Summarization win rates are relative to human reference summaries; T denotes transferred controllers, NT denotes trained on DPO.

These results indicate that the causal information provided by RCM can guide effective control; few local adjustments can influence content selection and compression throughout the entire response, supporting open-ended generation. Prefill intervention further demonstrates that early state adjustments can continuously affect subsequent generation.

Towards Adjustable Inference-Time Alignment

RCM and DHSA connect external preferences to measurable, intervenable internal computations. For target behaviors that are already reachable, researchers can track how native components promote or counteract their expression, and directly regulate generation along these causal clues. Mechanistic analysis thus enters practical alignment, providing a preference adaptation method for frozen bases that is independent of weights and can be enabled on demand.

Component effects vary with input, leaving room for more flexible control: Future work can explore dynamically selecting intervention positions, timing, and intensity based on current inputs and generation states. Training methods also have room for expansion: Applying reinforcement learning methods like PPO and GRPO, as well as distillation methods like MOPD, to train hidden state controllers remains to be explored and validated. Inference-time alignment is expected to adapt to more needs and further synergize with parameter optimization.

The project provides reproducible code and an extensible framework, decoupling intervention timing, position selection, control methods, and controller training, facilitating experiment orchestration.

Paper Title: Direct Hidden-State Alignment: Mapping and Controlling Preference Expression in LLMs Paper: https://arxiv.org/abs/2609.33298 Project: https://github.com/abovefiramament/DHSA

Like, Share, and Save!

We welcome your thoughts in the comments section!

— End —

[Academic Submissions] Please send emails on weekdays to: [email protected]. Include [Submission] in the subject line and tell us: Who you are, Where you come from, attach the submission content along with project/homepage links, and provide your contact information.

🎓 We will (try our best to) reply promptly :)

🌟 Star Us 🌟

Daily Updates on Cutting-Edge Tech Progress