Activation Steering for Masked Diffusion Language Models

AI summary

Overview

• Research area: Natural language processing; interpretability and inference-time control of masked diffusion language models (MDLMs), using safety refusal as the case study.
• Technical level: Intermediate. Readers need basic familiarity with transformer residual-stream activations, autoregressive LLMs, and iterative masked-token denoising.
• Scope: The paper adapts a contrastive activation-steering method (originally developed for autoregressive LLMs) to MDLMs, then characterizes where and when the extracted steering direction acts, and how far it transfers across languages and architectures.

What This Paper Is About

Autoregressive LLMs can be controlled at inference time by adding a low-dimensional "steering" direction to internal activations, but it was unknown whether the same trick works for masked diffusion language models, which generate text by iteratively denoising masked tokens under non-causal attention rather than predicting the next token. The authors extract a single refusal-related direction from contrastive harmful vs. harmless prompts and apply it globally during reverse diffusion, testing whether MDLMs admit a comparable one-dimensional control knob and how the diffusion architecture changes where and when that knob is accessible.

Key Contributions

1. Low-dimensional refusal control in MDLMs. The authors show that refusal behavior in MDLMs is governed by a single low-dimensional activation direction, and that applying it globally during reverse diffusion induces large, consistent shifts in refusal behavior.
2. Diffusion-specific steering from pre-instruction tokens. Unlike autoregressive models, effective refusal directions in MDLMs can be extracted not only from post-instruction tokens but also from pre-instruction tokens, reflecting the non-causal, parallel processing of diffusion models.
3. When and where guidance acts in diffusion. Ablations locate maximal steering leverage to early denoising steps and mid-to-late transformer layers, with the first diffusion block contributing disproportionately to behavioral control.
4. Transfer within MDLMs and architectural limits. Refusal directions transfer robustly between English and Chinese within the same MDLM, but do not reliably transfer to an autoregressive model, indicating architecture-dependent safety representations.

Main Findings

• Refusal is approximately one-dimensional. In multiple MDLMs, refusal behavior is governed by a consistent, approximately one-dimensional activation subspace. The direction is obtained from a single prompt-only forward pass with no optimization and without altering the diffusion sampling procedure.
• Steering beats prompt-based and optimization-based baselines. On JailbreakBench (100 harmful instructions across ten categories), activation steering produced the largest and most consistent shift across all three models, reducing keyword-based refusal scores on LLaDA from approximately 98% under direct prompting to 0%–16% and lowering LLaMA Guard 2 safety scores from approximately 100% to as low as 16%–25%.
• Baseline jailbreaks transfer poorly to diffusion. GCG (adversarial suffix of length 20, 200 optimization steps) barely changed behavior: on LLaDA-8B-Instruct it scored 96.0 keyword-refusal and 99.0 LLaMA Guard versus 97.0 and 99.0 for Direct. PAIR (N=5 attack streams, K=5 refinement iterations) reached 34.0 and 64.0. Slice reached 0.0 keyword-refusal but left 57.0 on the LLaMA Guard safety score.
• Full metric table (Keywords-Refusal / LLaMA Guard, lower is less safe). LLaDA-8B-Instruct: Direct 97.0/99.0; GCG 96.0/99.0; PAIR 34.0/64.0; Slice 0.0/57.0; steering from pre-instruction tokens 4.0/19.0; steering from post-instruction tokens 0.0/16.0. LLaDA-1.5: Direct 98.0/100.0; GCG 97.0/99.0; PAIR 35.0/79.0; Slice 0.0/63.0; pre-instruction steering 6.0/20.0; post-instruction steering 7.0/19.0. MMaDA-8B-MixCoT: Direct 26.0/50.0; GCG 12.0/45.0; PAIR 8.0/29.0; Slice 0.0/29.0; pre-instruction steering 12.0/25.0; post-instruction steering 16.0/31.0.
• Pre-instruction and post-instruction steering perform similarly. Across all three models the pre- and post-instruction variants produced highly similar outcomes, supporting the diffusion-specific claim that refusal-relevant information is accessible at structurally defined template tokens throughout the prompt, not only at the end-of-instruction region.
• Sensitivity is localized in layer–token space for LLaDA but not for MMaDA. Validation heatmaps for LLaDA-8B-Instruct and LLaDA-1.5 concentrate safety degradation in specific intermediate-to-late layers and a small subset of token positions, and the structure is highly consistent across the two LLaDA variants. MMaDA-8B-MixCoT instead shows broadly degraded safety with no well-defined localization. The best-performing layer–token pair per model always corresponded to a post-instruction token.
• Early denoising steps matter most. Applying steering only in the first κ steps of each generation block (First-κ) produced the strongest reduction in safety, with diminishing returns as κ grew. Last-κ schedules preserved high safety for small-to-moderate κ and degraded sharply only when κ became large enough to include relatively early steps. Periodic Every-κ interventions were intermediate and stabilized around κ=5, indicating the first denoising step accounts for most of the impact.
• The first diffusion block dominates. The First-κ-First-Block schedule (steering only the first κ steps of the first generation block) achieved an effect comparable to First-κ applied across all diffusion blocks.
• Prompt tokens drive the effect. Restricting steering to input (prompt) tokens clearly reduced keyword refusals and LLaMA Guard 2 safety scores, whereas output-only steering gave little to no improvement over the unsteered baseline. Joint input-and-output steering was strongest.
• Mid-to-late layers carry the effect. Using the best-performing vector (post-instruction tokens, layer 23), steering restricted to early layers (1–8) or early–mid layers (1–16) had minimal impact. Steering over layers 9–24 or 9–32 recovered most of the effect, while layers 25–32 alone were substantially less effective. Applying steering to layer 23 alone yielded only a partial reduction, indicating that the extraction layer is informative but not sufficient for control.
• Strong cross-lingual transfer. On LLaDA-8B-Instruct, EN→ZH steering (English vector, Chinese test set translated with NLLB-200) moved results from Direct 97.00/99.00 to 1.00/31.00 (Keywords-Refusal / LLaMA Guard). ZH→EN steering moved results from Direct 86.00/100.00 to 3.00/9.00. English- and Chinese-selected vectors also localized to nearby layer–token pairs.
• Shared intermediate representation across languages. A logit-lens-style analysis (following Schut et al., 2025) recorded top-1 predicted tokens per layer from a single forward pass and found a high fraction of Chinese-script predictions in intermediate layers for both languages, followed by a return to the input script near the final layers.
• No cross-architecture transfer. Applying the best LLaDA-8B-Instruct steering vector to the autoregressive Meta-Llama-3-8B-Instruct under the same global intervention produced no measurable change in keyword-based refusal rates or LLaMA Guard 2 safety scores versus direct prompting.

Methodology in Plain English

The authors take a recipe originally built for autoregressive models and adapt it to masked diffusion language models. First, they gather two contrastive sets of prompts — 128 harmful instructions (sampled from AdvBench, MaliciousInstruct, and TDC2023, keeping only those the model actually refuses) and 128 harmless instructions from ALPACA. They run a single forward pass over just the prompt, without simulating the denoising trajectory, and record residual-stream activations at every layer and at a small set of structural chat-template token positions. Positions before the user instruction form one set, positions after it form another; the indices are tied to the template and stay the same regardless of prompt length.

For each layer and candidate token position they average the activations separately for the harmful and harmless sets, take the difference, and normalize it to unit length. That gives a family of candidate directions. They pick the single best one using 32 harmful instructions from the HarmBench validation split (standard-behavior prompts only, excluding ones needing extra context or involving copyright), then evaluate on the held-out JailbreakBench set.

At generation time the chosen direction is applied globally: at every reverse-diffusion step, for every layer and every token position, the residual activation has its component along the steering direction subtracted out — a projection-based modification of hidden states. The same single direction is used everywhere, with no per-step optimization and no change to the diffusion sampling procedure itself.

Evaluation uses two binary metrics: a keyword-refusal score based on case-insensitive matching against a fixed list of twelve refusal substrings, and a LLaMA Guard 2 score indicating whether the completion is free of harmful content. Baselines are Direct prompting, GCG, PAIR, and Slice. The authors then run ablations that switch the intervention on and off at different denoising steps, token segments, and layer ranges, and run transfer tests across English and Chinese and across architectures.

Why This Matters

• Research impact. The work opens representation-level inference-time control for diffusion language models, a setting where prior control methods operated only at the sampling level (step-by-step guidance, auxiliary policies, repeated computation). It also identifies concrete architectural differences: pre-instruction tokens carry usable steering signal in MDLMs but typically not in autoregressive models under causal attention, and refusal directions do not carry over to an autoregressive model.

• Real-world applications:

  • Red-teaming and safety auditing of diffusion language models before deployment, by measuring how much a single direction can erode refusal behavior.
  • Building lightweight guardrails or controllable-generation hooks that bias outputs toward or away from a target behavior without retraining.
  • Cross-lingual safety evaluation, since a direction extracted in one language can be applied to another within the same MDLM.
  • Interpretability tooling that uses low-dimensional directions to probe what a model represents internally about harmful intent.

• Industry relevance. The method requires only a single prompt-only forward pass and no gradient-based optimization, so it is cheap relative to finetuning and directly relevant to teams deploying or evaluating MDLMs. The paper is explicitly dual-use: it can aid auditing but can also be misused to bypass safety, and the authors call for access controls, monitoring, and broad validation for any deployment beyond research.

Future Directions

• Determine why MMaDA-8B-MixCoT lacks the clear layer–token localization seen in the two LLaDA variants, and whether its weaker baseline refusal behavior explains the difference.
• Extend the approach beyond refusal to other controllable behaviors such as sentiment, topic, and style, which the introduction cites as motivating applications but which the paper does not test.
• Investigate why steering applied to the single extraction layer (layer 23) recovers only part of the effect, and characterize what propagation across mid-to-late layers contributes.
• Explore whether architecture-specific steering directions can be mapped or re-derived for autoregressive models, given that the MDLM-derived vector produced no measurable effect on Meta-Llama-3-8B-Instruct.

Target Audience

Researchers and engineers working on diffusion language models, activation steering, and LLM safety and alignment. It will be most useful to readers who already understand transformer residual streams and diffusion denoising and want a concrete, cheap inference-time control primitive, as well as to safety teams evaluating whether control mechanisms developed for autoregressive models carry over to the diffusion paradigm.

Authors’ abstract

Masked diffusion language models (MDLMs) generate text via iterative masked-token denoising, enabling mask-parallel decoding and distinct controllability and efficiency tradeoffs from autoregressive LLMs. Yet, efficient representation-level mechanisms for inference-time control in MDLMs remain largely unexplored. To address this gap, we introduce an activation steering primitive for MDLMs: we extract a single low-dimensional direction from contrastive prompt sets using one prompt-only forward pass, and apply a global intervention on residual-stream activations throughout reverse diffusion, without performing optimization or altering the diffusion sampling procedure. Using safety refusal as a deployment-relevant case study, we find that refusal behavior in multiple MDLMs is governed by a consistent, approximately one-dimensional activation subspace. Applying the corresponding direction yields large and systematic behavioral shifts and is substantially more effective than prompt-based and optimization-based baselines. We further uncover diffusion-specific accessibility: effective directions can be extracted not only from post-instruction tokens, but also from pre-instruction tokens that are typically ineffective in autoregressive models due to causal attention. Ablations localize maximal leverage to early denoising steps and mid-to-late transformer layers, with early diffusion blocks contributing disproportionately. Finally, in an MDLM trained on English and Chinese, extracted directions transfer strongly between English and Chinese, but do not reliably generalize to an autoregressive architecture, highlighting architecture-dependent representations of safety constraints.