Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Training Latent Diffusion Models with Interacting Particle Algorithms

Tim Y. J. Wang, Juan Kuntz, O. Deniz Akyildiz

TL;DR

This work reframes end-to-end training of latent diffusion models as the minimization of a free-energy functional and derives a gradient flow in a Wasserstein-2 geometry for latent distributions alongside Euclidean geometry for model parameters. By approximating this flow with an interacting-particle system, the authors obtain IPLD, an encoder-free, scalable training algorithm that supports subsampling, distributed computation, preconditioning, momentum, and KL annealing, with provable convergence and non-asymptotic error bounds. Theoretical results guarantee exponential convergence under standard assumptions and provide finite-sample error controls for the particle-based approximation. Empirically, IPLD matches or surpasses the closest VI analogue on synthetic tasks and demonstrates competitive image-modeling performance with improvements as the particle count grows, while highlighting practical trade-offs and potential for integration with broader end-to-end diffusion-model advances.

Abstract

We introduce a novel particle-based algorithm for end-to-end training of latent diffusion models. We reformulate the training task as minimizing a free energy functional and obtain a gradient flow that does so. By approximating the latter with a system of interacting particles, we obtain the algorithm, which we underpin theoretically by providing error guarantees. The novel algorithm compares favorably in experiments with previous particle-based methods and variational inference analogues.

Training Latent Diffusion Models with Interacting Particle Algorithms

TL;DR

This work reframes end-to-end training of latent diffusion models as the minimization of a free-energy functional and derives a gradient flow in a Wasserstein-2 geometry for latent distributions alongside Euclidean geometry for model parameters. By approximating this flow with an interacting-particle system, the authors obtain IPLD, an encoder-free, scalable training algorithm that supports subsampling, distributed computation, preconditioning, momentum, and KL annealing, with provable convergence and non-asymptotic error bounds. Theoretical results guarantee exponential convergence under standard assumptions and provide finite-sample error controls for the particle-based approximation. Empirically, IPLD matches or surpasses the closest VI analogue on synthetic tasks and demonstrates competitive image-modeling performance with improvements as the particle count grows, while highlighting practical trade-offs and potential for integration with broader end-to-end diffusion-model advances.

Abstract

We introduce a novel particle-based algorithm for end-to-end training of latent diffusion models. We reformulate the training task as minimizing a free energy functional and obtain a gradient flow that does so. By approximating the latter with a system of interacting particles, we obtain the algorithm, which we underpin theoretically by providing error guarantees. The novel algorithm compares favorably in experiments with previous particle-based methods and variational inference analogues.
Paper Structure (51 sections, 5 theorems, 55 equations, 10 figures, 2 tables, 2 algorithms)

This paper contains 51 sections, 5 theorems, 55 equations, 10 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Paper structure.
  4. Notation.
  5. Background on latent diffusion models
  6. Methods
  7. The gradient flow
  8. Approximating the flow and a simple algorithm
  9. A practical algorithm: IPLD
  10. Subsampling.
  11. Distributed Training.
  12. Preconditioning and momentum.
  13. Reweighting and annealing.
  14. Related Work
  15. Interacting Particle Algorithms.
  16. ...and 36 more sections

Key Result

Proposition 2.1

$(\theta_\star,\phi_\star)$ maximize $\tilde{\ell}(\theta,\phi):=M^{-1}\sum_{m=1}^M\log \tilde{p}_{\theta,\phi}(x_m)$ iff $(\theta_\star,\phi_\star,q_\star)$ minimize $\tilde{F}$ for some $q_\star$, where $\tilde{p}_{\theta,\phi}(x_m):=\int \tilde{p}_{\theta,\phi}(x,z_{0:K}) {\mathrm{d}} z_{0:K}$.

Figures (10)

  • Figure 1: Estimated MMD between the ground truth and the distribution learned with IPLD (green) and DiffusionVAE (red) for the GMM (top) and concentric circles (bottom) datasets.
  • Figure 2: Samples generated by IPLD trained with varying numbers of particles for 1,000 steps (first two components out of the $d_x=64$ many).
  • Figure 3: Samples generated with IPLD trained on CIFAR-10, CelebA64, and SVHN. The CelebA64 samples have been curated for better visualization. See \ref{['app:additional_samples']} for additional samples (including uncurated CelebA64 ones).
  • Figure 4: Linear interpolation between the particles learned on the CelebA64 training set.
  • Figure 5: Evolution of the particle cloud of IPLD trained on the CIFAR-10 and CelebA64 dataset with $5$ particles. Zoom in to view the details better.
  • ...and 5 more figures

Theorems & Definitions (7)

  • Proposition 2.1
  • Theorem 3.1
  • Theorem 3.2
  • Lemma 1
  • proof
  • Proposition A.1
  • proof