CDD integrates discrete diffusion models with differentiable optimization to enforce constraints on text generation without retraining . It achieves zero constraint violations across toxicity mitigation, molecular generation, and instruction following.
Discrete diffusion models (MDLM, UDLM) generate sequences by iteratively denoising from a noise distribution. At each reverse step, CDD inserts an Augmented Lagrangian Method (ALM) projection that modifies the denoiser's predicted token distributions to satisfy user-defined constraints:
using Gumbel-Softmax relaxation for differentiability and ALM for constraint enforcement.
cdd/
├── samplers/
│ └── cdd_sampler.py # Core ALM projection + sampling loop (§4)
├── constraints/
│ ├── toxicity.py # Toxicity mitigation constraint (§5.1)
│ ├── molecular.py # SA + Novelty constraints (§5.2)
│ └── instruction.py # Counting + Lexical constraints (§5.3)
├── models/
│ └── load_models.py # MDLM/UDLM model loading utilities
├── experiments/
│ ├── toxicity_mitigation.py # Full toxicity experiment
│ ├── molecular_generation.py # Full molecular experiment
│ ├── instruction_following.py # Full instruction experiment
│ └── train_surrogates.py # Surrogate model training
└── utils/
├── noise_schedule.py # Diffusion math (schedules, posteriors)
└── evaluation.py # Metrics (PPL, entropy, violation rate)1 pip install -e .
2
3 # For GPU models (required for MDLM/UDLM inference):
4 pip install flash-attn --no-build-isolation
5
6 # For molecular evaluation:
7 pip install rdkit-pypi
1 import torch
2 from cdd . models import load_mdlm
3 from cdd . samplers import CDDSampler , ALMConfig
4
5 # Load pretrained MDLM
6 model , tokenizer , info = load_mdlm ( device = "cuda" )
7
8 # Create sampler (no constraints)
9 sampler = CDDSampler (
10 model = model ,
11 tokenizer = tokenizer ,
12 constraint_fn = lambda x : torch . tensor ( 0.0 ) , # No constraint
13 diffusion_type = "mdlm" ,
14 num_timesteps = 1000 ,
15 seq_length = 128 ,
16 device = "cuda" ,
17 )
18
19 result = sampler . sample ( batch_size = 1 )
20 print ( result [ "text" ] [ 0 ] )
1 from cdd . samplers import CDDSampler , ALMConfig , TOXICITY_ALM_CONFIG
2
3 # Setup toxicity constraint (§5.1)
4 from cdd . constraints import create_toxicity_constraint
5 constraint = create_toxicity_constraint ( threshold = 0.5 )
6
7 sampler = CDDSampler (
8 model = model ,
9 tokenizer = tokenizer ,
10 constraint_fn = constraint ,
11 alm_config = TOXICITY_ALM_CONFIG , # λ=0, μ=1, K=1000, M=10, η=0.2
12 diffusion_type = "mdlm" ,
13 device = "cuda" ,
14 )
15
16 # Conditional generation from toxic prompt
17 prefix = tokenizer . encode ( "The politician was accused of" , return_tensors = "pt" ) . to ( "cuda" )
18 result = sampler . sample ( batch_size = 1 , prefix_ids = prefix )
19 print ( result [ "text" ] [ 0 ] ) # Non-toxic completion
1 from cdd . models import load_udlm
2 from cdd . samplers import CDDSampler , MOLECULAR_SA_ALM_CONFIG
3
4 model , tokenizer , info = load_udlm ( "kuleshov-group/udlm-qm9" , device = "cuda" )
5
6 sampler = CDDSampler (
7 model = model ,
8 tokenizer = tokenizer ,
9 constraint_fn = sa_constraint , # SA ≤ 3.5
10 alm_config = MOLECULAR_SA_ALM_CONFIG , # λ=0, μ=1, μ_max=1000, K=1000, M=100, η=1.0
11 diffusion_type = "udlm" ,
12 num_timesteps = 1000 ,
13 seq_length = 32 ,
14 device = "cuda" ,
15 )
16
17 result = sampler . sample ( batch_size = 32 )
1 python -m cdd.experiments.toxicity_mitigation \
2 --threshold 0.5 \
3 --num_samples 1000 \
4 --device cuda
1 python -m cdd.experiments.molecular_generation \
2 --sa_threshold 3.5 \
3 --num_samples 1000 \
4 --device cuda
1 # Counting task
2 python -m cdd.experiments.instruction_following \
3 --task counting --num_samples 100
4
5 # Lexical task
6 python -m cdd.experiments.instruction_following \
7 --task lexical --num_samples 100
1 # Toxicity surrogate (GPT-Neo 1.3B on Jigsaw)
2 python -m cdd.experiments.train_surrogates --task toxicity
3
4 # SA surrogate (GPT-2 on QM9)
5 python -m cdd.experiments.train_surrogates --task sa
λ ← λ + μ·g(ỹ)
μ ← min(2μ, μ_max)1 @article{cardei2025constrained,
2 title={Constrained Language Generation with Discrete Diffusion Models},
3 author={Cardei, Michael and Christopher, Jacob K and Hartvigsen, Thomas
4 and Bartoldson, Brian R. and Kailkhura, Bhavya and Fioretto, Ferdinando},
5 journal={arXiv preprint arXiv:2503.09790},
6 year={2025}
7 }
