ILRR: Inference-Time Steering Method for Masked Diffusion Language Models

TL;DR

ILRR introduces a learning-free, inference-time steering framework for discrete diffusion language models by aligning latent activations with a reference sequence, enabling controllable generation with minimal overhead. By operating in the model’s latent space and using a semantic extractor, ILRR steers high-level semantics while preserving syntax, achieving significant attribute control (e.g., toxicity and sentiment) under tight compute budgets. The Spatially Modulated Steering extension enables effective guidance for long texts from short references by distributing the steering signal across the sequence. Empirical results on LLaDA and MDLM show consistent improvements over baselines with modest overhead, highlighting latent-activation intervention as a practical, scalable approach to controllable non-autoregressive language generation.

Abstract

Discrete Diffusion Language Models (DLMs) offer a promising non-autoregressive alternative for text generation, yet effective mechanisms for inference-time control remain relatively underexplored. Existing approaches include sampling-level guidance procedures or trajectory optimization mechanisms. In this work, we introduce Iterative Latent Representation Refinement (ILRR), a learning-free framework for steering DLMs using a single reference sequence. ILRR guides generation by dynamically aligning the internal activations of the generated sequence with those of a given reference throughout the denoising process. This approach captures and transfers high-level semantic properties, with a tunable steering scale enabling flexible control over attributes such as sentiment. We further introduce Spatially Modulated Steering, an extension that enables steering long texts using shorter references by regulating guidance intensity across the sequence. Empirically, we demonstrate that ILRR achieves effective attribute steering on LLaDA and MDLM architectures with a minor computational overhead, requiring only one additional parallel forward pass per denoising step. Under the same compute budget, ILRR improves attribute accuracy over comparable baselines by 10$\%$ to 60$\%$ points, while maintaining high generation quality.

Figures (5)

• Figure 1: Overview of ILRR. (a) illustrates the steering process within the latent space; (b) example of positive and negative sentiment generations.
• Figure 2: Steering accuracy (positive sentiment) across steering scales $\alpha$ compared to 4-gram overlap with the reference text.
• Figure 3: Steering accuracy (positive sentiment) across steering step sets $T_s$ compared to 4-gram overlap with the reference text.
• Figure 4: Steering accuracy (negative sentiment) across model's layers.
• Figure 5: Steering accuracy (toxicity) across pooling kernel sizes $k$ compared to ppl (GPT2-L) and 4-gram overlap with the reference text.