On the Generalization of Diffusion Model

TL;DR

This paper investigates the generalization capabilities of diffusion probabilistic models, defining it via mutual information, and proposes a new training objective to improve their ability to generate unseen data.

Contribution

It introduces a formal measure of generalization for diffusion models and proposes a new training objective to enhance their extrapolation ability.

Findings

Deterministic samplers produce data highly related to training set, indicating poor generalization.

Sufficient training leaves slight differences critical for generalization.

New training objective improves the model's ability to generate unseen data.

Abstract

The diffusion probabilistic generative models are widely used to generate high-quality data. Though they can synthetic data that does not exist in the training set, the rationale behind such generalization is still unexplored. In this paper, we formally define the generalization of the generative model, which is measured by the mutual information between the generated data and the training set. The definition originates from the intuition that the model which generates data with less correlation to the training set exhibits better generalization ability. Meanwhile, we show that for the empirical optimal diffusion model, the data generated by a deterministic sampler are all highly related to the training set, thus poor generalization. This result contradicts the observation of the trained diffusion model's (approximating empirical optima) extrapolation ability (generating unseen data). To understand this contradiction, we empirically verify the difference between the sufficiently trained diffusion model and the empirical optima. We found, though obtained through sufficient training, there still exists a slight difference between them, which is critical to making the diffusion model generalizable. Moreover, we propose another training objective whose empirical optimal solution has no potential generalization problem. We empirically show that the proposed training objective returns a similar model to the original one, which further verifies the generalization ability of the trained diffusion model.