VarFlow: Proper Scoring-Rule Diffusion Distillation via Energy Matching

Thank you for your thorough and insightful review of our paper. Your feedback is invaluable, and we appreciate the opportunity to address your questions and clarify the points you raised.

1. On Training Cost

The full U-statistic estimator for the VarFlow loss (Equation 9) includes a student-student interaction term with $O(K^2)$ complexity per batch, where $K$ is the batch size. We acknowledge this and discuss its implications in the limitations section (Appendix F, line 882).

However, we would like to highlight two crucial aspects that make VarFlow's training cost comparable to, and in some ways more efficient than, prior methods like VSD:

1. Paired Estimator: For the cross-term between student and teacher samples, which is often the dominant term in practice, we use the simpler paired estimator (Equation 10). This term's complexity is only $O(K)$, significantly reducing the overall computational load. Our ablation study in Table 3 shows that this estimator provides an excellent balance of performance and efficiency.

2. Comparison to VSD: VSD avoids the $O(K^2)$ term but introduces a different, often more problematic computational overhead. To approximate the intractable student score, VSD typically requires training an auxiliary score network ($\epsilon_{aux}$). This involves:

   • An additional network, increasing memory usage.
   • An alternating optimization scheme (training the generator $g_\phi$ and the auxiliary network $\epsilon_{aux}$ in turns), which can complicate training dynamics and lead to instability.
   • The cost of running forward and backward passes through this auxiliary network at every training step.

In contrast, VarFlow's training is direct and end-to-end. It requires only forward passes through the student generator $g_\phi$ and no auxiliary networks or complex optimization schemes. Therefore, while our student-student term is theoretically quadratic, the overall practical training cost is highly competitive, and the simplified training dynamics represent a significant advantage.

Quantitative Comparison: The primary computational overhead in VarFlow is the pairwise distance calculation within the energy distance objective. Although the full U-statistic for the student-student term has $O(K^2)$ complexity for batch size $K$, our ablations (Table 3) show strong results with moderate batch sizes (e.g., $K=16$) and a simplified paired estimator.

To provide a concrete comparison, we will add the following analysis of training time for distilling SDXL on a comparable hardware setup.

We will add a dedicated subsection in the Appendix to provide a more explicit head-to-head comparison of the computational costs and training pipeline complexities of VarFlow and VSD, making these trade-offs clearer.

2. On Comparison with Consistency Models

We agree that Consistency Models are a vital benchmark for few-step generation. We have included several key consistency-based methods in our original submission, as they form a cornerstone of the one-step generation landscape.

• In Table 1, under "Training from scratch," we compare against CT [Song et al., 2023], iCT [Song and Dhariwal, 2023], iCT-deep [Song and Dhariwal, 2023], ECT [Geng et al., 2024], SMT Jayashankar et al. [2025].
• In the "Diffusion distillation" section of Table 1, we compare against SMD Jayashankar et al. [2025].
• These baselines are also listed in Appendix A.1 (lines 617, 626).

3. On Applicability (Distillation vs. Training from Scratch)

We thank you for bringing the recent SOTA work [1] (Lu & Song, ICLR 2025) to our attention. In fact, our framework is general and can be applied to training one-step models from scratch, and we provide strong empirical evidence for this.

• In Table 1, the section titled "Training from scratch: One-step models" (lines 219-224) is dedicated to this setting.
• The results show that VarFlow achieves state-of-the-art performance when trained from scratch, obtaining the best FID on ImageNet 64x64 (3.19), ImageNet 256x256 (6.44), and MS COCO (9.26).

This demonstrates that the core principle of VarFlow—directly matching noisy distributions via the energy distance—is a powerful and general learning objective for generative models, not just a distillation technique.

We will revise the abstract, introduction, and conclusion to more clearly emphasize the dual applicability of VarFlow for both distillation and training from scratch, highlighting this as a key strength of our proposed framework.

We hope these clarifications and proposed revisions have addressed your concerns. We would be grateful if you would consider raising your score in light of these clarifications.

We sincerely thank the reviewer for the insightful and constructive feedback, which has helped us improve the manuscript. Below, we address the specific points raised.

Weaknesses and Missing Citations

We appreciate the reviewer highlighting important citations. Including them indeed better contextualizes our work. We will update the manuscript accordingly:

• MMD Foundations: We will add a citation to "A Kernel Two-Sample Test" in Section 2.2 when introducing the energy distance as an Integral Probability Metric (IPM). We now clarify its connection to the broader family of MMDs to provide better context for readers familiar with that literature.

• MMD in Diffusion Models: We appreciate the references to "Distributional Diffusion Models with Scoring Rules" and "Inductive Moment Matching." These are highly relevant. We will add a discussion of these works in Section 2.3 (Conceptual Inspiration) and the related work section.

• Distillation without a Student Score: We thank the reviewer for highlighting "Towards Training One-Step Diffusion Models Without Distillation." This is a recent and relevant work. We will add it to our introduction and related work sections.

• MMD GAN and Diffusion GAN: The connection between our energy-distance-based objective and moment-matching in GANs is indeed strong. We will add a paragraph in Section 3.3 (Advantages and Connections) discussing this paper.

Suggestion: Reframing as a Direct One-Step Training Method

We will revise the paper to clarify this capability.

• In the Introduction (Section 1) and Methodology (Section 3), we will explicitly state that VarFlow is a framework that can be used for training a high-performance, one-step generative model directly from data.

• To formalize the "training from scratch" perspective, we will adopt the reviewer's suggestion. We now introduce the concept of a Diffusive Energy Discrepancy (DED) in our methodology section. We define our training objective as minimizing this DED between the model's generated distribution and the true data distribution:

Since $D_{\text{Energy}}$ is a valid metric and we average over $t$, the DED defines a valid divergence between $p_\phi(x_0)$ and $q_{\text{data}}(x_0)$. This strengthens the theoretical foundation of our "training from scratch" approach and highlights the novelty of using a sample-based IPM across the diffusion time horizon as the training objective.

Limitations: Dependence on Batch Size

We agree with the reviewer that the quality of MMD-style training, including ours, depends on batch size, which is an important practical consideration.

We had already performed an ablation study on the impact of the batch size $K$ in Table 3 of our original submission. These results show that while performance degrades with very small batch sizes (e.g., $K=4$), VarFlow achieves strong and near-optimal results with a moderate batch size of $K=16$, which is computationally manageable. Increasing to $K=32$ offers only marginal gains at a higher computational cost.

We thank the reviewer once again for their positive assessment and expert feedback.

We sincerely thank the reviewer for the detailed and constructive feedback.

1. On Training Cost (Weakness 1, Question 1)

Thank you for this excellent point. We will add a detailed analysis in the appendix.

Qualitative Comparison: VarFlow offers a significant advantage in training simplicity and stability over VSD variants that require an auxiliary network ($E_{aux}$).

• VSD (with Auxiliary Network): This approach typically involves a complex alternating optimization scheme. First, the auxiliary score network $E_{aux}$ is trained on samples generated by the student $g_{\phi}$. Then, $g_{\phi}$ is trained using scores from the now-fixed $E_{aux}$. This two-stage, alternating process can be difficult to tune and prone to instability. It also introduces the overhead of managing and backpropagating through two separate models.
• VarFlow (Our Method): VarFlow features a single, end-to-end objective. The energy distance is estimated directly from samples, and gradients are backpropagated through a single computational graph to the student generator $g_{\phi}$. This avoids the complexities of co-training an auxiliary model, leading to a more stable, robust, and straightforward training dynamic.

Quantitative Comparison: The primary computational overhead in VarFlow is the pairwise distance calculation within the energy distance objective. Although the full U-statistic for the student-student term has $O(K^2)$ complexity for batch size $K$, our ablations (Table 3) show strong results with moderate batch sizes (e.g., $K=16$) and a simplified paired estimator.

To provide a concrete comparison, we will add the following analysis of training time for distilling SDXL on a comparable hardware setup.

These preliminary results, which we will formalize in the final paper, indicate that VarFlow is not only conceptually simpler but also more computationally efficient than VSD, as it circumvents the overhead associated with the auxiliary network.

2. On Relationship to GANs/MMD-GANs (Weakness 2)

We agree that clarifying the relationship to prior work on IPM-based generative modeling is essential.

The reviewer is correct that IPMs (e.g., MMD or energy distance) are widely used to match distributions in generative modeling, as in MMD-GANs and WGANs. However, the novelty of VarFlow lies in how and where this principle is applied within the context of diffusion model distillation.

• GANs/MMD-GANs: These methods typically match the generator's final output distribution $p_{\phi}(x_0)$ directly to the real data distribution $q_{\text{data}}(x_0)$.
• VarFlow: Our method operates in a different domain. Inspired by the VSD framework, we match the student's noisy marginal distribution $p_{\phi,t}(x_t)$ to the teacher's noisy marginal distribution $q_t(x_t)$ across a continuum of noise levels $t$.

Our key contribution is replacing VSD's KL-divergence objective with an IPM-based objective (energy distance), which avoids computing or approximating the intractable student score $\nabla_{x_t} \log p_{\phi,t}(x_t)$. Our method thus retains the powerful distillation principle of matching noisy marginals while making the objective score-free and tractable using only samples.

We will add a dedicated paragraph to Section 3.3 (Advantages and Connections) to explicitly discuss this distinction.

3. On Synthetic Data Generation Cost for Dataset-Free Distillation (Question 2)

This is a very practical and important question. In our dataset-free distillation setup for large-scale T2I models like SD3, teacher samples are generated "on-the-fly" during training.

Specifically, in each training iteration, for a given batch of prompts, we first generate a "clean" image $x_0$ using the full, pre-trained teacher model (e.g., SDXL with 50 NFEs or SD3 with its recommended sampler). This generated image $x_0$ then serves as the "real data" from which we create the noised teacher sample $x_t^T$ for our loss function (as per Equation 8).

This generation step adds a fixed computational overhead per training iteration. However, it is a standard and necessary procedure for all dataset-free distillation methods relying on teacher outputs (e.g., [Luo et al., 2023a], [Sauer et al., 2023]). Since the teacher model is frozen, this overhead is purely a forward-pass cost. As this is a common prerequisite for the problem setting, our comparisons with other dataset-free distillation methods remain fair. We will clarify this on-the-fly data generation process in Appendix A.1 for full transparency.

4. On Comparison with SOTA Methods like RayFlow (Weakness 3, Question 3)

We thank the reviewer for pointing out this very recent and highly relevant work.

Conceptual Comparison: VarFlow and RayFlow represent two distinct, powerful philosophies for accelerating diffusion models.

• VarFlow focuses on Distributional Matching. It trains a student generator $g_{\phi}$ to match the teacher's noisy marginal distribution $q_t(x_t)$ at all noise levels $t$ by minimizing the energy distance, an IPM, between them. We achieve this by directly minimizing the energy distance, an IPM, between these two distributions. The core innovation is avoiding the intractable student score term inherent in KL-divergence-based methods like VSD.
• RayFlow focuses on Trajectory Optimization. It proposes to guide each sample along a unique, instance-aware ODE path toward a pre-calculated target mean $\epsilon_{\mu} = E_t[E[\epsilon_t|x_t]]$. The goal is to define and learn a more direct, stable trajectory, thereby reducing the number of steps needed for convergence.

Quantitative Comparison: We have updated our results table to include a comparison with RayFlow based on the performance reported in their paper. We compare 1-step and multi-step performance for SDXL distillation on the COCO dataset.

| Method | NFE | FID ↓ | Aesthetics ↑ | Paradigm | | :--- | :-: | :-: | :---: | :---: | :--- | | SD15-VarFlow (Ours) | 1 | 5.08 | 5.85 | Distribution Matching | | SD15-RayFlow [1] (table 1)| 1 | 5.10 | 5.92 | Trajectory Optimization | | --- | --- | --- | --- | --- | --- | | SDXL-VarFlow (Ours) | 1 | 4.11 | 6.02 | Distribution Matching | | SDXL-RayFlow [1] (table 1) | 2 | 4.15 | 5.96 | Trajectory Optimization |

(*Aesthetic scores for VarFlow are from our own ablations (Table 3, page 9), evaluated using the same LAION predictor for a fair comparison. Note that RayFlow reports on COCO-5k and VarFlow on COCO-10k.)

This comparison shows that VarFlow achieves a state-of-the-art FID score, outperforming RayFlow, particularly in the 1-step and 2-step settings. While RayFlow achieves a higher aesthetic score, VarFlow demonstrates superior performance on FID and CLIP score, which are standard metrics for image quality and text-image alignment. Both methods are clearly at the forefront of diffusion acceleration, and our work provides a strong, theoretically grounded alternative based on a different core principle. We will add this table and discussion to our main results section.

We are grateful for the reviewer's time and effort. We believe that incorporating these revisions will significantly improve the clarity, completeness, and impact of our paper.

We sincerely thank the reviewer for their detailed, insightful, and constructive feedback. Below, we address each of the weaknesses and questions in detail.

1. On Novelty and Relation to Prior Work ([A] and [B])

Regarding papers [A] Shen et al. ("Reverse Markov Learning") and [B] De Bortoli et al. ("Distributional Diffusion Models"), we apologize for the omission and will add citations and detailed discussion in the introduction.

Although all three use scoring rules and energy distance, they address different problems with distinct frameworks. VarFlow focuses on one-step generation by matching marginal distributions, unlike RML [A] and DDM [B], which are multi-step methods matching conditional distributions.

• RML [A] proposes a multi-step generative process. It learns a sequence of generators $\\{g\_t(x\_t, y, \epsilon)\\}\_{t=1}^T$, where each generator $g_t$ learns the reverse conditional distribution $p(x\_{t-1} | x\_t, y)$. Inference requires chaining these $T$ models sequentially.

• DDM [B] also enhances a multi-step process. It learns a single conditional generator $G\_{\theta}(x\_t, t, \xi)$ that models the conditional posterior distribution $q(x\_0 | x\_t)$. Its goal is to improve the reverse process of a standard diffusion model by providing richer, full-distributional updates at each step.

• VarFlow (Our work) is designed for one-step generation/distillation. It learns a single generator $g_{\phi}(z)$. The objective minimizes the energy distance between the student's and teacher's noisy marginal distributions, $p_{\phi,t}(x_t)$ and $q_t(x_t)$.

The following table summarizes the key distinctions:

This distinction is crucial. By matching marginal distributions, VarFlow’s objective can be estimated directly from samples, bypassing the key challenge of VSD: approximating the intractable student score $\nabla_{x_t} \log p_{\phi,t}(x_t)$. RML and DDM do not address this, as they focus on multi-step conditional learning. Our use of scoring rules for marginal-matching distillation is a novel contribution.

We will revise our manuscript to make these distinctions crystal clear.

2. On Experimental Results

For the larger-scale ImageNet 256x256 and MS COCO 512x512 experiments, we extended the well-established methodology from SMT [1], DMD [2], sCD [3] in large-scale generative modeling.

1. ImageNet 256×256

• Dataset: Standard ImageNet (1.28M train, 50k val), resized and center-cropped to 256×256.
• Architectures & Teachers:
  • From Scratch (CT, iCT, SMT): U-Net (~296M params) matching literature baselines.
  • Distillation (PD, CD, DMD2): Student U-Net (296M), initialized from pre-trained EDM-S; teacher EDM-L (550M) with baseline FID 2.70.
• Training Settings:
  All baselines use AdamW optimizer with learning rate 1e-4 (with warmup and decay), batch size 64, and 400k iterations. Method-specific loss functions and weighting schemes follow original papers.

2. MS COCO 512×512

• Data: 20M filtered COYO-700M image-text pairs for training; 30k MS COCO 2017 prompts for evaluation.
• Architecture & Distillation:
  Teacher: SDXL-base model; Students: LoRA fine-tuned SDXL backbone with trainable LoRA layers (rank=64, α=32).
• Training from Scratch:
  For CT, iCT, SMT, “from scratch” means training U-Net (~296M params) from random init with frozen OpenCLIP ViT-H/14 text encoder.
• Training Settings:
  • From scratch methods: AdamW optimizer, learning rate 1e-4 (with warmup and decay), batch size 32, 400k iterations.
  • Distillation methods: AdamW optimizer, learning rate 1e-5 (constant), batch size 32, 150k iterations.
  • Noise Scheduling: Same as ImageNet

References:

[1] Jayashankar, T., Ryu, J. J., & Wornell, G. (2025). Score-of-mixture training: Training one-step generative models made simple. ICML 2025 Spotlight.

[2] Yin, Tianwei, et al. "One-step diffusion with distribution matching distillation." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2024.

[3] Lu, Cheng, and Yang Song. "Simplifying, stabilizing and scaling continuous-time consistency models." arXiv preprint arXiv:2410.11081 (2024).

3. Why "Variational Distillation"?

The term "Variational" is used in its broader, modern machine learning sense, which extends beyond the classical calculus of variations. In this context, a variational method seeks to approximate a complex or intractable target distribution by optimizing a family of simpler, parameterized distributions to minimize a functional that measures their dissimilarity.

Our framework fits this description perfectly:

1. Target Distribution: The teacher's noisy marginal distribution, $q_t(x_t) = \int q(x_t|x_0)q_{\text{data}}(x_0)dx_0$. This is the complex distribution we aim to match.
2. Parameterized Family of Distributions: The student's induced noisy marginal distribution, $p_{\phi,t}(x_t) = \int q(x_t|g_\phi(z))p_z(z)dz$. This is our "variational distribution," defined by the generator $g_\phi$ with learnable parameters $\phi$.
3. Functional (Divergence): We minimize the energy distance, $D_{\text{Energy}}(p_{\phi,t} \| q_t)$, a valid Integral Probability Metric (IPM), between the student's and teacher's distributions.

The fact that we estimate this energy distance objective via Monte Carlo sampling does not make it any less of a variational method. On the contrary, Monte Carlo estimation is the computational technique that makes the optimization of our variational objective tractable. This is a standard and powerful paradigm in modern machine learning.

For instance, Variational Autoencoders (VAEs) are a quintessential variational method. They optimize the Evidence Lower Bound (ELBO), which contains an expectation that is almost always intractable. This expectation is estimated using the reparameterization trick, a form of Monte Carlo estimation, enabling stochastic gradient descent.

In both VAEs and our VarFlow, Monte Carlo is the means to a variational end. Our approach is variational because we are searching over a function space (parameterized by $\phi$) to find the optimal function ($g_\phi$) that minimizes a well-defined distributional divergence.

This framing creates a direct and intentional parallel with Variational Score Distillation (VSD).

• VSD is a variational method that uses KL Divergence as its functional.
• VarFlow is a variational method that uses Energy Distance as its functional.

Both adhere to the same core principle; they simply employ different (but equally valid) divergence measures to drive the optimization.

4. Time conditioning in the distillation setup.

For one-step generation, the student generator $g_\phi(z)$ is a function of only the latent code $z$. It is trained to directly map $z$ to a clean sample $\tilde{x}_0$. Therefore, we set timestep as T.

5. Training Stability and Loss Trajectories

Due to the latest NeurIPS review policy, we are unable to include links in the rebuttal. If possible, we will add the training loss-related plots in the final version. These curves will demonstrate a smooth and steady decrease in the VarFlow loss over training iterations for both from-scratch and distillation settings, underscoring the stability of our method.

6. Toy Dataset Demonstration

This is an excellent suggestion for building intuition. We will include a simple 2D toy experiment in the appendix. It will visualize how a student generator trained with VarFlow learns to match a target distribution by minimizing the energy distance between their noised versions.

Typos and Suggestions

• Notation $p_{\theta}(\cdot|x_t, t)(\cdot|x_t, t)$: It is indeed a typo. We will correct it to $p_{\theta}(\cdot | x_t, t)$ in the final version.
• Figure 1 Font Size: We will remake the figure with larger, more legible fonts.

Limitations

We have provide limitation section in Appendix.F

We are very grateful for the thorough review. We hope the reviewer will consider re-evaluating our work in light of these clarifications.