FastFLUX: Pruning FLUX with Block-wise Replacement and Sandwich Training

TL;DR

The paper addresses the high computational cost of diffusion transformers by introducing FastFLUX, an architecture-level pruning framework that combines Block-wise Replacement with Linear Layers (BRLL) and Sandwich Training (ST) to preserve generation quality. BRLL replaces heavy residual branches with linear layers, initialized via a closed-form least-squares solution, while ST provides LoRA-based localized fine-tuning around replaced blocks to mitigate degradation. The approach maintains performance with significant inference speedups at modest pruning ratios (e.g., 20%), and demonstrates generalization to architectures beyond FLUX, including Stable Diffusion v3.0. Overall, FastFLUX offers a scalable, low-training-cost path to deployable diffusion models with competitive quality.

Abstract

Recent advancements in text-to-image (T2I) generation have led to the emergence of highly expressive models such as diffusion transformers (DiTs), exemplified by FLUX. However, their massive parameter sizes lead to slow inference, high memory usage, and poor deployability. Existing acceleration methods (e.g., single-step distillation and attention pruning) often suffer from significant performance degradation and incur substantial training costs. To address these limitations, we propose FastFLUX, an architecture-level pruning framework designed to enhance the inference efficiency of FLUX. At its core is the Block-wise Replacement with Linear Layers (BRLL) method, which replaces structurally complex residual branches in ResBlocks with lightweight linear layers while preserving the original shortcut connections for stability. Furthermore, we introduce Sandwich Training (ST), a localized fine-tuning strategy that leverages LoRA to supervise neighboring blocks, mitigating performance drops caused by structural replacement. Experiments show that our FastFLUX maintains high image quality under both qualitative and quantitative evaluations, while significantly improving inference speed, even with 20\% of the hierarchy pruned. Our code will be available soon.

Figures (9)

• Figure 1: Overview of the FastFLUX Framework. (1) Importance Estimation: Each block is ranked by estimated importance to determine the pruning order. (2) Block-wise Replacement with Linear Layers (BRLL): Selected blocks have their residual branches replaced with lightweight linear layers, whose parameters are initialized using a closed-form least-squares fit to the original block outputs. (3) Sandwich Training (ST): For non-overlapping cases, we train the replaced block along with its immediate neighbors enhanced by LoRA. For overlapping cases, we search upward or downward to locate the nearest unpruned blocks, apply LoRA to those, and jointly train all intermediate replaced blocks. Input-output data for each training step are re-collected by loading all previously trained sandwich structures. (4) LoRA Strategy: LoRA modules are selectively inserted into either single or double blocks during ST to provide local supervision while keeping the rest of the model frozen. This framework enables efficient and structured pruning with minimal performance degradation.
• Figure 2: Illustration of the Block-wise Replacement with Linear Layers (BRLL) method. Each ResBlock consists of a residual branch and a shortcut branch. The residual branch is replaced with a linear layer, whose parameters are computed using input-output features collected during the forward pass. The shortcut connection is preserved to maintain stable information flow.
• Figure 3: Illustration of the LoRA insertion strategy used in our Sandwich Training (ST). For each block selected for replacement with a linear layer, we identify its closest unpruned neighboring blocks before and after. LoRA modules are inserted into these adjacent blocks to provide localized supervision, enabling efficient and targeted fine-tuning of the replaced block.
• Figure 4: Qualitative comparison between FastFLUX and other mainstream T2I models under the same prompts. FastFLUX consistently generates images of higher visual fidelity. Even at different pruning ratios, the outputs of FastFLUX remain similar to those of the original FLUX.1-dev, while outperforming non-FLUX baselines in overall quality.
• Figure 5: Qualitative comparison between FastFLUX and L1-norm pruned models under similar latency settings. Within each latency group, FastFLUX achieves better alignment with the original FLUX.1-dev output, preserving more texture and visual detail. In contrast, L1-norm pruned models tend to lose fine-grained information, especially under higher pruning ratios.