Diffusion language models stop reusing one steering vector: feeding the previous prediction back lifts TruthfulQA multiple-choice average from 30.89 to 45.42

A diffusion language model generates text by repeated denoising, and each step produces a tentative prediction; TraceSteer makes the correction intervention follow that denoising trajectory instead of applying one direction throughout. Earlier activation steering methods largely reuse a predefined vector and a fixed strength designed for autoregressive models, without exploiting the way a diffusion model revises its hypotheses step by step. On LLaDA-8B-Instruct, the TruthfulQA multiple-choice average rises from 30.89 without intervention to 45.42, and open-ended TRUEINFO from 40.31 to 51.25; on Dream-7B-Instruct the average rises from 32.94 to 40.79 and TRUEINFO from 34.67 to 48.65.

The intervention direction is set jointly by the previous tentative prediction and the current hidden states, while its strength varies with prediction uncertainty and denoising progress. DF-Steering encodes the previous complete tentative response into contextual representations and relates them to the current hidden states to build a step-specific steering vector; TS-Gate learns a gate from the entropy of the previous prediction and the current denoising timestep to scale that vector. The two modules are trained jointly with factual target answers from TruthfulQA while the backbone stays frozen, plus a magnitude regularizer on the steering vector relative to the hidden states; training uses all 817 TruthfulQA instances split 70:15:15.

Steering parameters trained only on TruthfulQA transfer directly to five downstream benchmarks. Earlier methods typically need adjustment for the target task, whereas this setup adds no target-task adaptation. On COPA, StoryCloze, MMLU, WinoGrande and BoolQ, both backbones beat the no-intervention version; the five-benchmark average rises from 71.14 to 89.18 on LLaDA and from 73.93 to 87.87 on Dream.

The report covers two diffusion language model backbones, LLaDA-8B-Instruct and Dream-7B-Instruct, evaluated on TruthfulQA multiple-choice and open-ended generation plus five downstream benchmarks: COPA, StoryCloze, MMLU, WinoGrande and BoolQ. The intervention is applied at Transformer layer 16, and training keeps the backbone frozen while optimizing only DF-Steering and TS-Gate. It is aimed at readers interested in inference-time intervention and factual reliability for diffusion language models. The authors propose exploring weakly supervised and self-supervised strategies to reduce reliance on factual annotations.

The method relies on factual target answers to train the steering and gating modules; such high-quality annotations are expensive and time-consuming to obtain, especially in specialized domains, which may limit applicability to broader scenarios. Ablations show that removing magnitude regularization sharply lowers TRUE and INFO in open-ended generation, indicating that intervention magnitude needs constraining. Longer prediction histories do not help: using only the immediately preceding prediction gives the best overall performance.