Steering Without Breaking: Mechanistically Informed Interventions for Discrete Diffusion Language Models

In plain words

Diffusion language models write text by starting from a scrambled version of the whole sentence and cleaning every word up in parallel over many rounds. Controlling such a model usually means nudging its internal state the same amount in every round. This paper finds that the nudge is wasteful: a topic such as sports gets decided in the first few rounds, while a feeling such as positive sentiment builds up over a much longer stretch. Pushing at the same strength the whole time over-pushes the topic and damages the text, and it also disturbs the other attributes.

The authors first look inside four such models with a tool that splits the internal state into readable features. They measure when each feature family switches on. Then they steer only during the window when each attribute is actually forming, and use more push for the attributes that are harder to move. The paper reports better-reading text at similar steering strength.

TL;DR

• Uniform steering of diffusion language models applies the same intervention at every denoising step, which the authors show lowers text quality.
• Sparse-autoencoder features reveal that attributes commit on different schedules: topic early and sharply, sentiment gradually, with model-dependent profiles §3.2.
• Adaptive steering weights each step by how much of an attribute's emergence falls in its block §4.1 eq. 4.
• Results are author-reported and mixed in places: adaptive steering often has lower target confidence than uniform, and fluency is worse on DREAM multi-attribute runs tab. 1.

Problem

Controlled generation methods for DLMs were borrowed from autoregressive models and apply one intervention at every denoising step §1. The authors argue this spends budget on steps where the target attribute is already fixed or has not started forming, and that the cost grows when several attributes are steered at once §1 §4.2. Existing alternatives either use residual-stream directions, prompts, or classifier guidance; classifier guidance needs differentiable classifiers in the denoising loop and is stated not to apply to masked discrete diffusion §5.1.

The idealised allocation argument in §4.2 assumes per-step gain and cost are known exactly. The authors state that neither is observable in practice, and that the proxy used in place of them rests on two conditions that are plausible but not formally verified §4.2.

Method

Pipeline (offline and inference)

1. Train SAEs per layer. TopK sparse autoencoders with k = 32 on residual-stream activations of four DLMs, sampled across uniformly drawn masking rates §3.1 app. C.1 app. C.3.
2. Find attribute features. Noise each contrastive text at mask ratios 0.3, 0.5 and 0.7, encode the hidden states with ReLU (not TopK), compute Cohen's d per feature, keep the top 300 by absolute d, run Mann-Whitney U tests with Bonferroni correction, then keep the top 50 per direction. Steering uses the top 20 per attribute and direction §3.1 app. C.7.
3. Store contrastive shifts. For each steered feature j, the shift is the difference of class means, Δ_j = mean(class A) − mean(class B) app. C.7.
4. Track dynamics. Generate N = 5,000 samples per model with SAE hooks at every step, keep classifier-confident samples, average the per-step activations, split the trajectory into B blocks and compute block fractions f_b, the share of total non-negative activation growth in each block §3.2 app. C.8.
5. Build the schedule. Normalise each attribute's block fractions so the peak block has weight 1: w_dyn(t) = f_b / max f_b for step t in block b §4.1 eq. 4.
6. Steer at inference. At each step and layer, encode the residual, add the scaled shifts to the selected features, clamp to non-negative, decode, and add back the SAE error §4.1 eq. 1 eq. 2 eq. 3.

Four modes come from crossing two schedules (uniform or adaptive) with two strength rules (raw α or effectiveness-ratio scaling, "+E"). The effectiveness ratio for attribute a is r(a) = α*_a / max over attributes of α*, where α*_a is the smallest strength that reaches a reference control threshold in single-attribute sweeps. The hardest attribute gets r = 1 §4.1 app. C.9.

For multiple attributes, the shifts are summed per feature before a ReLU clamp, so the update stays in the SAE's non-negative domain §4.1 eq. 3. The authors report near-disjoint feature sets across attributes, but not fully disjoint ones §4.1.

# Block fractions and temporal weights  [§3.2] [§4.1] [eq. 4]
# Input:  mean_act[t, j] (mean activation per step and feature)
#         blocks: list of (start_step, end_step) pairs
#         features: indices in F_a for one attribute
# Output: one weight per block, peak block = 1

def temporal_weights(mean_act, blocks, features):
    # Net positive emergence of the attribute's features in each block
    growth = []
    for start, end in blocks:
        growth.append(
            sum(
                max(0.0, mean_act[end, j] - mean_act[start, j])
                for j in features
            )
        )
    # Each block's share of the total emergence  [§3.2]
    total = sum(growth)
    fractions = [g / total for g in growth]
    # Scale so the strongest block gets weight 1  [eq. 4]
    peak = max(fractions)
    return [f / peak for f in fractions]

# One denoising step of adaptive steering  [§4.1] [eq. 1] [eq. 2] [eq. 3]
# Input:  x (residual at layer l, one position per token), step t
# Output: steered residual x'

def steer_step(x, t, attributes, alpha):
    # Encode with the layer's SAE  [app. C.9]
    h = relu(encode(x))
    # Start the update from the unmodified code
    h_new = h.copy()
    for a in attributes:
        # Overall scale times attribute weight times temporal weight  [eq. 2]
        strength = alpha * r[a] * w_dyn[a][block_of(t)]
        for j in F[a]:
            # Add the contrastive shift to each selected feature  [eq. 1]
            h_new[j] = h_new[j] + strength * delta[a][j]
    # Clamp to the non-negative domain after summing  [eq. 3]
    h_new = relu(h_new)
    # Keep the SAE error term and swap only the decoded part  [§4.1]
    return x - decode(h) + decode(h_new)

Analysis (proxy and bound)

• Efficiency theorem. Under the stated assumptions (s_t ≥ 0, c_t > 0, E_target > 0), the ratio of the optimal to the uniform maximum attribute shift, squared, equals (Σ s_t²/c_t)(Σ c_t) / (Σ s_t)², which equals 1 + CV² of s_t/c_t under the cost-weighted measure [thm. 1]. The adaptive gain therefore grows with cost-weighted dispersion of s_t/c_t §4.2. The identity is consistent with the Lagrangian solution in the appendix proof app. A.
• Concentration statistic. η = N Σ w_i² over normalised block fractions, where η = 1 is uniform and η = N is full concentration in one block §4.2 app. A. The paper reports η between 1.44 and 5.01 across the 12 model-attribute conditions §4.1.
• Interference bound. A decoder-Gram bound on cross-attribute coupling is stated only for disjoint feature sets. The paper says deployed sets are not fully disjoint and reports realised coupling at most 0.15 at |F| = 20 app. A.

Hyperparameters

Results

Author-reported. Setup: four DLMs; single-attribute results are means over sentiment, topic and style; multi-attribute results are means over four combinations tab. 1. Confidence is the target classifier's mean probability; PPL is Qwen2-7B perplexity (lower is better); Flu is the LLM-judge fluency on 1 to 5 (higher is better). Step is the step-function baseline with τ = 0.80.

Pairwise schedule comparison at matched points, percentage of comparisons where the first method is better tab. 2:

Other reported results:

• Adaptive beats the second-best SAE method on quality at best operating points in 22 of 28 conditions (Mann-Whitney U, p < 0.05); the four exceptions are all DREAM §5.2.
• Fixing one α across all seven conditions costs Adaptive 0.2 to 1.2 pp of mean confidence, and Uniform 1.4 to 9.2 pp §5.2 tab. 3.
• Cross-attribute interference is lower for Adaptive, with the largest gap on MDLM (Uniform up to 57 pp, Adaptive up to 33 pp) §5.2.
• Per-condition η ranking matches the within-model gain ranking for MDLM, SEDD and LLaDA (Spearman 1.0 each; pooled permutation p ≈ 0.007); DREAM 0.5 §5.2 app. A.
• E-ratio calibration helps Adaptive in 12 of 16 multi-attribute conditions, with SEDD the main exception §5.2 tab. 4.
• SAE steering at a single operating point adds 5% to 51% wall-clock time, from 1.05× on SEDD to 1.51× on LLaDA; contrastive vectors add about 1.00× app. E.6.

Paper vs code

No official code.

The arXiv abstract page lists no code, project page or comments link. The TeX sources contain no URL. GitHub repository searches for the title and the method terms, and org-scoped searches in awslabs and aws-samples (the authors are at AWS AI Labs), returned no matching repositories. Nothing was cloned.

Pitfalls

• Steering write-back is described two ways. The method writes x' = x − x̂ + x̃, keeping the SAE error term §4.1 eq. 3. The appendix's mechanism step says the hidden state is replaced by the decoded modified code app. C.9. Implement the residual-correcting form the method section gives and check the appendix against it.
• Encoder at inference is not settled. The method uses ReLU encoding §4.1 and the appendix's steering mechanism also uses ReLU app. C.9. The extraction note says SAEs use TopK in training and normal inference app. C.7. Whether steering encodes with TopK or ReLU is not stated.
• Block normalisation differs. The method divides block fractions by the peak, so the peak block gets weight 1 §4.1 eq. 4. The appendix says steering is scaled by the block fraction without that normalisation app. C.9. The two give different step strengths.
• Block count is inconsistent. Figure 2 uses B = 8 fig. 2, the dynamics table uses B = 16 for MDLM and SEDD app. C.8, and the η check uses B = 8 app. A. The step-function baseline uses 8 blocks app. C.10. The block count used for steering should be taken from the dynamics table.
• The symbol for η's block count varies. The method defines η with N, the results text uses K, and N also means number of samples elsewhere §4.2 §5.2 app. A. Assume the block count, not the sample count.
• E-ratio threshold is ambiguous. The method calls it a reference control threshold §4.1; the appendix says α* is where each attribute reaches about 80% classifier rate app. C.9. The evaluation reports continuous confidence, not rate §5.1 app. C.11. The threshold is not otherwise specified.
• PPL sign convention is unclear. The η table defines the gain as (Adaptive − Uniform) / Uniform app. A. The text reads positive gains as improvements, and PPL improvements are lower values, so the stated formula would give negative numbers. The sign convention for PPL is not stated.
• PPL is not diagnostic on LLaDA. Aggressive steering on LLaDA produces repetitive text with low PPL, and the paper reports the PPL–fluency correlation falls to 0.21 there §5.1. The paper uses fluency as the quality gate for LLaDA §5.1. Table 1 keeps LLaDA PPL for reference tab. 1.
• Adaptive is not stronger steering. Adaptive has lower target confidence than Uniform in six of eight single and multi rows of Table 1 tab. 1. The advantage is at matched confidence and matched α, not at max strength tab. 2.
• DREAM multi-attribute fluency is worse. Multi-attribute fluency is 1.34 for Adaptive against 1.69 for Uniform tab. 1. The abstract's "consistent" claim does not hold for this cell.
• Proxy assumptions are not verified. The optimality argument needs (i) active-set growth to track marginal sensitivity and (ii) cost to vary slowly. The paper calls both plausible but unverified and relies on empirical quality gains §4.2.
• Decoder-interference bound is conditional. It assumes disjoint feature sets; the deployed sets are not disjoint, so interference control rests on measurements app. A §4.1.
• Dynamics depend on classifier filtering. Block fractions come only from samples the attribute classifier confidently labels, retaining 7% to 64% app. C.8. The same classifier families are used for evaluation app. C.11. The filtering threshold for informality is 0.6, not 0.8 app. C.8.
• Sentiment timing differs between abstract and text. The abstract says sentiment emerges over about 20% of the process; the body gives 23% for the 50% mark on MDLM §3.2.
• Small-n correlation. The η ordering test uses three attributes per model; no single model's rank correlation can reach significance app. A. The pooled p-value rests on four models.
• Judge validity. The LLM judge agrees with human consensus 73% of the time on a 5-point fluency scale; disagreements cluster mid-scale §5.1. The judge is a proprietary model, so exact re-runs are hard.
• Cost. Each steering sweep needs 20 α values × 200 samples per condition and seven conditions per model; dynamics need 5,000 traced samples per model app. C.8 app. C.9. Steering adds 5% to 51% wall-clock time app. E.6.
• Model-specific settings. Remasking differs by model (low-confidence for LLaDA, entropy for DREAM) and LLaDA and DREAM use neutral prompts app. C.9. Keep these per-model when reproducing.

Lineage

• Builds on: TopK sparse autoencoders app. C.1 and SAE steering of autoregressive models (Goyal et al. 2025; Yeon et al. 2025; Cho et al. 2025; Arad et al. 2025) §1.
• Builds on: contrastive steering vectors (Zou et al. 2023; Rimsky et al. 2024) and probe-based steering (Belinkov 2022) used as baselines §5.1.
• Compared against: step-function steering over a commitment-derived interval, following Kynkäänniemi et al. 2024 §5.1 app. C.10.
• Concurrent work: SAE analysis of DLMs (Wang et al. 2026) and steering DLMs by reference-sequence alignment (Avrahami et al. 2026, ILRR, which needs a reference sequence) §1 §5.1.