Source-linked AI summary
TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillation
David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap, Shuangfei Zhai, Siyuan Hu, Daniel Zheng, Walter Talbott, Eric Gu
TL;DR
High-quality diffusion sampling can require thousands of denoising steps. TRACT distills teacher outputs across time intervals using a teacher update followed by a student update, and reports improved single-step FID without architecture changes.
Problem
Diffusion sampling can require up to thousands of denoising steps to achieve high quality, limiting its favorability in resource-constrained settings.
Source: §1 · p. 1Method
TRACT trains a student to approximate a teacher’s output from t to t′ with one teacher inference update followed by a student update, using one or two phases.
Source: §1 · p. 2Results
Single-step FID improves by up to 2.4× without architecture changes.
Source: §5 · p. 9Takeaways & Limitations
Single-step sampling can greatly improve the tractability of diffusion models.
Source: §5 · p. 9Distillation from teachers above 8192 steps remains unexplored; sampling costs have typically capped DDIM/DDPM distillation at 8192 steps.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Denoising Diffusion models have demonstrated their proficiency for generative sampling. However, generating good samples often requires many iterations. Consequently, techniques such as binary time-distillation (BTD) have been proposed to reduce the number of network calls for a fixed architecture. In this paper, we introduce TRAnsitive Closure Time-distillation (TRACT), a new method that extends BTD. For single step diffusion,TRACT improves FID by up to 2.4x on the same architecture, and achieves new single-step Denoising Diffusion Implicit Models (DDIM) state-of-the-art FID (7.4 for ImageNet64, 3.8 for CIFAR10). Finally we tease apart the method through extended ablations. The PyTorch implementation will be released soon.
1 Introduction
Diffusion models can generate high-quality samples, but inference often requires many denoising steps. TRACT extends time-distillation to target one-step inference and reports improved single-step results.
1 Introduction
- Diffusion inference can require thousands of denoising steps, making it less favorable than one-shot methods in resource-constrained settings.Source: §1 · p. 1
- Inference-speeding approaches reduce input dimensionality, improve ODE solvers [32], or distill teachers into students with fewer steps [34].Source: §1 · p. 1
- BTD divides distillation into log2(T) phases, each replacing two teacher steps, and can reach four steps with minor performance loss on CIFAR10 and 64x64 ImageNet.Source: §1 · p. 2
- BTD’s phase-by-phase process can accumulate approximation error and prevents aggressive stochastic weight averaging for generalization.Source: §1 · p. 2
- TRACT trains a student to approximate a teacher’s output from time t to t′ using one teacher update followed by a student update, and can finish in one or two phases.Source: §1 · p. 2
- FID scores of 7.4 on 64x64 ImageNet and 3.8 on CIFAR10 accompany TRACT’s reported improvement over state-of-the-art one- and two-step results.Source: §1 · p. 2
2 Related Work
This section outlines DDIM inference, numerical ODE solvers, and progressive diffusion-model distillation. It describes their roles and constraints in reducing sampling steps.
2 Related Work
- DDIMs reuse the original noise at every inference step; the paper describes signal prediction as one way to estimate the clean sample and then infer samples at other timesteps.Source: §2 · p. 2
- Probability-flow ODEs permit larger steps than SDEs and can use numerical solvers, but truncation error limits the accuracy of approximate trajectories.Source: §2 · p. 2
- Euler and Runge-Kutta are numerical ODE solvers; Heun’s second-order method balances truncation error against the number of function evaluations [23].Source: §2 · p. 3
- Existing ODE solvers struggle to produce high-quality samples in the few-step regime, while TRACT is presented as an orthogonal approach whose outputs can be further refined with higher-order methods.Source: §2 · p. 3
- Progressive distillation transfers a pretrained diffusion model’s outputs to a student with fewer steps, extending earlier single-step distillation approaches [33] [34].Source: §2 · p. 3
- In BTD, a student replaces two teacher denoising steps; repeated student-to-teacher phases halve sampling steps, requiring log2 T phases to reach one step.Source: §2 · p. 3
3 Method
TRACT extends binary time-distillation by reducing its number of phases and using self-teaching to approximate a transitive closure across timestep groups. The method is motivated by accumulated teacher error and the interaction between training length and weight averaging, and is also adapted to a variance-exploding setting with a Runge–Kutta teacher.
3 Method
- TRACT reduces binary time-distillation from log2 T phases to typically one or two; it is demonstrated with signal prediction but also works with noise-predicting objectives.The method focuses first on the variance-preserving setting and is described as independent of that setting.Source: §3 · p. 3
3.1 Motivation
- The authors conjecture that distilled sample quality depends on both the number of phases and the length of each phase, and say experiments later validate this conjecture.Source: §3.1 · p. 3
Objective degeneracy
- In binary time-distillation, each student has positive loss and becomes an imperfect teacher, so errors accumulate across successive generations and cause objective degeneracy.Source: Objective degeneracy · p. 3
Generalization
- EMA momentum is constrained by training length: high momentum can improve DDPM results but over-regularize models when training is too short, and total training length scales with phase count.Source: Generalization · p. 3
3.2 TRACT
- Each phase partitions T steps into T′ contiguous groups, and EMA self-teaching approximates each group’s closure with a teacher step followed by a student jump.The partitioning strategy is open; experiments use equally sized groups.Source: §3.2 · p. 4
3.3 Adapting TRACT to a Runge-Kutta teacher and Variance Exploding noise schedule
- For variance-exploding noise, TRACT adapts an EDM Runge–Kutta teacher to a DDIM student and uses self-teaching to learn the transitive closure.The VE schedule is parameterized by noise standard deviations, with σ_0 = 0.Source: §3.3 · pp. 4–5
4 Experiments
Experiments evaluate TRACT on CIFAR-10 and class-conditional 64x64 ImageNet, reporting single-step results and ablations of EMA settings and distillation schedules.
4.1 Image generation results with BTD teachers
- CIFAR-10 1-step FID was 5.02 versus BTD’s 9.12 at 96M samples (4.45 at 256M), while 64x64 ImageNet reached 7.43, 2.4x better than BTD.The CIFAR-10 results are in Table 1 and ImageNet results in Table 2.Source: §4.1 · p. 6, Table 1 · p. 6, Table 2 · p. 7
4.2 Image generation results with EDM teachers
- TRACT-EDM runs use released EDM checkpoints, based on NCSN++ for CIFAR-10 and ADM for 64x64 ImageNet; results are reported for both datasets.The results are presented in Tables 1, 2, 7, and 8.Source: §4.2 · p. 7
4.3 Stochastic Weight Averaging ablations
- Self-teaching EMA performs well for µS∈[0.1, 0.9], while values ≤0.01 are worse and high-variance and values above about 0.9 degrade performance; inference EMA also strongly affects results.The self-teaching ablation fixes the schedule at 1024→32→1, uses 48M samples per phase, and sets µI=0.99995.Source: §4.3 · p. 7
4.4 Influence of the number of distillation phases
- With fixed overall training length, TRACT performs best with two distillation phases, while results worsen as the number of phases increases.The fixed-overall-length ablation uses 96M samples, µS=0.5, and µI=0.99995.Source: §4.4 · p. 8, Table 5 · p. 8
- With 48M samples per phase, 3–5-phase TRACT remained worse than the 2-phase schedule; on a shared 10-phase schedule, BTD scored FID 5.95 versus TRACT’s 6.8.The authors suggest objective degeneracy is not fully solved by increasing per-phase training, and that BTD’s overall disadvantage may reflect its inability to use a 2-phase schedule.Source: §4.4 · pp. 8–9, Table 6 · p. 8
4.5 Beyond time distillation
- Distilling a 60.0M-parameter model to 19.4M parameters raises CIFAR-10 1-step FID from 5.02 to 6.47.This experiment explores using TRACT to distill into a smaller architecture.Source: §4.5 · p. 9
5 Conclusion
TRACT improves sample quality in few-step diffusion through fewer distillation phases and stronger stochastic weight averaging, and the authors identify application to other data types as future work.
5 Conclusion
- Up to 2.4×, TRACT improves single-step FID without architecture changes and can distill to smaller student architectures.The method is demonstrated on image datasets; applying it to other data types is left to future work.Source: §5 · p. 9
- TRACT distills with fewer phases and stronger stochastic weight averaging than prior methods.Source: §5 · p. 9
- The authors propose exploring teachers with more than 8192 steps, since TRACT permits arbitrary step reductions between training phases.They describe this as a potential extension to improve the quality-efficiency trade-off.Source: §5 · p. 9
Contributions
The contributions section credits the team with developing TRACT, running experiments and ablations, building training infrastructure, and preparing data, models, and evaluation pipelines.
Contributions
- David Berthelot led the research and developed the transitive-closure method and working code prototypes.Source: Contributions · p. 9
- Arnaud Autef obtained CIFAR-10 results and led ablations and distributed training setup, while Walter Talbott contributed to ablations and writing.Source: Contributions · p. 9
- Daniel Zheng built cloud compute infrastructure and ran experiments; Siyuan Hu implemented FID and conducted smaller-architecture distillation experiments.Source: Contributions · p. 9
- Jierui Lin finalized data and evaluation pipelines, obtained ImageNet64 results, and reproduced binary distillation variants; Dian Ang Yap implemented EDM variants and designed TRACT experiments.Source: Contributions · p. 9
- Shuangfei Zhai contributed discussions, writing, and ablations, while Eric Gu contributed writing and smaller-architecture distillation experiments.Source: Contributions · p. 9
A.1 Deriving the distillation target ˆx
The derivation expresses the student’s target as a closure of teacher steps, with the student matching a teacher update followed by a student update.
A.1 Deriving the distillation target ˆx
- The student is trained to match the closure of teacher steps through self-teaching.Source: §A.1 · p. 13
- The displayed equality relates the student’s direct mapping from t to t_i to a teacher update from t to t−1 followed by a student mapping to t_i.Source: §A.1 · p. 13
- The right-hand expression is denoted x_ti and used to develop equations for the training target of a perfect student.Source: §A.1 · p. 13
A.2 Deriving the distillation loss
The section derives a signal training loss from a noise training loss and notes a commonly used implementation form for numerical reasons.
A.2 Deriving the distillation loss
- The signal training loss is derived from the noise training loss using expected noise and predicted noise.Source: §A.2 · p. 13
- The derivation frames the relationship as one between expected noise ϵ and predicted noise ϵ|t.Source: §A.2 · p. 13
- The implementation commonly uses an alternative form for numerical reasons.Source: §A.2 · p. 13
A.3 Experimental details for TRACT
The experiments use separate CIFAR-10 and ImageNet training configurations, with distinct inference-time EMA settings and phase allocations. Table 7 reports CIFAR-10 FID results with markers distinguishing architecture/data and teacher differences.
A.3 Experimental details for TRACT
- CIFAR-10 uses self-teaching EMA momentum µS = 0.5 and inference EMA momentum 0.99997 for TRACT-96M or 0.99999 for TRACT-256M, with 1/6 of training samples for 1024 →32 and 5/6 for 32 →1.The authors note the optimal phase balance may depend on total training length.Source: §A.3 · p. 13
- Table 7 reports CIFAR-10 FID and marks models using a different student architecture and synthetic training data with †, and TRACT-EDM models using better teachers with ∗.The † marker applies to Diffusion Fourier Neural Operators.Source: Table 7 · p. 14
- For 64x64 ImageNet, best-performing TRACT-96M uses global batch size 256 across 8 GPUs, Adam at constant learning rate 2 × 10−4, no weight decay or dropout, and gradient clipping to norm 1.0.Source: §A.3 · p. 14
- ImageNet uses self-teaching EMA momentum µS = 0.5, inference EMA momentum 0.99995, equal training budgets for 1024 →32 and 32 →1, and predicts both signal x0 and noise ϵ0.The signal-and-noise prediction follows the BTD setup.Source: §A.3 · p. 14
A.4 Experimental details for TRACT with EDM teachers
The appendices detail TRACT configurations, additional sampling results, an inference-EMA heuristic, EDM-teacher distillation, and architecture transfer. The reported 4- and 8-step results are slightly worse than BTD, while the EMA heuristic performs better than a direct grid search in the tested settings.
A.4 Experimental details for TRACT with EDM teachers
- CIFAR-10 EDM-teacher experiments use NCSN++ with 56 million parameters, batch size 512, and Adam; the 40-step teacher setting, with 79 Runge-Kutta NFEs, is distilled to one step.The chosen timestep setting balances teacher-target quality with student learning ease.Source: §A.4 · p. 14
- The 64x64 ImageNet EDM-teacher setup uses a 296-million-parameter ADM model and one-phase distillation from 128 timesteps, or 255 Runge-Kutta NFEs, to one or multiple steps.The model is used with weights from Karras et al. and no changes.Source: §A.4 · p. 14
A.5 More experimental results
- 4- and 8-step student performance on CIFAR-10 and 64x64 ImageNet is slightly worse than BTD, possibly because self-teaching is less efficient than supervised training.Tables 7 and 8 report these results.Source: §A.5 · p. 14, Table 8 · p. 15
A.6 Generated samples
- Across varying-step sample figures, the deterministic noise-to-sample mapping is mostly preserved, though one-step students show slightly degraded image quality.Figures 2–4 show fixed-noise samples across sampling-step counts on CIFAR-10 and 64x64 ImageNet.Source: Figure 2 · p. 15, Figure 3 · p. 15, Figure 4 · p. 15
A.7 A heuristic to pick the EMA momentum parameter
- Keeping ε fixed as training length changes provides an inference-EMA momentum heuristic that performs better than direct grid search in the tested CIFAR-10 runs.The experiments use ε values from 10−1 to 10−5 and training lengths from 16M to 64M samples; all ε choices reach reasonable performance, and longer training always improves performance.Source: §A.7 · p. 16, Figure 5 · p. 17
A.8 TRACT with EDM teachers
- For EDM teachers, TRACT keeps the VE noise schedule, uses a Runge-Kutta teacher sampler, and distills into a DDIM student by self-teaching the closure of teacher steps.The training target applies one teacher RK update followed by a student update toward the target time.Source: §A.8 · pp. 16–17
A.9 Knowledge Distillation
- TRACT can transfer knowledge from architecture A to architecture B with one additional training phase.The authors identify heterogeneous student and teacher architectures, such as Transformers, as future validation work.Source: §A.9 · pp. 17–18