Decoupling Training-Free Guided Diffusion by ADMM

TL;DR

This paper introduces a novel ADMM-based framework for conditional diffusion model generation that decouples guidance and unconditional models, improving sample quality and adherence to conditions.

Contribution

It proposes a new decoupling approach using ADMM for guided diffusion, with theoretical convergence analysis and superior experimental performance.

Findings

Consistently high-quality sample generation.

Outperforms existing methods in various tasks.

Strong adherence to conditioning criteria.

Abstract

In this paper, we consider the conditional generation problem by guiding off-the-shelf unconditional diffusion models with differentiable loss functions in a plug-and-play fashion. While previous research has primarily focused on balancing the unconditional diffusion model and the guided loss through a tuned weight hyperparameter, we propose a novel framework that distinctly decouples these two components. Specifically, we introduce two variables ${x}$ and ${z}$, to represent the generated samples governed by the unconditional generation model and the guidance function, respectively. This decoupling reformulates conditional generation into two manageable subproblems, unified by the constraint ${x} = {z}$. Leveraging this setup, we develop a new algorithm based on the Alternating Direction Method of Multipliers (ADMM) to adaptively balance these components. Additionally, we establish the equivalence between the diffusion reverse step and the proximal operator of ADMM and provide a detailed convergence analysis of our algorithm under certain mild assumptions. Our experiments demonstrate that our proposed method ADMMDiff consistently generates high-quality samples while ensuring strong adherence to the conditioning criteria. It outperforms existing methods across a range of conditional generation tasks, including image generation with various guidance and controllable motion synthesis.