GAS: Improving Discretization of Diffusion ODEs via Generalized Adversarial Solver

AI summary

Overview

Research area: Generative modeling / computer vision — specifically, accelerating the sampling process of diffusion models by training lightweight ODE solvers.

Technical level: Advanced. The paper assumes familiarity with probability flow ODEs, numerical ODE solvers (Euler, Runge-Kutta, linear multistep methods), knowledge distillation, and GAN training.

Scope: The paper introduces a new parameterization of a diffusion ODE sampler ("Generalized Solver") plus an adversarial training extension ("Generalized Adversarial Solver"), and evaluates them against existing solver-optimization and timestep-schedule-training methods on pixel-space and latent-space image benchmarks.

What This Paper Is About

Diffusion models produce state-of-the-art images but are slow to sample, requiring dozens of function evaluations. A recent line of work trains a small "student" solver to imitate the full "teacher" sampling trajectory using only a few steps. This paper argues that existing solver-distillation methods rely on awkward training tricks (unstable loss scaling, limited parameter spaces, separately trained parameter subsets) and lose fine-grained detail. The authors propose a simpler, more expressive solver parameterization and pair it with an adversarial loss to improve both convergence and visual fidelity.

Key Contributions

1. The Generalized Solver (GS): a new parameterization of the ODE sampler in which each step is a weighted sum of all previous points and velocity directions — removing the fixed solver-order restriction — with trainable additive corrections on top of theoretically derived coefficients from a base solver (e.g., DPM-Solver++(3M)).
2. Theoretical guidance as a backbone: rather than learning coefficients from scratch, the method trains additive corrections to time-dependent theoretical coefficients, initialized at zero, so that changing the timestep schedule does not destroy solver performance.
3. Adversarial training added to distillation: the distillation loss is combined with a relativistic GAN loss (R3GAN-style, with discriminator gradient penalties) to reduce artifacts and improve detail in low-NFE regimes.
4. Empirical demonstration: GS and GAS outperform existing solver/timestep-training methods across CIFAR10, FFHQ, AFHQv2, LSUN-Bedroom, ImageNet-256, and MS-COCO with Stable Diffusion, under similar resource constraints, plus ablations on parameterization, dataset size, and training stability.

Main Findings

• Consistent gains across datasets: GAS improves FID over prior solver-optimization methods at every evaluated NFE. On FFHQ at NFE=4, S4S Alt reports 10.63 while GAS reaches 7.86. On AFHQv2 at NFE=4, GAS reaches 4.48; on FFHQ at NFE=6, 3.79; on conditional ImageNet-256 at NFE=4, 5.38; on LSUN-Bedroom at NFE=5, 4.60; on MS-COCO with Stable Diffusion v1.5 at NFE=4, 14.71.

• Pixel-space CIFAR10 results: GAS achieves 4.05 at NFE=4, 2.49 at NFE=6, 2.24 at NFE=8, and 2.17 at NFE=10, versus a teacher FID of 2.03 (teacher model FIDs reported are those used during training).

• Parameterization matters for training stability: in a controlled comparison with equal trainable parameters (batch size 24, training dataset of 49k samples), the authors' parameterization reaches FID 4.39 on CIFAR10 and 10.79 on FFHQ at NFE=4, versus 31.44 and 24.24 for the S4S variant. Figure 3 reports more stable LPIPS training curves and faster convergence for the proposed parametrization.

• Adversarial loss improves realism even when regression loss worsens: in Table 4 (10k training iterations, 1000 validation samples), FFHQ GS scores FID 10.70 with distillation loss 0.116, while GAS scores FID 7.86 with a higher distillation loss of 0.127. On LSUN, GS scores 9.64 / 0.172 and GAS 7.54 / 0.174. The authors attribute this to the student being unable to reproduce teacher images exactly in low-NFE settings, which causes artifacts that the adversarial term suppresses.

• Efficient data usage: the authors report that 1400 teacher samples suffice to train GS without the adversarial loss, but that the optimization becomes harder in low-NFE scenarios with the adversarial loss, requiring expansion from 1400 to 5000 samples (the "full" scenario uses 49000 samples).

• Distillation distances used: LPIPS in pixel-space and L1 in latent-space experiments; the "soft" variant of the distillation objective is not used.

Methodology in Plain English

The authors start from the observation that most diffusion solvers — linear multistep methods and their diffusion-specific variants — all share the same basic update form: a coefficient times the current point, plus a weighted sum of recent velocity evaluations.

They change this in three ways. First, they let each new step depend on all previous points and all previous velocity evaluations, not just the last K. Second, instead of learning solver coefficients freely, they take the coefficients of a strong existing solver (DPM-Solver++(3M)) as a mathematical starting point and learn small additive corrections to them, initialized at zero. This keeps training stable even when the timestep schedule shifts. Third, they parameterize the timestep schedule itself via a "stick-breaking" transformation of trainable logits, plus small per-timestep offsets used when evaluating the diffusion model.

Training is done by distillation: a high-quality "teacher" sampler produces target outputs, and the student is trained to match them. On top of this regression objective, the authors add a GAN-style adversarial loss with a discriminator, using a relativistic formulation and gradient penalties, so the student's outputs look like the teacher's distribution rather than merely being close in pixel or latent distance.

Why This Matters

Impact on research: Automatic sampler learning is a cheap alternative to full diffusion distillation, and this paper shows that much of the remaining gap comes from how the solver is parameterized, not just how it is trained. The additive-guidance idea — learning corrections on top of theoretical coefficients rather than replacing them — is a general recipe that could be applied to other numerical scheme learning problems. The finding that lower regression loss does not imply lower FID (and can correspond to visible artifacts) is a useful caution for the distillation literature.

Real-world applications:

• Faster text-to-image generation, demonstrated here with Stable Diffusion v1.5 on MS-COCO at 512×512.
• On-device or low-compute image generation, where NFE savings translate directly to latency and cost.
• Class-conditional and unconditional image synthesis pipelines (ImageNet, LSUN-Bedroom, FFHQ, AFHQv2) used for data augmentation and content creation.
• Domains the introduction names as resource-constrained, such as text-to-video and text-to-3D generation, where full distillation is described as computation- and memory-heavy.

Industry relevance: the method is presented as inference-time acceleration that works on top of an existing pretrained diffusion model rather than requiring a new model to be trained. That makes it attractive for deployments that need to squeeze more quality out of a fixed number of network evaluations without retraining the backbone.

Future Directions

• Whether the reported dataset-size findings extrapolate: the available content is truncated in Section 4.3, so the full conclusion about how few teacher samples suffice with adversarial training is not reported here.
• Extension to video and 3D generation: the introduction identifies these as the settings where full distillation becomes infeasible, but no video or 3D experiments are reported in the available content.
• Interaction with other inference-time acceleration methods: caching intermediate steps and quantization are named in the introduction as alternative approaches, but the paper's available content does not report combining them with GS or GAS.
• How much of the gain comes from adversarial training versus parameterization at very low NFE: the ablation table shows adversarial training is decisive at NFE=4, while the gap narrows at higher NFE (for example, LSUN-Bedroom GS 3.34 versus GAS 3.36 at NFE=7). Whether similar dynamics hold for other model families is not reported in the available content.

Target Audience

Researchers and engineers working on diffusion model inference acceleration, few-step sampling, and solver distillation. It is most valuable to readers already comfortable with probability flow ODEs and numerical integration, and to practitioners who want a drop-in improvement over training-free solvers such as UniPC and DPM-Solver++ without retraining the underlying diffusion model.

Authors’ abstract

While diffusion models achieve state-of-the-art generation quality, they still suffer from computationally expensive sampling. Recent works address this issue with gradient-based optimization methods that distill a few-step ODE diffusion solver from the full sampling process, reducing the number of function evaluations from dozens to just a few. However, these approaches often rely on intricate training techniques and do not explicitly focus on preserving fine-grained details. In this paper, we introduce the Generalized Solver: a simple parameterization of the ODE sampler that does not require additional training tricks and improves quality over existing approaches. We further combine the original distillation loss with adversarial training, which mitigates artifacts and enhances detail fidelity. We call the resulting method the Generalized Adversarial Solver and demonstrate its superior performance compared to existing solver training methods under similar resource constraints. Code is available at https://github.com/3145tttt/GAS.