ALTER: All-in-One Layer Pruning and Temporal Expert Routing for Efficient Diffusion Generation

Thank you for your review and for recognizing the novelty and solid performance of our work. We have performed several new experiments to address your primary concerns.

W1 - Generalization to DiT Architectures

We agree that demonstrating generality is crucial. To address this, we have conducted a new experiment applying ALTER to a DiT-XL/2 backbone. The results show that our method achieves a ~1.92x speedup while maintaining an FID score (2.36) comparable to the full model (2.27).

With approximately half the layers active on average, our method significantly outperforms static pruning (TinyFusion), especially on key quality metrics like IS (254.29 vs. 234.50) and FID (2.36 vs. 2.86). When comparing to dynamic methods like DyDiT, ALTER brings higher speedup. It's important to note that our approaches are orthogonal and complementary. ALTER performs coarse-grained layer-level pruning, while DyDiT performs fine-grained pruning within each block. These new results demonstrate that ALTER is a general and effective framework for accelerating various diffusion architectures.

We have not implemented ALTER on SD3 and FLUX because of the limited computational resources and rebuttal time, but our successful application to DiT provides strong evidence that our approach is a general framework for accelerating various diffusion backbones.

W2 - Comparison with Other Baselines

Thank you for providing this comprehensive list. Here is how we position our work with respect to these methods:

• DyDiT++ [4] & TPDM [5] : As shown above, our DiT experiment now includes a comparison with DyDiT. But it's important to note that our approaches are orthogonal and complementary. ALTER performs coarse-grained layer-level pruning, while DyDiT performs fine-grained pruning within each block. TPDM is an orthogonal approach that accelerates inference by predicting and skipping entire timesteps. Our method, which prunes layers within a step, could be combined with TPDM to accelerate the remaining, non-skipped steps.

• Few-step & Distilled Models (SD-turbo, FLUX-turbo): These methods represent a different family of acceleration that uses distillation to train a specialized model for extremely few-step (e.g., 4 step) inference. Our work focuses on accelerating the original, pre-trained model for general-purpose use across various step counts (e.g., 20-50 steps) without requiring a full distillation process. Besides, step distillation methods are computationally intensive, often demanding thousands of GPU hours, whileour ALTER method is lightweight, requiring only about 48 A100 hours to fine-tune the SDv2.1 model.

• FLUX-lite: This is a statically pruned version of the FLUX architecture. Our primary experiments focus on accelerating the public SD2.1 and DiT backbones. Where we do include a statically pruned version of the SDv2.1 (BK-SDM-v2) and the DiT-XL/2 (TinyFusion).

We have updated our related work section to include a detailed discussion of these methods, clarifying the different axes of acceleration and contextualizing our contribution.

Q1 - Why the speedup with 25 steps is still significant in Tab. 1.

As shown in Fig.4, the pruned layers are almost at the last part, which is at high-resolution considering the UNet's architecture. The significant speedup arises from the highly non-uniform computational cost of layers in a U-Net. Although the pruned layer's number is not big, prunning the high-resolution could benefit more for acceleration compared with the middle layers.

Q2 - GenEval [6] and Human Evaluation

Thank you for the valuable suggestion to use GenEval. We made a significant effort to incorporate it but encountered unresolved technical issues with the official GenEval repository's environment setup, which we were unable to fix within the limited rebuttal period. And these technical problems are very common as shown in the github repository's opening issues. Thanks for your understanding.

However, to follow the spirit of your suggestion for a more comprehensive evaluation we have benchmarked our method using ImageReward [7] and Human Preference Score v2 (HPSv2) [8], two state-of-the-art models designed to measure human aesthetic and semantic preferences. We calculate these two metrics on their own benchmark datasets. Since APTP is prompt-based pruning, we use the prompt split method they provided to divide the benchmark's prompts to different experts finetuned on MS-COCO and CC3M.

As shown in the following two tables, the results are highly encouraging. On ImageReward, our ALTER models rank #1 (0.65 MACs) and #2 (0.60 MACs), outperforming prior methods and even the original SDv2.1 baseline, which indicates a higher alignment with human preference. Similarly, on the fine-grained HPSv2 benchmark, our method consistently achieves the highest average score, demonstrating robust performance across diverse visual styles from animation to photo-realism.

References:

[1] Peebles et al. DiT: Diffusion Models with Transformer Backbones.

[2] Li et al. TinyFusion: Layer Fusion for Diffusion Speed-up.

[3] Zhao et al. Dynamic Diffusion Transformer.

[4] Lao et al. DyDiT++: Dynamic Diffusion Transformers for Efficient Visual Generation.

[5] Ye et al. Schedule On the Fly: Diffusion Time Prediction for Faster and Better Image Generation.

[6] Ghosh et al. GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment.

[7] Xu et al. ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation.

[8] Wu et al. Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis.

Thanks to the detailed rebuttal. I shall maintain my positive rating.

We thank you for the insightful review and for recognizing the strengths of our work. Below we clarify our novelty compared to prior art and provide new results to address your questions.

W1 - Novelty and Comparison to APTP

We thank the reviewer for this point, and we have revised the paper to better articulate our novelty. While both methods use experts, ALTER introduces two fundamental differences from APTP: a dynamic, temporal routing mechanism and a unified, single-stage training framework.

1. Routing Mechanism & Model Utilization: APTP uses static, prompt-based routing. Each prompt is assigned to one expert sub-network for all timesteps. This means an image is generated using only a static portion of the model's total capacity. In contrast, ALTER uses dynamic, temporal routing. Every prompt is processed by a sequence of different experts, one for each timestep. This allows every image generation to leverage the full capacity of the model, assigning the most suitable specialized expert (e.g., "structure expert" vs. "texture expert") to each stage of the diffusion process as illustrated in Figure 1 of our paper.
2. Training & Deployment Framework: APTP requires a complex, multi-stage pipeline involving dataset clustering, experts determining and separate fine-tuning for each expert, resulting in multiple model checkpoints. ALTER is an end-to-end framework that learns the routing and experts jointly in a single stage, producing a single, unified model that is far simpler to manage and deploy.

These distinctions are critical to our method's ability to improve upon the baseline's generation quality, a result not achieved by prior pruning work.

W2 - Generalizability to Other Diffusion Architectures

We agree that demonstrating generality is crucial. To address this, we have conducted a new experiment applying ALTER to a DiT-XL/2 [1] backbone. The results show that our method achieves a ~1.92x speedup while maintaining an FID score (2.36) comparable to the full model (2.27).

With approximately half the layers active on average, our method significantly outperforms static pruning (TinyFusion [2]), especially on key quality metrics like IS (254.29 vs. 234.50) and FID (2.36 vs. 2.86). When comparing to dynamic methods like DyDiT [3], ALTER brings higher speedup. It's important to note that our approaches are orthogonal and complementary. ALTER performs coarse-grained layer-level pruning, while DyDiT performs fine-grained pruning within each block.

This empirically confirms that our framework is effective beyond U-Nets and is applicable to modern Transformer-based architectures.

W3 - The Choice of Experts Number

Thank you for this advice. We agree with the reviewer that claiming 10 experts as universally "optimal" is an overstatement. Our intention was to identify a practical and effective configuration for the SDv2.1 model.

We have revised the manuscript to tone down this claim, and refer to 10 experts as a "well-performing configuration" that demonstrates a strong efficiency-quality trade-off. Furthermore, our new experiments on the DiT architecture also successfully utilized 10 experts, achieving the strong results reported (responses to W2). This suggests that 10 is a reasonable choice across different architectures, though the ideal number may vary with the specific model and dataset.

Q1 - Why does pruning primarily happen at the final layers?

This is an excellent question. Our hypothesis is that it reflects the different roles of layers in the U-Net. Early/middle layers capture foundational spatial structures and low-level features that are essential throughout the entire diffusion process. Pruning them would be too destructive. In contrast, the final decoder layers are responsible for high-level semantic refinement and adding fine-grained details, tasks that are more stage-specific. Our time-aware router learns that these specialized refinement tasks can be effectively handled by smaller experts or even pruned at certain timesteps without compromising the core structure. This observation is also consistent with prior work like APTP, which noted that deeper layers were more amenable to pruning than the foundational ones.

Q2 - Is batch-parallelism possible with this approach?

Thank you for this critical question. Yes, batch-parallelism is possible for our framework. Our routing mechanism is deterministic, which means that for a given prompt batch and a specific timestep $t$, the routing decision is identical for all samples in a batch. The entire batch can therefore be processed efficiently through the same selected expert sub-network. This provides a significant practical advantage over prompt-routing methods like APTP. In APTP, different prompts within a single batch may be routed to different experts, breaking efficient batching and requiring serialization or complex batch management that can harm real-world throughput.

However, to ensure fair and reproducible comparisons, our latency benchmarks in the paper follow prior work and use a batch size of 1.

Q3 - How much of the performance is accounted for by the final Fine-tuning after T-end?

Thanks for your suggestion. To show the contribution of the final fine-tuning stage, we have conducted the requested ablation. The results below are for our ALTER (0.65) model on MS-COCO 2017 dataset.

As the results demonstrate, the model's performance after T$_{end}$ is already very strong, achieving the majority of the final quality. The final fine-tuning provides a consistent and important boost across both metrics and step counts, further improving both FID and CLIP scores.

References:

[1] Peebles et al. DiT: Diffusion Models with Transformer Backbones.

[2] Li et al. TinyFusion: Layer Fusion for Diffusion Speed-up.

[3] Zhao et al. Dynamic Diffusion Transformer.

We sincerely thank you for your detailed review and positive assessment of our work's novelty and motivation. We are pleased to present new results that address your concerns.

W1 & Q2 - More Comprehensive Empirical Evaluation

Thank you for suggesting a more comprehensive evaluation with modern metrics. We agree this is crucial for assessing generative quality. In response, we have benchmarked our method using ImageReward [1] and Human Preference Score v2 (HPSv2) [2], two state-of-the-art models designed to measure human aesthetic and semantic preferences. We calculate these 2 metrics on their own benchmark datasets. Since APTP is prompt-based pruning, we use the prompt split method they provided to divide the benchmark's prompts to different experts finetuned on MS-COCO and CC3M.

As shown in the following two tables, the results are highly encouraging. On ImageReward, our ALTER models rank #1 (0.65 MACs) and #2 (0.60 MACs), outperforming prior methods and even the original SDv2.1 baseline, which indicates a higher alignment with human preference. Similarly, on the fine-grained HPSv2 benchmark, our method consistently achieves the highest average score, demonstrating robust performance across diverse visual styles from animation to photo-realism.

W2 - Behavior of the Experts

This is an excellent question that makes us clarify the relationship between our dynamic pruning method and caching. The core difference is that caching methods reuse a layer's output features from a previous step, whereas our method, ALTER, learns to skip the entire computation of a layer.

Caching methods like DeepCache [3] show that the outputs of early-to-mid layers can be reused across steps, as they capture low-frequency information that evolves slowly. However, our results show that foundational early and middle layers are rarely pruned, suggest that while their outputs might be stable enough to be cached, the computations performed within them are still essential at every step to progressively refine the feature maps for later layers. Skipping these layers entirely is too detrimental to the final image quality. This observation is also consistent with prior work like APTP, which noted that deeper layers were more amenable to pruning than the foundational ones.

We agree that encouraging greater routing diversity is a valuable direction for future work, potentially using regularization techniques. We also strongly agree that caching methods could complement ALTER. This highlights a powerful scheme: ALTER could first prune the model's architecture, and a caching scheme could then be applied to the remaining active layers for further speedups.

Q1 - Does the performance hold for other diffusion backbones?

To empirically validate that our framework generalizes beyond U-Nets, we have conducted a new experiment applying ALTER to a DiT-XL/2 [4] backbone on ImageNet (256 $\times$ 256 resolution generation, 250 steps). For the DiT architecture, each Transformer block was treated as a prunable unit, and the ALTER framework was applied without fundamental architectural changes.

With approximately half the layers active on average, our method significantly outperforms static pruning (TinyFusion [6]), especially on key quality metrics like IS (254.29 vs. 234.50) and FID (2.36 vs. 2.86). It is nearly comparable with the full DiT-XL model, demonstrating a good quality-efficiency trade-off. When comparing to dynamic methods like DyDiT [5], ALTER brings higher speedup. It's important to note that our approaches are orthogonal and complementary. ALTER performs coarse-grained layer-level pruning, while DyDiT performs fine-grained pruning within each block.

This successful application of ALTER to a DiT backbone and its orthogonality to fine-grained methods, validates that ALTER is a general and effective framework for accelerating various diffusion architectures.

References:

[1] Xu et al. ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation.

[2] Wu et al. Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis.

[3] Ma et al. DeepCache: Accelerating Diffusion Models for Free.

[4] Peebles et al. DiT: Diffusion Models with Transformer Backbones.

[5] Zhao et al. Dynamic Diffusion Transformer.

[6] Li et al. TinyFusion: Layer Fusion for Diffusion Speed-up.

We sincerely thank you for your constructive feedback and for recognizing the novelty and solid empirical results of ALTER. Below we address each question in turn.

Q1 - Add SSIM and PSNR to measure the generation deviation from the pre-trained model

We provide the SSIM and PSNR to show the deviation of different methods from the original SDv2.1 25 step inference.

As shown in the above table, our method ALTER not only matches the MACs of prior work but also achieves significantly higher SSIM and PSNR scores, which demonstrates that our approach better preserves the generative characteristics of the original model.

To provide an even more comprehensive evaluation of generation quality, we also benchmarked our models using advanced preference models, ImageReward [1] and Human Preference Score v2 (HPSv2) [2] (all at 25 steps). These metrics assess perceptual quality and human aesthetic preference. As shown below, our method ranks highest on both metrics, even outperforming the original SDv2.1, confirming that our efficiency gains do not come at the cost of visual quality.

Q2 - Scalability / Orthogonality to other accelerators

We appreciate the reviewer for highlighting this key advantage of our work. We strongly agree that ALTER’s design makes it orthogonal to and compatible with other diffusion acceleration methods such as Diffusion Cache and token reduction.

Here we conduct a new experiment combining ALTER with a token reduction method ToMeSD [3] to demonstrate ALTER's scalability. For this combination, we follow the default setting of ToMeSD. As shown in the following table, we boost the total speedup from 3.64x to an impressive 6.54x with only a minor, controllable trade-off in performance scores. This empirically validates the excellent scalability of our method.

Regarding caching methods like DeepCache [4], we agree this is another exciting direction. Since ALTER dynamically prunes layers, a compatible cache would only need to store features and reuse for current active layers, potentially enabling a more memory-efficient caching scheme. We will leave this as future work.

Q3 - Rename "Temporal Expert Routing" to "Step Expert Routing"

Thank you for this suggestion. We appreciate you pointing out that "Step Expert Routing" might be more immediately familiar to some readers in the diffusion community. We have carefully considered this change. However, we have a deliberate reason for our choice of "Temporal" and would prefer to retain it. Our motivation is that the routing mechanism is conditioned on the timestep variable $t$ of the diffusion process, which is represented by the time embedding. While it is discretized into steps during diffusion sampling, our framing is rooted in the underlying temporal nature of the process.

Q4 - Standardize symbol usage

Thank you for your careful reading and pointing out this redundancy. We have clearly checked the notation definition and standardized all notation in the revised manuscript. We now use $t$ exclusively for the diffusion timestep and have introduced $k$ to denote the training iterations.

References:

[1] Xu et al. ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation.

[2] Wu et al. Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis.

[3] Bolya et al. Token Merging for Fast Stable Diffusion.

[4] Ma et al. DeepCache: Accelerating Diffusion Models for Free.