Title: Online Budget-Aware Guidance for Blockwise Discrete Text Diffusion
Primary area: generative models
Keywords: Diffusion Language Models, Constrained Generation, Online Control, Scheduler-free, BALE

Abstract: Diffusion language models (DLMs) are emerging as strong alternatives to autoregressive decoders, combining iterative refinement, global context access, and natural controllability. Yet controllable text diffusion remains underexplored: unlike vision, where plug-and-play guidance is central, existing methods for text rely on fixed or decaying step-size schedules that are brittle across prompts, constraints, and remasking strategies. The problem is amplified in modern blockwise DLMs, where gradient curvature and coverage fluctuate across spans, making schedule-based heuristics unstable. 
We present BALE(Budget-Aware Logit Editing), a plug-and-play controller that replaces heuristic schedules with an online, budget-aware rule. At each step, BALE fuses multi-constraint gradients and selects guidance strength by balancing a residual-targeted movement rule, stability caps based on an EMA Lipschitz proxy and a KL surrogate. Beyond per-step control, BALE allocates a fixed sequence-level budget across blocks and reallocates dynamically at boundaries. A discrete-time KL bound and a continuous-time SDE view explain why BALE stabilizes blockwise decoding. 
On LLaDa-8B, BALE improves lexical (must-have, forbidden), semantic (sentiment, formality), and combined controls over prompt tuning and schedule-based baselines, yielding higher satisfaction rates at matched or lower perplexity and stronger human/LLM preference. BALE is lightweight, scheduler-free, and advances controllable text diffusion toward practical deployment.

TL;DR: We present BALE, a plug-and-play controller for diffusion LMs that adaptively allocates guidance, improving controllability without hurting fluency.