Mean-field Chaos Diffusion Models

TL;DR

MF-CDMs address the curse of dimensionality in score-based generative modeling by representing high-cardinality data as a large interacting-particle system and leveraging propagation of chaos to connect finite-N training to a mean-field limit. The method replaces standard score matching with a mean-field score matching objective defined on Wasserstein space, and introduces a subdivision of chaotic entropy with particle branching to enable scalable training. Theoretical analysis provides concentration and chaos-stability results for reducible score networks, and empirically MF-CDMs deliver robust, scalable generation on synthetic data and large 3D point clouds, often surpassing traditional diffusion models at high cardinalities. This framework broadens SGMs’ applicability to unstructured, high-resolution data such as large-scale point clouds and complex stochastic systems.

Abstract

In this paper, we introduce a new class of score-based generative models (SGMs) designed to handle high-cardinality data distributions by leveraging concepts from mean-field theory. We present mean-field chaos diffusion models (MF-CDMs), which address the curse of dimensionality inherent in high-cardinality data by utilizing the propagation of chaos property of interacting particles. By treating high-cardinality data as a large stochastic system of interacting particles, we develop a novel score-matching method for infinite-dimensional chaotic particle systems and propose an approximation scheme that employs a subdivision strategy for efficient training. Our theoretical and empirical results demonstrate the scalability and effectiveness of MF-CDMs for managing large high-cardinality data structures, such as 3D point clouds.

Figures (5)

• Figure 1: 3D representations of $(\nu^N, \mu)$.
• Figure 2: Illustrative Overview of Denoising MF-SDEs/WGFs. MF-SDEs governing $M$ particles are evolved with respect to vector fields $f_t^{\otimes M} + \mathbf{s}_{\theta}^{\otimes M}$ over the interval $[t_{k-1}, t_k]$, interacting with proximate particles lying in $\mathbb{B}_R$. The illustration depicts the scenario in which the particle branching function $\Psi^{\theta}$ transforms the density of $M=3$ particles into an expanded density for $N=6$ particles $(\textit{e}.\textit{g}., \text{branching ratio }~\mathfrak{b} = 2)$ following the time interval $t_k$ and result in the joint density $\varrho_t^N$.
• Figure 3: (Left)Scalability to Data Complexity. Performance comparisons with varying data dimensionality $d$ and cardinality $N$. (Right)Ablation Study on Hyperparameters. Performance variation of MF-CDMs with respect to different hyperparameters; branching ratio $b \in \{1, 2, 4, 8\}$ and number of particle branching $|\mathbb{K}^\prime| \in \{1, 2, 4, 8\}$.
• Figure 4: Qualitative Results on MedShapeNet Dataset. Both $\mu_T^{\otimes N}$ and $\nu_T^{N_{-K}}$ illustrate the target and generated 3D shapes, where displayed liver object in MedShapeNet dataset comprises a high-cardinality point-set of nearly $2.0 E^{+4}$ points.
• Figure 5: Additional Qualitative Results on MedShapeNet Dataset. We display reconstructed 3D shapes Spine L3 vertebra and Colon in MedShapeNet dataset which comprise $2.0e^{+3}$ points.

Theorems & Definitions (36)

• Definition 2.1
• Definition 2.2
• Theorem 3.1
• Corollary 3.2
• Definition 3.3
• Proposition 3.4
• Proposition 4.1
• Theorem 4.2
• Theorem 4.3
• Lemma 1.1