Source-linked AI summary
GAS: Improving Discretization of Diffusion ODEs via Generalized Adversarial Solver
Aleksandr Oganov, Ilya Bykov, Eva Neudachina, Mishan Aliev, Alexander Tolmachev, Alexander Sidorov, Aleksandr Zuev, Andrey Okhotin, Denis Rakitin, Aibek Alanov
TL;DR
Few-step diffusion-solver training can reduce sampling cost, but existing approaches can be resource-intensive or complex and may not preserve fine details. GAS introduces a generalized sampler parameterization and combines distillation with adversarial training. The paper reports superior performance to existing solver-training methods across pixel- and latent-space datasets.
Problem
Few-step solver training can reduce inference time, but its compute and memory demands, complicated training schemes, and fine-detail preservation concerns remain challenges.
Source: §1 · pp. 1–2Method
GAS learns a generalized sampler that uses predefined-solver guidance for its weights and combines distillation with adversarial training.
Source: §1 · p. 2Results
GAS outperforms existing solver/timestep-training methods across several pixel- and latent-space datasets.
Source: §1 · p. 2Takeaways & Limitations
Adversarial loss is reported to boost performance and address artifacts from simple solver distillation.
Source: §5 · p. 10Backpropagating through the whole solver inference may limit scalability to larger images and models, and separate training for each preferred inference NFE remains a concern.
Source: §5 · p. 10
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Diffusion sampling can be costly, while solver and sampler-training approaches face resource, tuning, training-complexity, or detail-preservation challenges. The paper introduces a generalized sampler parameterization and adversarial training, reporting better results than existing solver-training methods.
1 INTRODUCTION
- Few-step training methods can reduce inference time but may be compute- and memory-heavy, while solver methods require substantial model-specific hyperparameter search.Source: §1 · pp. 1–2
- Existing sampler-training approaches can involve complicated training schemes, and straightforward distillation may inadequately preserve fine-grained detail.Source: §1 · p. 2
- The Generalized Solver uses a weighted sum of current and previous-step points and directions, with a predefined solver guiding learned corrections to theoretical weights.Source: §1 · p. 2
- Adding adversarial training to distillation is reported to improve fine-grained generation details.Source: §1 · p. 2
- GAS outperforms existing solver/timestep-training methods across pixel- and latent-space datasets; at equal NFE, it is reported superior to UniPC and nearly teacher-like.Source: §1 · p. 2, Figure 1 · p. 2
2 BACKGROUND
The background covers diffusion generation through a probability-flow ODE, numerical ODE solvers, optimization of student samplers against teacher outputs, and adversarial training.
2.1 DIFFUSION MODELS
- Diffusion sampling reverses a noising process by numerically solving the probability-flow ODE using an estimated score or velocity field.Source: §2.1 · pp. 2–3
2.2 ODE SOLVERS
- Linear multistep solvers reuse previously computed points and directions, while Runge–Kutta methods use midpoint evaluations that can hurt low-NFE performance.Source: §2.2 · p. 3
2.3 SOLVER AND SCHEDULE DISTILLATION
- Solver and schedule distillation optimize a lightweight student sampler’s parameters to match outputs from a high-quality teacher sampler.Source: §2.3 · pp. 3–4
2.4 ADVERSARIAL TRAINING
- Adversarial training uses a discriminator to distinguish real from generated samples while the generator aims to fool it; relativistic loss is designed to discourage mode dropping.Source: §2.4 · p. 4
3 METHOD
The Generalized Solver learns a flexible multistep sampler, combining its timestep schedule and coefficients with guidance from a theoretical solver. The Generalized Adversarial Solver adds adversarial training to distillation, which the paper reports helps remove artifacts at low NFE.
3.1 GENERALIZED SOLVER (GS)
- Each Generalized Solver step weights the current velocity direction and all previous points, extending the solver signature beyond a restricted order.The design also includes previously computed velocity directions.Source: Figure 2 · p. 5, §3.1 · p. 5
- GS learns a timestep schedule through stick-breaking portions and adds corrections to the timesteps used for model predictions.The coefficient corrections are initialized to zero, preserving meaningful time dependence even when timesteps change suddenly.Source: §3.1 · pp. 5–6
- A pre-defined solver supplies theoretical coefficients, which GS augments with trainable corrections and scalars for velocity coefficients.The theoretical coefficients form the basis for the learned coefficients.Source: §3.1 · p. 5
- GS combines theoretical coefficients, linear-multistep history, and additive parameterization to connect theoretical and trainable solver coefficients.The paper identifies theoretical coefficients as improving convergence.Source: §3.1 · p. 6
3.2 GENERALIZED ADVERSARIAL SOLVER (GAS)
- GAS augments GS distillation with adversarial training, using LPIPS for pixel-space experiments and L1 for latent-space experiments.The paper frames solver distillation as learning a mapping from paired input/output samples.Source: §3.2 · p. 7
- The paper reports that adversarial loss is effective at removing generation artifacts in low-NFE regimes.It notes that the regression task becomes harder in these regimes.Source: §3.2 · p. 7
4 EXPERIMENTS
Experiments evaluate the methods across pixel-space and latent-space image-generation tasks. GAS outperforms previously proposed methods across the evaluated datasets, with a reported FFHQ advantage over S4S Alt at four NFEs.
4 EXPERIMENTS
- Evaluations cover pixel-space CIFAR10, FFHQ, and AFHQv2, plus latent-space LSUN Bedroom, class-conditional ImageNet, and Stable Diffusion on MSCOCO.The pixel-space evaluations use 32×32 or 64×64 images, while the latent-space evaluations range from 256×256 to 512×512.Source: §4 · p. 7
- FID is calculated from 50,000 samples unless stated otherwise, with LPIPS used for pixel-space distillation and L1 for latent-space distillation.The experiments initialize timesteps with a time-uniform schedule and use DPM-Solver++(3M) coefficients as theoretical guidance.Source: §4 · p. 7
4.1 MAIN RESULTS
- FID 7.86 for GAS versus 10.63 for S4S Alt on FFHQ at NFE=4; GAS outperforms all previously proposed methods across evaluated datasets.The authors report that GS and GAS systematically improve sampling quality across solvers, especially at low NFE.Source: §4.1 · p. 7, Table 2 · p. 8
AFHQv2
The ablations report stronger solver performance and steadier optimization, while adversarial training improves image quality in low-NFE settings. Efficiency tests find that smaller training datasets can match full-data results, with adversarial training taking longer.
4.2 ABLATION STUDY
- The Generalized Solver outperforms S4S across datasets and NFE; Table 3 compares CIFAR10 and FFHQ using FID and LPIPS.The comparison uses 49k training samples and batch size 24.Source: §4.2 · p. 9, Table 3 · p. 9
- GS and GAS training improves FID over time, and the parameterization shows faster convergence and more stable LPIPS training than S4S.Figure 3 compares LPIPS evaluation loss across training iterations.Source: §4.2 · p. 9, Figure 3 · p. 9
- Adversarial loss improves generation quality and reduces artifacts at low NFE, even when LPIPS or L1 losses are larger.The authors report that lower regression loss does not always correspond to lower FID; adversarial training moves student generations closer to the teacher distribution.Source: §4.2 · p. 9, Figure 4 · p. 9
4.3 METHOD EFFICIENCY
- Training GS without adversarial loss needs 1,400 teacher samples, while GAS needs 5,000 to match full-data results across datasets and settings.The full-dataset scenario uses 49,000 samples; the smaller GAS dataset is used for the more challenging low-NFE settings with adversarial loss.Source: §4.3 · p. 10
- GS converges in 1–2.5 hours, comparable to LD3 and S4S, while GAS takes 2–9 hours, still within a similar order of training time.The reported durations vary by dataset.Source: §4.3 · p. 10
5 DISCUSSION
The discussion reports that GAS outperforms solver and timestep training approaches across six datasets, while identifying scalability and low-NFE constraints. Its appendices report benefits from theoretical coefficient guidance and adversarial loss choices.
5 DISCUSSION
- GAS combines theoretical guidance for solver coefficients with distillation and adversarial training, and outperforms solver/timestep training approaches on six datasets.The evaluation spans pixel- and latent-space datasets from 32 × 32 CIFAR10 to 512 × 512 MS-COCO with Stable Diffusion.Source: §5 · p. 10
- Backpropagation through the whole solver inference may limit scalability to larger images and models; whether separate training is needed for each NFE remains an open concern.The authors leave lightweight modifications to future work.Source: §5 · p. 10
- With pretrained diffusion weights frozen, GS and GAS better preserve the original model’s generative capabilities than distillation-based methods, but have lower quality at NFE=1 and 2.The paper attributes the low-NFE result to the limited number of trainable parameters in those settings.Source: §5 · p. 10
A RELATED WORK
- Prior acceleration work includes solver selection, timestep or coefficient optimization, caching, quantization, pruning, and few-step distillation.The related-work discussion describes both trajectory-matching and distribution-based distillation approaches, some of which use adversarial training.Source: A RELATED WORK · p. 16
B.1 THEORETICAL GUIDANCE FOR THE COEFFICIENTS
- On FFHQ, adding theoretical guidance to GS yields a substantial FID improvement over the version without guidance.The ablation uses identical hyperparameters and coefficients initialized to DPM-Solver++(3M).Source: §B.1 · p. 16
B.2 ADVERSARIAL LOSS
- On ImageNet at NFE=4, both traditional and relativistic GAN losses improve final output quality, with the relativistic loss producing a superior model.The comparison uses discriminator gradient penalties and matched hyperparameters.Source: §B.2 · pp. 16–17
B.3 FID PROGRESSION DURING TRAINING
Training curves show that GAS takes longer to converge than GS on FFHQ, while reaching a lower FID. At low NFE, GAS's FID continues to decrease through 30k iterations, whereas higher-NFE runs converge earlier.
B.3 FID PROGRESSION DURING TRAINING
- GAS requires more training iterations than GS to converge on FFHQ, but reaches a significantly lower FID.The authors describe this as a better trade-off between generation quality and a slight increase in training time.Source: §B.3 · p. 17, Figure 5 · p. 17
- At NFE 4 and 5, GAS's FID continues to decrease after 30k training iterations, suggesting further training could improve results.Source: §B.3 · p. 17, Figure 5 · p. 17
- Higher-NFE GAS runs do not show the same continued FID decline because their greater student capacity makes optimization easier and convergence earlier.Source: §B.3 · p. 17, Figure 5 · p. 17
B.4 GENERALIZATION ACROSS DATASETS
The appendix tests transfer between FFHQ and AFHQv2, adversarial-loss weighting and mode coverage, then compares solver properties and low-NFE performance. Results report consistent baseline improvements in several settings, with weak absolute FID at NFE 1 and 2.
B.4 GENERALIZATION ACROSS DATASETS
- FFHQ-to-AFHQv2 and AFHQv2-to-FFHQ transfer both yield GS' and GAS' results that consistently outperform baseline methods.Source: Table 7 · p. 18, §B.4 · p. 18, §B.5 · p. 18
B.5 ADVERSARIAL LOSS WEIGHT
- GAS is reported to achieve similar FID across adversarial-loss weights, indicating insensitivity to GAN-weight selection without hyperparameter tuning.The adversarial-loss weight was fixed to 1.0 for all datasets from 64 × 64 resolution.Source: Figure 6 · p. 18, §B.5 · p. 18
B.7 COMPARISON WITH CONSISTENCY DISTILLATION
- GS and GAS trained on CIFAR10 have higher coverage than Consistency Distillation, whose coverage is low relative to the teacher solver.The authors interpret CD's lower coverage as deterioration in generative properties; GS and GAS are described as more flexible and property-preserving acceleration methods.Source: Table 9 · p. 18, §B.7 · p. 19
B.6 MODE COLLAPSE
- On ImageNet at NFE 4, GAS does not suffer mode collapse and increases mode coverage and recall compared with GS.The authors chose this setting because adversarial loss was significant and could raise concerns about mode collapse.Source: Table 8 · p. 18, §B.6 · p. 18
B.8 PERFORMANCE AT LOW NFE
- On FFHQ at NFE 1 and 2, GS and GAS improve on DPM-Solver but still have high FID scores.Source: Table 11 · p. 19, §B.8 · p. 19
B.9 COEFFICIENTS PARAMETRIZATION
- On FFHQ, GS achieves better generation quality than LMS and LMS+PC at every evaluated NFE under identical training settings.The methods were initialized so GS and LMS were initially equivalent to DPM-Solver++(3M).Source: Table 10 · p. 19, §B.9 · p. 19
C.1 TRAINING DATASET SIZE
The dataset-size experiments find that GS can converge with far fewer teacher-generated training samples than the 49,000-sample baseline. GAS needs more data in some settings, where simultaneous discriminator training can overfit.
C.1 TRAINING DATASET SIZE
- The simultaneously trained GAS discriminator can overfit, requiring a larger dataset to alleviate the problem.Source: §C.1 · p. 20
- For GAS on higher-dimensional FFHQ, 5,000 training samples suffice to match the performance of training on 49,000.The passage identifies adversarial training as especially challenging in this setting.Source: §C.1 · pp. 19–20
- GS reaches equivalent or better FID than the larger-data baseline using 1,400 training images and 10,000 iterations, regardless of NFE.The comparison includes a 49,000-sample baseline.Source: §C.1 · p. 19, Table 12 · p. 20
C.2 TRAINING TIME
GS and GAS training takes hours or less in the reported comparisons, with GAS adding some training memory but not changing inference speed. GS uses the same peak training memory as LD3, and GAS adds less than 4 GB for its discriminator loss.
C.2 TRAINING TIME
- On CIFAR10, GS trains in under an hour and GAS in up to 2 hours; in latent diffusion, GS and GAS take up to 3 hours versus LD3's under an hour.The authors describe the latent-diffusion training times as within the same order.Source: §C.2 · p. 20, Table 14 · p. 21
C.3 MEMORY USAGE
- GS uses the same peak training memory as LD3, while GAS adds less than 4 GB for discriminator-loss training.The paper characterizes the additional memory as minor; storing prior states does not add peak-memory overhead.Source: §C.3 · p. 20, Table 15 · p. 21
C.4 INFERENCE TIME
- GS inference time is comparable to solver baselines, and GAS and GS have identical inference times with no added inference-time overhead across the reported model and dataset settings.The extra operations use prior states and are described as much simpler computationally than a diffusion-model function evaluation.Source: Table 16 · p. 22, §C.4 · p. 22
D EXPERIMENTAL DETAILS
Appendix D documents the baseline schedules, teacher setup, solver parameterization, and training protocols for pixel, latent, and Stable Diffusion experiments. The timestep visualization reports that learned decoupled schedules remain positive and monotonically decreasing without explicit constraints.
D.1 BASELINE DISCRETIZATION HEURISTICS
- Baseline heuristics include polynomial time schedules, logSNR schedules, and GITS optimized noise-level sequences for low-NFE sampling.The polynomial exponent is commonly 1 or 2, corresponding to time-uniform or time-quadratic discretization.Source: §D.1 · p. 22
D.2 TEACHER SOLVER
- Teacher data use UniPC, while Stable Diffusion students are trained against a one-plus teacher to reduce the gap in noise dynamics and support smoother convergence.The one-plus setup trains an NFE=n student with an NFE=n+1 teacher.Source: §D.2 · pp. 22–23
D.3 SOLVER COEFFICIENTS PARAMETERIZATION
- With all parameters ϕ set to zero, the Generalized Solver reduces exactly to DPM-Solver++(3M).The appendix provides the solver step in Algorithm 1.Source: §D.3 · p. 23
D.4 PRACTICAL IMPLEMENTATION DETAILS
- The appendix assigns θ, ϕ, and ξ to schedule, solver-coefficient, and evaluation-time timestep parameters, and uses FID evaluation on 50,000 samples (30,000 for MS-COCO).ImageNet evaluation uses an equal number of generated samples per class.Source: §D.4 · pp. 23–24
D.4.1 PIXEL SPACE DIFFUSION ON CIFAR10, FFHQ, AND AFHQV2
- Pixel-space setups use LPIPS distillation and adversarial-loss weights of 1.0 for FFHQ and AFHQv2; GS/GAS use 10k iterations, while GAS is listed as 5000/1000 for those datasets.A pre-trained FFHQ-64 checkpoint is used for both FFHQ and AFHQv2.Source: §D.4.1 · pp. 24–25
D.4.2 LATENT SPACE DIFFUSION ON LSUN-BEDROOM AND IMAGENET
- Latent-space experiments use L1 distillation, UniPC teachers, and 30,000 solver-training iterations; GAS uses 5,000 training examples versus 1,400 for GS.The experiments cover LSUN-Bedroom and ImageNet.Source: §D.4.2 · p. 25
D.4.3 TEXT-TO-IMAGE GENERATION WITH STABLE DIFFUSION
- The experiments use the IPNDM(2M) solver with GITS, and GAS is trained for 2,000 iterations on 5,000 training examples with 128 validation examples.The teacher uses NFE = n + 1, where n is the student NFE.Source: §D.4.3 · pp. 25–26
E TIMESTEPS VISUALIZATION
- Learned decoupled timesteps remain positive and monotonically decreasing without explicit constraints, and are typically slightly lower than their corresponding unconstrained timesteps.One first timestep slightly exceeded 1; the authors report that continuous positional embeddings made the model robust to this small shift.Source: E TIMESTEPS VISUALIZATION · p. 26
F ADDITIONAL SAMPLES
Additional samples show that adversarial training’s visual effect is most prominent at low NFE: many selected GAS samples look better than GS at NFE 4, while same-seed outputs become closer to the teacher from NFE 6. Random samples show smaller differences, and the authors report no mode collapse.
F ADDITIONAL SAMPLES
- Random samples show only minor fine-grained differences between GS and GAS, which the authors associate with a positive effect on FID.The cited passage refers to Table 2 for the FID effect; no table values are supplied here.Source: F ADDITIONAL SAMPLES · p. 26
- Figure 7 visualizes final timestep schedules with and without additive corrections on FFHQ and ImageNet.The additional comparison figures show GS and GAS samples against the teacher and UniPC across FFHQ, AFHQv2, LSUN-Bedroom, and ImageNet.Source: F ADDITIONAL SAMPLES · p. 26, Figure 7 · p. 27
- The authors report no mode collapse with adversarial training, citing random-sample diversity and low resulting FID values.They attribute their explicit response to this concern to the relativistic GAN loss.Source: F ADDITIONAL SAMPLES · p. 27
- At NFE = 4, most selected GAS samples show superior image quality to GS but are farther from the teacher; same-seed outputs become closer to the teacher starting at NFE = 6.The authors say the adversarial effect is most prominent at low NFE, where matching teacher performance is harder.Source: F ADDITIONAL SAMPLES · p. 27, Figure 8 · p. 29, Figure 9 · p. 29
- Stable Diffusion examples compare GAS with LD3 using identical initial latent noise, showing six random and six selected images.Figures 24 and 25 cover MS-COCO at NFE = 5 and NFE = 6, respectively.Source: F ADDITIONAL SAMPLES · pp. 27–28, Figure 24 · p. 37, Figure 25 · p. 37