U-Turn Diffusion

TL;DR

This paper introduces U-Turn diffusion, a method that shortens diffusion processes to improve sample generation, revealing critical times where samples diverge or change class representation, and analyzing score function non-linearity effects.

Contribution

We propose U-Turn diffusion, a novel augmentation that shortens diffusion processes and maintains detailed balance, providing insights into memorization and class separation times.

Findings

Identified critical Memorization Time T_m where samples diverge.

Discovered Speciation Time T_s where samples begin representing different classes.

Showed the score function becomes affine beyond T_s and approximately affine between T_m and T_s.

Abstract

We investigate diffusion models generating synthetic samples from the probability distribution represented by the Ground Truth (GT) samples. We focus on how GT sample information is encoded in the Score Function (SF), computed (not simulated) from the Wiener-Ito (WI) linear forward process in the artifical time $t\in [0\to \infty]$, and then used as a nonlinear drift in the simulated WI reverse process with $t\in [\infty\to 0]$. We propose U-Turn diffusion, an augmentation of a pre-trained diffusion model, which shortens the forward and reverse processes to $t\in [0\to T_u]$ and $t\in [T_u\to 0]$. The U-Turn reverse process is initialized at $T_u$ with a sample from the probability distribution of the forward process (initialized at $t=0$ with a GT sample) ensuring a detailed balance relation between the shorten forward and reverse processes. Our experiments on the class-conditioned SF of the ImageNet dataset and the multi-class, single SF of the CIFAR-10 dataset reveal a critical Memorization Time $ T_m $, beyond which generated samples diverge from the GT sample used to initialize the U-Turn scheme, and a Speciation Time $ T_s $, where for $ T_u > T_s > T_m $, samples begin representing different classes. We further examine the role of SF non-linearity through a Gaussian Test, comparing empirical and Gaussian-approximated U-Turn auto-correlation functions, and showing that the SF becomes effectively affine for $ t > T_s $, and approximately affine for $t\in [T_m,T_s]$.