Enhancing diffusion models with Gaussianization preprocessing

TL;DR

This work tackles slow diffusion-model sampling caused by a bifurcation in reverse-time dynamics by introducing an invertible Gaussianization preprocessing that reshapes data toward an independent standard Gaussian. The method combines Independent Component Analysis with one-dimensional Gaussianization (via KDE, CDF, and the inverse Gaussian CDF) and supports an inverse Gaussianization to recover original data; iterative variants further suppress residual non-Gaussian structure. In experiments with synthetic Gaussian Mixture data, Gaussianization leads to faster inference convergence and modest training-speed gains while preserving or improving alignment with the true distribution, especially for small networks. The findings underscore data-space transformations as a practical complement to trajectory design in diffusion models, while acknowledging scalability limits and proposing avenues for deeper Gaussianization and conditional extensions.

Abstract

Diffusion models are a class of generative models that have demonstrated remarkable success in tasks such as image generation. However, one of the bottlenecks of these models is slow sampling due to the delay before the onset of trajectory bifurcation, at which point substantial reconstruction begins. This issue degrades generation quality, especially in the early stages. Our primary objective is to mitigate bifurcation-related issues by preprocessing the training data to enhance reconstruction quality, particularly for small-scale network architectures. Specifically, we propose applying Gaussianization preprocessing to the training data to make the target distribution more closely resemble an independent Gaussian distribution, which serves as the initial density of the reconstruction process. This preprocessing step simplifies the model's task of learning the target distribution, thereby improving generation quality even in the early stages of reconstruction with small networks. The proposed method is, in principle, applicable to a broad range of generative tasks, enabling more stable and efficient sampling processes.

Figures (9)

• Figure 1: The data transformation process. From left to right: Original data distribution, ICA transformed data distribution, Gaussianized data distribution, and reconstructed data distribution. Each step visualizes the evolution of the data through the proposed pipeline.
• Figure 2: Illustration of the one-dimensional Gaussianization process. Top Left: Original bimodal distribution. Bottom Left: $u = F(x)$. Bottom Right: $y = G^{-1}(u)$. Top Right: Gaussianized distribution.
• Figure 3: Iterative Gaussianization Process. The figure demonstrates the transformation of a two-dimensional dataset through the Gaussianization pipeline across multiple iterations. Each row corresponds to a specific iteration, showing the original data distribution $x^{(0)}$, the Gaussianized data $\mathcal{Z}^{(k)}$, the partially reconstructed data $x^{(k)} = \mathcal{Z}^{(k-1)}$, and the corresponding probability density functions (PDFs) $\hat{p}^{(k)}_i(z)$ for both dimensions. The iterative process ensures the reconstructed data aligns closely with the original data, while the Gaussianized data approaches a standard Gaussian distribution.
• Figure 4: Snapshots of reconstruction steps by the diffusion model using network width 16. Orange dots represent training data sampled from the Gaussian mixture model. The blue dots are reconstruction by the diffusion model at different time points. (a) Baseline, (b) Gaussianized.
• Figure 5: Snapshots of reconstruction steps by the diffusion model using the network width 32. (a) Baseline, (b) Gaussianized. Presentation follows Fig. \ref{['fig:snapshots_16']}