Activation Steering for Masked Diffusion Language Models

TL;DR

The paper tackles the challenge of controlling masked diffusion language models at inference time without retraining. It introduces activation steering, a lightweight framework that derives low-dimensional steering directions from a single forward pass on contrastive prompts and injects them throughout the reverse-diffusion process to steer generation. The approach is evaluated on LLaDA-8B-Instruct, showing large, reliable shifts in harmful vs. harmless behavior, with residual-stream activations—especially post-MLP residuals—being the primary conduit for control. The work demonstrates practical, interpretable, and scalable controllability for MDLMs and outlines future work to broaden applicability and robustness across architectures and configurations.

Abstract

Masked diffusion language models (MDLMs) generate text through an iterative denoising process. They have recently gained attention due to mask-parallel decoding and competitive performance with autoregressive large language models. However, effective mechanisms for inference-time control and steering in MDLMs remain largely unexplored. We present an activation-steering framework for MDLMs that computes layer-wise steering vectors from a single forward pass using contrastive examples, without simulating the denoising trajectory. These directions are applied at every reverse-diffusion step, yielding an efficient inference-time control mechanism. Experiments on LLaDA-8B-Instruct demonstrate reliable modulation of high-level attributes, with ablations examining the effects of steering across transformer sub-modules and token scope (prompt vs.\ response).

Figures (6)

• Figure 1: Example generations illustrating activation steering in MDLMs. Outputs without steering are shown in green, and outputs with steering are shown in red.
• Figure 2: A general steering framework for inference-time scaling of MDLMs. Steering vectors are extracted from contrastive prompts in a single forward pass and applied during reverse diffusion.
• Figure 3: The predefined set of "refusal substrings" used to assess the refusal score for a given completion.
• Figure 4: PCA projections of prompt-level residual representations at selected transformer layers. PCA is fit using unsteered prompts only (harmless training and harmful test), and steered harmful test prompts are projected into this space without refitting. Shaded regions denote kernel density estimates, markers indicate class centroids, and arrows show centroid shifts induced by steering.
• Figure 5: Refusal and safety behavior across LLaDA layers under activation steering, varying token-level scope and layer-wise application. The left panel shows Keywords-based refusal scores and the right panel shows LLaMA Guard 2 safety scores.