用轻量线性层替换残差块，加速FLUX模型推理
FastFLUX: Pruning FLUX with Block-wise Replacement and Sandwich Training
- 用线性层替代复杂残差分支，保留快捷连接保稳定
- 20%层级剪枝后仍保持高图像质量，推理速度显著提升
- 局部微调策略缓解结构替换导致的性能下降，适合部署场景
文本到图像生成模型如扩散Transformer（DiTs）在表达能力上取得突破，以FLUX为代表。然而其庞大的参数量导致推理缓慢、内存占用高且难以部署。现有加速方法（如单步蒸馏和注意力剪枝）常伴随严重性能下降及高昂训练成本。为此，我们提出FastFLUX——一种面向FLUX的架构级剪枝框架。核心是块级替换线性层（BRLL）方法，将残差块中复杂的残差分支替换为轻量线性层，同时保留原始快捷连接以确保稳定性。此外，引入夹心训练（ST）策略，利用LoRA监督邻近块进行局部微调，缓解结构替换带来的性能损失。实验表明，即使剪枝20%层级，FastFLUX仍能保持高质量图像生成能力，并显著提升推理速度。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。