Research map / Image generation / Fast diffusion sampling
Composition and Alignment of Diffusion Models using Constrained Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a constrained-optimization framework unifying alignment (reverse-KL-constrained reward alignment) and composition (reverse-KL product composition, forward-KL mixture composition) of diffusion models, solved with Lagrangian primal-dual training.
Problem
Alignment and composition of pretrained diffusion models rely on ad hoc weighted averages of rewards/score functions, making trade-offs hard to tune and not guaranteeing that all desired properties are satisfied.
Approach
Recasts alignment as minimizing reverse KL to the pretrained model subject to expected reward constraints, and composition as minimizing the maximum reverse KL to each pretrained model; derives closed-form tilted solutions (reward-tilted distribution and weighted geometric product) with strong duality, and trains score models via dual sub-gradient ascent.
Key results
- Constrained alignment with MPS, saturation, and local contrast constraints achieves near-feasible reward levels where equal-weight training lets MPS decrease.
- In multi-aesthetic alignment, constrained training improves all rewards up to thresholds including hps and yields smaller KL to the pretrained model than unconstrained weighting.
- Constrained product composition of LoRA models finetuned on different rewards improves rewards across all models versus equal weights, which drop up to 30% on some rewards.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use equal weights for MPS, saturation, and local-contrast objectives instead of constrained alignment | Stable Diffusion v1.5 alignment · MPS, local contrast, and saturation constraints | MPS reward | Hurt |
| enforce reward constraints for MPS, saturation, and local contrast | Stable Diffusion v1.5 alignment · MPS, local contrast, and saturation constraints | constraint satisfaction: a 50\% improvement | Helped |
| constrain all aesthetic rewards rather than assign fixed weights | Stable Diffusion v1.5 alignment · Multiple aesthetic constraints | reward scores | Helped |
| constrain all aesthetic rewards rather than assign fixed weights | Stable Diffusion v1.5 alignment · Multiple aesthetic constraints | KL divergence to pretrained model | Helped |
| compose reward-finetuned models with arbitrary weights rather than constrained weights | Stable Diffusion product composition · Models finetuned on different rewards | individual reward scores: up to 30\% | Hurt |
| select composition weights with constrained composition rather than equal weights | Stable Diffusion product composition · Models finetuned on different rewards | reward scores across models | Helped |
| select composition weights with constrained composition rather than equal weights | Stable Diffusion concept composition · Two text-to-image concepts | minimum CLIP and BLIP text-to-image similarity | Helped |
Surprising findings
- Equal-weight composition of multiple models finetuned on different rewards can drop some reward values by up to 30%.
- With equal weights, the MPS reward decreases during alignment when regularizers are added.
Limitations
- Theory relies on convexity in distribution space; non-convex score-space optimality and convergence are not established.
- Experiments limited to image generation; mixed alignment and composition constraints not tested.
- Point-wise KL in composition uses a surrogate score for high-dimensional images, and the point-wise KL estimator is approximate with discretization error O(1/T).
Experimental setup
Metrics: reward values KL divergence to pretrained model CLIP similarity BLIP similarity
Baselines: equal-weight composition equal-weight reward regularization single combined prompt generation
Builds on: AlignProp DDIM LoRA
Closest papers by meaning
- Constrained Diffusion Models via Dual Training (NeurIPS 2024 · similarity 0.91)
- Test-time Alignment of Diffusion Models without Reward Over-optimization (ICLR 2025 · similarity 0.90)
- Inference-Time Alignment of Diffusion Models with Direct Noise Optimization (ICML 2025 · similarity 0.90)
- Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards (CVPR 2025 · similarity 0.89)
- Constrained Diffusion with Trust Sampling (NeurIPS 2024 · similarity 0.89)
- Fast constrained sampling in pre-trained diffusion models (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.