BiGain:面向生成与分类联合任务的统一词元压缩
BiGain: Unified Token Compression for Joint Generation and Classification
对做扩散模型加速的人,这篇给出了一个免训练、即插即用的词元压缩方案:用拉普拉斯门控的词元合并保留高频边缘纹理,用最近邻与平均池化之间的可控内插外推来压缩 KV,从而在同等加速比下同时保住生成质量与判别能力,可直接套用到 DiT/U-Net 骨干的推理管线中。
原文依据 a training-free, plug-and-play framework that preserves generation quality while markedly improving classification in accelerated diffusion models
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论扩散模型的加速方法(如词元合并或下采样)通常以降低计算量下的合成质量为优化目标,却往往忽视模型潜在的判别能力。我们以联合目标重新审视词元压缩,提出 BiGain——一个免训练、即插即用的框架,在保持生成质量的同时显著提升加速扩散模型的分类性能。我们的核心洞见是频率分离:将特征空间信号映射到频率感知表示,可将精细细节与全局语义解耦,从而使压缩同时兼顾生成保真度与判别效用。BiGain 以两个频率感知算子体现这一原则:(1) 拉普拉斯门控词元合并,鼓励在频谱平滑的词元之间合并,同时抑制高对比度词元的合并,从而保留边缘与纹理;(2) 内插外推 KV 下采样,通过在最近邻池化与平均池化之间进行可控的内插外推来对键/值下采样,同时保持查询不变,从而无需重新训练即可保持注意力精度。在基于 DiT 与 U-Net 的骨干网络以及 ImageNet-1K、ImageNet-100、Oxford-IIIT Pets 和 COCO-2017 等多个数据集上,我们提出的算子持续改善基于扩散的分类的速度—精度权衡,同时在可比加速下保持、有时甚至提升生成质量。例如,在 ImageNet-1K 上,对 Stable Diffusion 2.0 采用 70% 的词元合并比例时,BiGain 将分类准确率提升 7.15%,同时将生成的弗雷歇初始距离改善 0.34(1.85%)。我们的综合分析表明,平衡的频谱保留——在保留高频细节的同时保留低/中频语义内容——是扩散模型词元压缩的可靠设计准则。据我们所知,BiGain 是首个在加速扩散模型下联合研究并同时推进生成与分类的框架,支持双用途生成系统的低成本部署。
英文原文摘要 2114 字符
Acceleration methods for diffusion models (e.g., token merging or downsampling) typically optimize for synthesis quality under reduced compute, yet they often ignore the model's latent 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 markedly 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 without retraining. Across DiT- and U-Net-based backbones and multiple datasets of ImageNet-1K, ImageNet-100, Oxford-IIIT Pets, and COCO-2017, our proposed operators consistently improve the speed-accuracy trade-off for diffusion-based classification, while maintaining, sometimes even enhancing generation quality under comparable acceleration. For instance, on ImageNet-1K, with 70% token merging ratio on Stable Diffusion 2.0, BiGain increases classification accuracy by 7.15% while also improving FID for generation by 0.34 (1.85%). Our comprehensive analyses indicate that balanced spectral retention, preserving high-frequency detail alongside low/mid-frequency semantic content 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 of dual-purpose generative systems.