Improved Training Technique for Shortcut Models

Thank you for recognizing the core strengths of our work, including the significant FID improvements, enhanced flexibility, and the systematic approach to solving multiple challenges. We appreciate your suggestions for improvement.

Note: References to tables, figures, or propositions without the 'Rebuttal' prefix refer to content located in the main paper.

Q1: ImageNet 512x512 and different backbone

Our work introduces a robust, architecture-agnostic training methodology for shortcut models. While we report results on SiT-XL/2 to align with the Shortcut Model baseline, our technique generalizes beyond Transformer-based models. Due to limited time, we tackle both backbone and larger image size within one experiment. Specifically, we scaled our approach to ImageNet 512×512 using XL/2 variant of FlowDCN, a fully convolutional generative model with group-wise deformable convolution blocks instead of a DiT-based model. Given the rebuttal time, our models have not been fully trained, and we report their scores at the latest iteration (300K) in the table below. We achieved a substantial FID improvement compared to the original shortcut model training, demonstrating both scalability and broad applicability.

Rebuttal Table 1: ImageNet 512x512 experiment.

Q2: Text-to-image and compare multi-step model

We thank the reviewer for pointing out the value of a broader multi-step comparison. In fact, in Table 1 we compared iSM against the recent multi-step class-conditioned diffusion and flow models on ImageNet 256×256 (e.g. DiT-XL/2 (250 steps), SiT-XL/2 (250 steps), U-ViT-H (50 steps), SimDiff (512 steps),...).

Regarding text-to-image (T2I) models such as Stable Diffusion, we note the following:

1. Our work is scoped to class-conditioned ImageNet generation (as do prior methods like DiT, SiT, SM, IMM).
2. Extending to T2I would require both a text encoder and retraining on large scale dataset like LAION datasets, far beyond our current compute budget.
3. Unlike many text-to-image models, most class-conditional methods do not report, or perform poorly at, the 50-step sampling budget. To our knowledge, U-ViT-H is the only class-to-image model with competitive FID at 50 steps, and we included its result in Table 1 for direct comparison.
4. Architecturally, we build on DiT-based backbones, which have recently been demonstrated (e.g., in PixArt-Alpha) to scale effectively to text conditioning. We therefore believe our improvements would transfer directly to the T2I setting.

In summary, Table 1 shows iSM competitive performance with respect to SOTA class-conditional generative models on ImageNet, and while a direct text-to-image model comparison is an exciting future direction, it lies outside our current class-to-image benchmark and compute resources.

Q3: Scaling law experiments for iSM

Due to limited time of rebuttal and computational budget, we couldn't scale up the model to a larger size than DiT-XL, hence we follow SM to run the scaling law experiment from smaller model (DiT-B) to our final model (DiT-XL) as a proxy for larger-scale behavior. The results in Rebuttal Table 2 clearly show that the performance gap between our iSM and the baseline SM widens significantly as model size increases, suggesting strong scaling properties.

Rebuttal Table 2: Scaling law experiment in ImageNet 256x256.

Q4: Provide complete training resources and reproducible steps?

Yes. We have already included our implementation details in the Appendix and a checklist in the main paper. All experimental workflows, including training and evaluation, were conducted on 8 x NVIDIA A100-40GB GPUs. We are committed to open-sourcing our code and releasing model checkpoints upon acceptance to ensure full reproducibility for the community.

The author has answered my questions. I believe this method demonstrates good generality and scalability, so I will be raising my score.

Thank you for the positive and constructive review. We appreciate that you found the paper well-written, the components novel, and the improvements significant. We address your points below.

Note: References to tables, figures, or propositions without the 'Rebuttal' prefix refer to content located in the main paper.

Q1: Experiments on ImageNet 512x512, other backbones, and generalizability.

Our work introduces a robust, architecture-agnostic training methodology for shortcut models. While we report results on SiT-XL/2 to align with the Shortcut Model baseline, our technique generalizes beyond Transformer-based models. Due to limited time, we tackle both backbone and larger image size within one experiment. Specifically, we scaled our approach to ImageNet 512×512 using XL/2 variant of FlowDCN, a fully convolutional generative model with group-wise deformable convolution blocks instead of a DiT-based model. Given the rebuttal time, our models have not been fully trained, and we report their scores at the latest iteration (300K) in the table below. We achieved a substantial FID improvement compared to the original shortcut model training, demonstrating both scalability and broad applicability.

Rebuttal Table 1: ImageNet 512x512 experiment.

In line with common practice in diffusion model research, our work focuses on class-to-image tasks, as scaling to full text-to-image generation is challenging under current resource and time constraints. Many influential papers in this field, such as DiT, SiT, EDM, Consistency Models, and Shortcut Models, have established their contributions through experiments on label-to-image generation, thereby laying the groundwork for future scaling to text-to-image diffusion models.

Q2: Detailed algorithm for sOT

The procedure is as follows:

1. Aggregate Batches: For every K training steps, we collect and hold the noise samples and data latents from K separate mini-batches of size M.
2. Compute Large-Scale OT Plan: We compute a single OT plan over the aggregated set of K*M samples.
3. Redistribute and Train: The resulting K*M optimally paired samples are then split back into K mini-batches, which are used for the subsequent K model training steps.

This process allows the model to benefit from a large-scale transport plan without the memory burden of a large training batch size. We will add pseudo-code algorithm for Scaling Optimal Transport (sOT) into the paper.

Q3: Computational and memory overhead of sOT

The computational overhead from sOT is a modest ~4% increase in total training time, while the memory overhead is negligible.

• Memory Efficiency: Our sOT implementation is memory-efficient because it decouples the OT batch size (K*M) from the model's training batch size (M). We only need to store the K*M latent vectors, while the memory-intensive model weights and gradients are processed with the standard, smaller batch size M.
• Computational Cost: For K=32 and a per-GPU batch size of M=64, the OT plan is computed over 2048 samples every 32 steps. This adds only ~4% to the total training time on an A100 GPU, a small price for the significant gains in trajectory stability and final FID scores. We report the training iteration speed for different values of K (where K=0 indicates no OT used).

Rebuttal Table 2: Training speed with sOT in ImageNet 256x256.

Q4: I wonder if Multi-Level Wavelet Function is scalable to resolution, e.g., if it still works with ImageNet 512x512.

Yes, the Multi-Level Wavelet Loss is highly scalable and particularly well-suited for higher resolutions. The Discrete Wavelet Transform (DWT) is an inherently multi-scale algorithm. Its ability to disentangle features across frequency bands becomes even more critical for high-resolution images, where preserving fine-grained textures (high-frequency) alongside global composition (low-frequency) is essential for quality. This makes it a robust choice ([1], [2]) for scaling to 512x512 and beyond. Our result in Rebuttal Table 1 on ImageNet 512x512 also strengthen this argument, since iSM significantly outperforms the original shortcut model.

Q5: Robustness of the interval guidance hyperparameter

We affirm that $t_{\text{interval}}$ is a robust hyperparameter, a finding supported by theoretical motivation, our internal ablation studies, and external validation from concurrent work.

1. SNR-Based Principle: The parameter is not arbitrary but is grounded in the signal-to-noise ratio (SNR) dynamics of diffusion models. Notably, previous work [1] found that guidance is ineffective or harmful at very high noise levels (low $t$), and $t_{\text{interval}}$ simply implements this principle.

2. Empirical Robustness (Table 2): Our ablation study demonstrates a wide effective range, not a single brittle value.

• Even a small cutoff at $t_{\text{interval}} = 0.1$ substantially improves performance ($\text{FID}_{N=1}$ drops from 9.62 to 8.58), showing the primary benefit is avoiding guidance in the earliest, highest-noise steps.
• Performance degrades only when the cutoff is overly aggressive (e.g., $t_{\text{interval}} = 0.5$), confirming a stable optimal window rather than a specific magic number.
• Furthermore, we have trained a new model (different architecture, FlowDCN) and new dataset (ImageNet 512x512) with the same $t_{\text{interval}}=0.3$ and the performance is still very competitive (in Rebuttal Table 1)

3. Consistency with Prior Work: Our finding is strongly validated by prior work [1], show this principle is critical for state-of-the-art results on diverse setups like ImageNet-512 and Stable Diffusion XL. Both our work (integrating interval guidance into the training loss) and theirs (applying it at inference) identify a similar optimal range around $t=0.3$.

This convergence of evidence confirms that $t_{\text{interval}}=0.3$ is a robust and principled default, not a dataset-specific choice.

Q6: Additional evaluation metrics (FD-DINOv2, IS, Precision, Recall)

We report FD-DINOv2, Inception Score (IS), Precision, and Recall for the baseline Shortcut Model (SM), Inductive Moment Matching (IMM), and our improved Shortcut Model (iSM). The results below show that iSM consistently outperforms the baseline SM across multiple NFE settings.

Rebuttal Table 3: Additional evaluation metrics for ImageNet 256x256.

Q7: Regarding Formatting Concerns:

Thank you for these helpful suggestions. We will them in the revision:

• We will cite REPA [3] in our component-wise analysis section as an inspiration for the evaluation structure.
• We will reposition Table 2 to the top of Page 9 for improved readability and flow.

[1] Applying Guidance in a Limited Interval Improves Sample and Distribution Quality in Diffusion Models, NeurIPS 2024.

[2] Diffusion-4K: Ultra-High-Resolution Image Synthesis with Latent Diffusion Models, CVPR 2025.

[3] Representation alignment for generation: Training diffusion transformers is easier than you think, ICLR 2025.

Thank you for your valuable feedback and questions. We are glad you found the paper easy to follow and the problem important.

Note: References to tables, figures, or propositions without the 'Rebuttal' prefix refer to content located in the main paper.

Q1: Additional evaluation metrics (FD-DINOv2, IS, Precision, Recall)

We report FD-DINOv2[1], Inception Score[2] (IS), Precision[3], and Recall[3] for the baseline Shortcut Model (SM), Inductive Moment Matching (IMM), and our improved Shortcut Model (iSM). The results below show that iSM consistently outperforms the baseline SM across multiple NFE settings.

Rebuttal Table 1: Additional evaluation metrics for ImageNet 256x256.

Q2: Is the "accumulated guidance" issue a well-known problem? Is it a contribution?

To the best of our knowledge, the explicit formalization and analysis of the accumulation of guidance effects in shortcut models is a novel contribution of our paper. While CFG is well-studied, its compounding effect in shortcut models has been overlooked. We are the first to identify, formalize (Proposition 1), and solve this critical flaw. We’ll revise the introduction to emphasize this contribution, making its importance immediately clear.

Q3: Significance of OT gains and intuition for high-dimensional OT

1. OT operates in Latent Space: We agree that OT in 256x256 pixel space may be intractable. In fact, all our operations, including OT matching, are performed in the 32x32 latent space of a pre-trained VAE. We apologize if this was not clear.
2. Significance of Gains: In the high-performance regime (FID < 3.0), improvements are non-linear and dropping from 2.64 to 2.23 is a major leap, since each FID reduction in this regime becomes increasingly challenging (for example, in the paper SiT's Table 1, last two rows, when applying CFG, the improvement drop from 2.27 to 2.06)

Q4: If I understand well your benchmark, GANs is still the best 1-step generation model for conditional generation?

It’s true that state-of-the-art GANs achieve the lowest FID in a single step, but they’re notoriously hard to train and often suffer from diversity collapse. In contrast, our diffusion-based iSM trains stably (no adversarial loss) and delivers significantly better sample diversity (Recall of 0.61 vs. StyleGAN-XL of 0.50), narrowing the FID gap while ensuring a stronger diversity–fidelity trade-off. Furthermore, iSM’s flexible sampling schedules and efficient guidance sampling set it apart from other diffusion models.

Q5: Are Sections 4-6 a contribution of the paper?

Absolutely. In fact, in the Introduction (Sec. 1, lines 53–76) we explicitly identify these shortcomings of vanilla shortcut models: (1) low-frequency bias, (2) high-variance/curved trajectories, and (3) EMA-lagged self-consistency, and then state our three complementary contributions to address them:

• Multi-Level Wavelet Loss (Section 4) to correct frequency bias;
• Scaling Optimal Transport Matching (Section 5) to straighten and stabilize trajectories;
• Twin EMA Strategy (Section 6) to remove the temporal lag in self-consistency targets.

These sections detail each of our contributions. As shown in Table 2, progressively adding the components from Sec. 4–6 steadily improves the model’s performance, confirming their novelty and importance.

Q6: Could authors provide additional information on the "Multi-Level Wavelet Function"?

Certainly. The Multi-Level Wavelet Loss functions as a multi-scale objective that tackles the frequency bias. Instead of naively applying L2 loss in the spatial domain, it recursively applies Discrete Wavelet Transform (DWT) to the prediction and target. This decomposes the error signal across various frequency bands, forcing the network to minimize errors in both low-frequency structures and high-frequency details simultaneously. For reference, Appendix B provides the pseudocode (Algorithm 1).

Q7: Do you use a latent space?

Yes. Following previous works like DiT, SiT, SM, our model operates entirely in the latent space of the pre-trained sd-vae-ft-mse autoencoder. We will make this detail more prominent in the main text to avoid any confusion.

[1] Exposing flaws of generative model evaluation metrics and their unfair treatment of diffusion models, NeurIPS 2023.

[2] Improved techniques for training gans, NeurIPS 2016.

[3] Improved Precision and Recall Metric for Assessing Generative Models, NeurIPS 2019.

Thank you for your detailed review and constructive feedback. We appreciate that you found our motivation clear, our proposed solution cohesive, and our experiments thorough. Below, we provide a point-to-point response and a summary of the corresponding revisions.

Note: References to tables, figures, or propositions without the 'Rebuttal' prefix refer to content located in the main paper.

Q1: Experiments on ImageNet 512x512, other backbones, and generalizability.

Our work introduces a robust, architecture-agnostic training methodology for shortcut models. While we report results on SiT-XL/2 to align with the Shortcut Model baseline, our technique generalizes beyond Transformer-based models. Due to limited time, we tackle both backbone and larger image size within one experiment. Specifically, we scaled our approach to ImageNet 512×512 using XL/2 variant of FlowDCN, a fully convolutional generative model with group-wise deformable convolution blocks instead of a DiT-based model. Given the rebuttal time, our models have not been fully trained, and we report their scores at the latest iteration (300K) in the table below. We achieved a substantial FID improvement compared to the original shortcut model training, demonstrating both scalability and broad applicability.

Rebuttal Table 1: ImageNet 512x512 experiment.

In line with common practice in diffusion model research, our work focuses on class-to-image tasks, as scaling to full text-to-image generation is challenging under current resource and time constraints. Many influential papers in this field, such as DiT, SiT, EDM, Consistency Models, and Shortcut Models, have established their contributions through experiments on label-to-image generation, thereby laying the groundwork for future scaling to text-to-image diffusion models.

Q2: Comparison with other cutting-edge models

A direct FID comparison between iSM and models like REPA can be misleading, as they operate in fundamentally different inference-efficiency regimes. REPA is a state-of-the-art full-trajectory diffusion model requiring 250 NFE to achieve its FID of 1.42. Our iSM achieves a competitive FID of 1.93 with just 8 NFE. This represents a >30x reduction in computational cost for the trade-off in final FID. Our contribution lies in improving shortcut models across at extreme NFE efficiency, achieving a significant FID reduction in these settings over the original baseline and competitive performance with other model families.

Q3: Additional evaluation metrics (FD-DINOv2, IS, Precision, Recall)

We report FD-DINOv2, Inception Score (IS), Precision, Recall for the baseline Shortcut Model (SM), Inductive Moment Matching (IMM), and our improved Shortcut Model (iSM). The results below show that iSM consistently outperforms the baseline SM across multiple NFE settings.

Rebuttal Table 2: Additional metrics for ImageNet 256x256.

Q4: The Scaling OT (sOT) method is very interesting. Have you investigated the effect of the update frequency?

Yes, we performed an ablation over the update interval K to measure its impact (see Table 2 in main paper). We found:

• Better training: Increasing K yields a closer match to the global transport plan, which straightens training trajectories and boosts stability.
• Minimal overhead: Even at a larger K, the extra overhead remains modest (~4% in our final setup, K = 32).
• Saturation point: As shown in Table 2, the FID scores improvements plateau past a certain K, so there’s little benefit in updating more frequently.

Q5: Does the accumulated guidance problem persist for other base step sizes (N)?

Yes, and we emphasize that this is one of our key contributions. The accumulated guidance problem is a systemic flaw in the recursive formulation of shortcut models. Our Proposition 1 provides the first formalization of this compounding error, showing that the effective guidance scales undesirably as $w^{log_2(N)}$.

Our Integrated Guidance framework resolves this issue by making the guidance as an explicit model input w, we train the network to learn the mapping (x, t, c, d, w) -> v_guided directly. This bypasses the flawed recursive formulation entirely, making our solution inherently robust and independent of the base step size N. We’ll underscore this contribution in Section 3.2.

Q6: Intuition for Twin EMA decay rate (0.95)

The selection for decay rate for fast-decay EMA ($\theta_{target}$) needs to satisfy the following:

• Low enough to keep the self-consistency target network closely aligned with the online network, avoiding the distributional lag as training progresses.
• High enough to preserve the benefits of EMA, which smooths noisy mini-batch gradients and prevents training instability.

Therefore, a much lower rate decay for $\theta_{target}$ risks instabilility and negates the benefits of EMA, while a much higher rate negates the benefit of the twin strategy and reintroduces distributional lags. Our ablation in Table 2 empirically validates that 0.95 provides the optimal trade-off. The decay rate of 0.95 allows the self-consistency target to closely track the online network, but with just enough lag to smooth out high-frequency noise in the training signal.

Q7: The gap between one-step and multi-step performance

We would like to clarify that our paper does not claim to completely eliminate the gap between one-step and multi-step generation. Instead, our central contribution is to significantly reduce the performance degradation at lower NFEs compared to previous shortcut models, while also providing a single flexible model that performs well across a wide range of sampling settings. We identified the degradation as primarily resulting from: (1) accumulated guidance, (2) low-frequency bias, (3) high-variance/curved trajectories, and (4) EMA-lagged self-consistency. For instance, our iSM model improves the FID from a much higher value in SM (10.60) down to 5.27 for one-step and further down to 2.05 for four-step generation. As noted in our conclusion (L381-383), while bridging this gap remains a challenging task, it offers an exciting and promising avenue for future research.