Primal-Dual Guided Decoding for Constrained Discrete Diffusion

2026-05-10

AI Summary

This work addresses the challenge of effectively enforcing global attribute constraints during generation in discrete diffusion models. The authors propose a dual-guided decoding method at inference time, formulating constrained generation as a KL-regularized optimization problem. By employing mirror descent to adaptively update Lagrange multipliers, the approach applies constraint-aware additive biases to token logits at each denoising step. Notably, this method requires neither retraining nor additional model calls, supports joint satisfaction of multiple constraints, provides theoretical bounds on constraint violation, and incurs negligible degradation in generation quality. Experiments on topic-controlled text, molecular design, and music playlist generation demonstrate substantial improvements in constraint satisfaction rates while preserving domain-specific quality metrics.

Abstract

Discrete diffusion models generate structured sequences by progressively unmasking tokens, but enforcing global property constraints during generation remains an open challenge. We propose primal-dual guided decoding, an inference-time method that formulates constrained generation as a KL-regularised optimisation problem and solves it online via adaptive Lagrangian multipliers. At each denoising step, the method modifies token logits through an additive, constraint-dependent bias, with multipliers updated by mirror descent based on constraint violation. The bias arises as the optimal KL-regularised projection of the constraint, so the constrained distribution remains as close as possible to the model's unconstrained distribution while still satisfying the constraint. The method requires no retraining and no additional model evaluations beyond standard sampling, supports multiple simultaneous constraints, and provides formal bounds on constraint violation. We evaluate our approach on topical text generation, molecular design, and music playlist generation, showing that a single algorithm instantiated via domain-specific scoring functions improves constraint satisfaction while preserving relevant domain-specific quality metrics.

Problem

• constrained generation
• discrete diffusion
• global constraints
• structured sequences
• token unmasking

Innovation

• primal-dual optimization
• constrained generation
• discrete diffusion models
• KL regularization
• mirror descent