arXiv:2603.12240cs.CVcs.LG2026-03

让加速后的扩散模型既生成清晰图像,又能准确分类。

BiGain: Unified Token Compression for Joint Generation and Classification

  • 通过频率分离设计,区分细节与语义,实现更合理的压缩。
  • 在ImageNet-1K上提速70%时,分类准确率提升7.15%,生成质量更好。
  • 无需训练、可直接接入现有模型,适合部署优化场景。

扩散模型的加速方法(如令牌合并或下采样)通常只关注生成质量的保持,却忽视判别能力。本文提出BiGain,一种无需训练、即插即用的统一框架,在加速扩散模型的同时兼顾生成与分类性能。核心思想是频率分离:将特征空间信号映射为频率感知表示,解耦细粒度细节与全局语义,使压缩同时保留生成保真度与判别效用。BiGain包含两个频率感知算子:(1) 拉普拉斯门控令牌合并,鼓励频谱平滑的令牌合并,抑制高对比度令牌合并,从而保留边缘与纹理;(2) 插值-外推键值下采样,通过可控插值与平均池化间过渡下采样键值,保持查询完整,维持注意力精度。在基于DiT和U-Net的骨干网络及ImageNet-1K、ImageNet-100、Oxford-IIIT Pets、COCO-2017数据集上,该框架持续改善加速下的速度-准确率权衡,生成质量不降反升。例如在ImageNet-1K上,对Stable Diffusion 2.0进行70%令牌合并时,分类准确率提升7.15%,FID改善0.34(相对提升1.85%)。分析表明,平衡保留高频细节与低/中频语义是扩散模型令牌压缩的可靠设计准则。据我们所知,BiGain是首个在加速扩散模型下联合优化生成与分类性能的框架,支持更低成本部署。

原文摘要 · Abstract (English)

Acceleration methods for diffusion models (e.g., token merging or downsampling) typically optimize synthesis quality under reduced compute, yet often ignore discriminative capacity. We revisit token compression with a joint objective and present BiGain, a training-free, plug-and-play framework that preserves generation quality while improving classification in accelerated diffusion models. Our key insight is frequency separation: mapping feature-space signals into a frequency-aware representation disentangles fine detail from global semantics, enabling compression that respects both generative fidelity and discriminative utility. BiGain reflects this principle with two frequency-aware operators: (1) Laplacian-gated token merging, which encourages merges among spectrally smooth tokens while discouraging merges of high-contrast tokens, thereby retaining edges and textures; and (2) Interpolate-Extrapolate KV Downsampling, which downsamples keys/values via a controllable interextrapolation between nearest and average pooling while keeping queries intact, thereby conserving attention precision. Across DiT- and U-Net-based backbones and ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO-2017, our operators consistently improve the speed-accuracy trade-off for diffusion-based classification, while maintaining or enhancing generation quality under comparable acceleration. For instance, on ImageNet-1K, with 70% token merging on Stable Diffusion 2.0, BiGain increases classification accuracy by 7.15% while improving FID by 0.34 (1.85%). Our analyses indicate that balanced spectral retention, preserving high-frequency detail and low/mid-frequency semantics, is a reliable design rule for token compression in diffusion models. To our knowledge, BiGain is the first framework to jointly study and advance both generation and classification under accelerated diffusion, supporting lower-cost deployment.

扩散模型令牌压缩多任务优化高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。