Consistent Diffusion Meets Tweedie: Training Exact Ambient Diffusion Models with Noisy Data

Consistent Diffusion Meets Tweedie: Training Exact Ambient Diffusion Models with Noisy Data

2024 | Giannis Daras, Alexandros G. Dimakis, Constantinos Daskalakis
The paper presents a novel framework for training diffusion models using noisy data, addressing the open problem of sampling from the uncorrupted distribution. The key technical contributions include: 1. **Double Application of Tweedie’s Formula**: A method to learn optimal denoisers for noise levels $\sigma \geq \sigma_n$, where $\sigma_n$ is the standard deviation of the noise in the training data. This is achieved by applying Tweedie’s formula twice, which allows for efficient computation. 2. **Consistency Loss Function**: A consistency loss function to learn optimal denoisers for noise levels $\sigma \leq \sigma_n$. This function ensures that the model's predictions are consistent with the true conditional expectation, improving sampling performance at noise levels below the observed data noise. The paper also highlights the issue of memorization in diffusion models, demonstrating that even highly corrupted images can be almost perfectly reconstructed, indicating significant memorization of the training set. The proposed framework is evaluated on the Stable Diffusion XL model, showing that fine-tuning with corrupted data reduces memorization while maintaining competitive performance. Additionally, the method is shown to be effective in generating samples from distributions with different noise levels and in handling datasets with significantly different distributions. The paper concludes by discussing limitations and future directions, including the need to address linearly corrupted data and the impact of training time due to the use of consistency loss.The paper presents a novel framework for training diffusion models using noisy data, addressing the open problem of sampling from the uncorrupted distribution. The key technical contributions include: 1. **Double Application of Tweedie’s Formula**: A method to learn optimal denoisers for noise levels $\sigma \geq \sigma_n$, where $\sigma_n$ is the standard deviation of the noise in the training data. This is achieved by applying Tweedie’s formula twice, which allows for efficient computation. 2. **Consistency Loss Function**: A consistency loss function to learn optimal denoisers for noise levels $\sigma \leq \sigma_n$. This function ensures that the model's predictions are consistent with the true conditional expectation, improving sampling performance at noise levels below the observed data noise. The paper also highlights the issue of memorization in diffusion models, demonstrating that even highly corrupted images can be almost perfectly reconstructed, indicating significant memorization of the training set. The proposed framework is evaluated on the Stable Diffusion XL model, showing that fine-tuning with corrupted data reduces memorization while maintaining competitive performance. Additionally, the method is shown to be effective in generating samples from distributions with different noise levels and in handling datasets with significantly different distributions. The paper concludes by discussing limitations and future directions, including the need to address linearly corrupted data and the impact of training time due to the use of consistency loss.
Reach us at info@study.space