Primal-Dual Guided Decoding for Constrained Discrete Diffusion

TL;DR

This paper introduces primal-dual guided decoding, an inference-time method for constrained discrete diffusion generation that enforces global constraints without retraining, using KL-regularised optimisation and adaptive Lagrangian multipliers.

Contribution

It formulates constrained generation as a KL-regularised optimisation problem and solves it online, supporting multiple constraints with formal bounds on violations.

Findings

Improves constraint satisfaction in text, molecular, and music generation.

Supports multiple simultaneous constraints without retraining.

Provides formal bounds on constraint violations.

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.