Ambient Dataloops: Generative Models for Dataset Refinement

TL;DR

Ambient Dataloops tackles learning a data distribution when training data vary in quality by coupling dataset refinement with diffusion-model training. The approach iteratively trains a diffusion model on a noisy dataset using a corruption-aware Ambient objective and then performs posterior sampling to partially denoise the data, creating progressively higher-quality training sets for subsequent loops. Theoretical analysis shows that, under reasonably accurate score estimates, dataset looping can reduce estimation error, and experiments demonstrate state-of-the-art gains in unconditional and text-conditioned image generation and in de novo protein design. The method is particularly impactful for learning from noisy or data-scarce domains and suggests broad applicability beyond vision to scientific datasets and structured design tasks.

Abstract

We propose Ambient Dataloops, an iterative framework for refining datasets that makes it easier for diffusion models to learn the underlying data distribution. Modern datasets contain samples of highly varying quality, and training directly on such heterogeneous data often yields suboptimal models. We propose a dataset-model co-evolution process; at each iteration of our method, the dataset becomes progressively higher quality, and the model improves accordingly. To avoid destructive self-consuming loops, at each generation, we treat the synthetically improved samples as noisy, but at a slightly lower noisy level than the previous iteration, and we use Ambient Diffusion techniques for learning under corruption. Empirically, Ambient Dataloops achieve state-of-the-art performance in unconditional and text-conditional image generation and de novo protein design. We further provide a theoretical justification for the proposed framework that captures the benefits of the data looping procedure.

Figures (9)

• Figure 1: Dataset and model evolution across loops of our method.$D_0$ shows synthetically generated images from DiffusionDB diffdb, a dataset used for text-to-image generative modeling. These images have artifacts due to learning errors of the underlying model. We train a model on this dataset, $M_1$, that we use to improve its own training set, leading to a "restored" dataset $D_1$. Successive iterations of this process lead to a co-evolution of both the model and the dataset -- see dataset $D_2$ and model $M_1$ respectively. We avoid catastrophic self-consuming loops by accounting for learning errors at each iteration using Ambient Diffusion omniambient_diffusion techniques for learning from imperfect data.
• Figure 2: Illustration of the Ambient Dataloops framework.At each loop, we are given points that can be used to train the diffusion model at certain noise levels. We train a model on this noisy dataset using Ambient Diffusion (green), and then we use it to improve the dataset through posterior sampling (blue).
• Figure 3: Multiple loops and ablation on the rate of progress. The horizontal axis is the noise level we denoise a corrupted CIFAR-10 dataset after $k$ loops, where $k$ changes for each one of the lines. Going too fast or too slow is suboptimal. There is a point after which reducing the dataset further only hurts (madness regime) because the current model has reached its denoising capacity. FID is always computed wrt to the original clean CIFAR-10.
• Figure 4: (a) Pareto frontier for protein backbone design. (b) Example point refinement procedure.
• Figure 5: EDM loss vs noise level for Ambient Diffusion Omni omni vs Ambient Loops (Loop 1) models across four corruptions on Cifar-10. Loops models have increased denoising performance across all noise levels for all four corruptions.

Theorems & Definitions (4)

• Lemma 1: Contractive transformations lead to better learning
• Lemma 2: Non-Expansion and Contraction under Log-Sobolev
• Lemma 1: Contractive transformations lead to better learning
• Lemma 2: Non-Expansion and Contraction under Log-Sobolev