Composition and Alignment of Diffusion Models using Constrained Learning

TL;DR

This paper introduces a unified constrained optimization framework for aligning and composing diffusion models, ensuring they meet multiple reward constraints and combine desirable attributes effectively.

Contribution

It proposes a novel theoretical and algorithmic approach to unify alignment and composition of diffusion models through constrained optimization.

Findings

The framework guarantees models satisfy multiple reward constraints.

Empirical results show effective alignment and composition in image generation.

The approach outperforms existing methods in maintaining desired properties.

Abstract

Diffusion models have become prevalent in generative modeling due to their ability to sample from complex distributions. To improve the quality of generated samples and their compliance with user requirements, two commonly used methods are: (i) Alignment, which involves finetuning a diffusion model to align it with a reward; and (ii) Composition, which combines several pretrained diffusion models together, each emphasizing a desirable attribute in the generated outputs. However, trade-offs often arise when optimizing for multiple rewards or combining multiple models, as they can often represent competing properties. Existing methods cannot guarantee that the resulting model faithfully generates samples with all the desired properties. To address this gap, we propose a constrained optimization framework that unifies alignment and composition of diffusion models by enforcing that the aligned model satisfies reward constraints and/or remains close to each pretrained model. We provide a theoretical characterization of the solutions to the constrained alignment and composition problems and develop a Lagrangian-based primal-dual training algorithm to approximate these solutions. Empirically, we demonstrate our proposed approach in image generation, applying it to alignment and composition, and show that our aligned or composed model satisfies constraints effectively. Our implementation can be found at: \href{https://github.com/shervinkhalafi/constrained_comp_align}{https://github.com/shervinkhalafi/constrained\_comp\_align}