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

TL;DR

ILRR is a learning-free inference-time steering method for masked diffusion language models that aligns internal activations with a reference to control generated text attributes effectively.

Contribution

Introduces ILRR, a novel inference-time steering framework for DLMs that uses internal activation alignment with a reference, enabling flexible attribute control without additional training.

Findings

ILRR achieves 10-60% improvement in attribute accuracy over baselines.

ILRR requires only one extra forward pass per denoising step.

Effective attribute steering with minimal computational overhead.

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.