跳转至

BiGain: Unified Token Compression for Joint Generation and Classification

会议: CVPR 2025
arXiv: 2603.12240
代码: https://github.com/Greenoso/BiGain
领域: 图像生成 / 扩散模型加速
关键词: Token压缩, 扩散模型分类, 频率感知, Token合并, KV下采样

一句话总结

BiGain 首次将扩散模型的 token 压缩重新定义为生成+分类的双目标优化问题，提出拉普拉斯门控 token 合并（L-GTM）和插值-外推 KV 下采样（IE-KVD）两个频率感知算子，在保持生成质量同时显著提升分类准确率（ImageNet-1K 70%合并比下 Acc +7.15%，FID -0.34）。

研究背景与动机

  1. 领域现状：扩散模型不仅是生成主力，还被复用为分类器（逐类去噪似然评分），同一 backbone 支持生成和判别任务。
  2. 现有痛点：现有 token 压缩（ToMe、ToDo）仅优化生成质量（FID），忽略判别能力。几乎不影响生成的压缩可能严重损害分类。
  3. 核心矛盾：生成依赖低/中频全局语义，分类依赖高频细节。朴素压缩恰好移除分类最需要的结构。
  4. 本文要解决什么：设计 training-free、plug-and-play 的 token 压缩，同时保护生成和判别。
  5. 切入角度：频率分离——将中间特征映射到频率感知表示，解耦高频与低/中频。
  6. 核心idea一句话：平衡频谱保留是 token 压缩的可靠设计原则。

方法详解

整体框架

两个互补的 training-free 算子：L-GTM 用于 token 合并，IE-KVD 用于 attention K/V 下采样。兼容 DiT 和 U-Net。

关键设计

  1. 拉普拉斯门控 Token 合并（L-GTM）
  2. 做什么：优先合并频谱平滑 token，保留高频 token
  3. 核心思路：对隐状态做空间拉普拉斯滤波得逐位置频率分数；最低分做 destination，其余为 source，匹配合并

  4. 设计动机：标准 ToMe 不区分频率，高频 token 被优先合并损害分类

  5. 插值-外推 KV 下采样（IE-KVD）

  6. 做什么：对 K/V 做空间下采样，保持 Q 全分辨率
  7. 核心思路：在最近邻和均值池化间做可控插值/外推，分类偏近邻保高频，生成从低频到高频线性变化

  8. 设计动机：保留 Q 维持 token 精细感受野，对分类 MC 评分精度至关重要

  9. 扩散分类器兼容性

  10. 做什么：确保不破坏 diffusion classifier 的估计有效性
  11. 核心思路：两算子 timestep-local 且确定性，所有类别共享同一压缩调度

损失函数 / 训练策略

Training-free，无需训练。推理时直接插入注意力层。分类采用分阶段剪枝加速。

实验关键数据

主实验

方法 FLOPs 减少 Acc@1 (Pets) vs Baseline
Baseline - 81.03 -
ToMe 10% 72.96 -8.07
BiGain_TM 10% 78.38 -2.65
ToDo 14.2% 79.15 -1.88
BiGain_TD 14.2% 79.90 -1.13

消融实验

  • 去掉 Laplacian gating：分类掉点显著
  • IE-KVD alpha=0.9 分类最佳，生成时 0.8到1.2 线性变化最佳
  • DiT 和 U-Net 最优 alpha 不同，需架构适配

关键发现

  • 基线压缩对分类伤害远早于、远大于对生成的伤害
  • 保留 Q 全分辨率是保护分类精度关键
  • 频率平衡在两种架构、四个数据集上一致有效
  • ImageNet-1K 70%合并比：Acc +7.15%, FID -0.34

亮点与洞察

  • 首次系统研究 token 压缩对扩散模型分类能力的影响
  • 频率分离 insight 简洁有力——"看起来好"不等于"分类好"
  • 两个算子都是 training-free 和 plug-and-play
  • 实现帕累托改善而非简单权衡

局限性 / 可改进方向

  • 最优参数与架构/数据集相关
  • 拉普拉斯是手工设计的频率代理
  • 未探索与步数减少方法的组合

相关工作与启发

  • ToMe（Bolya et al. 2023）首次引入 training-free token 合并，BiGain 在其基础上加入频率感知
  • ToDo 提出 token 下采样替代合并，BiGain 的 IE-KVD 在其基础上引入可控插值/外推
  • 扩散分类器（Li et al., Chen et al.）证明扩散 backbone 的判别能力，本文是首个关注压缩对分类影响的工作
  • 频率分离思想可推广到视频生成、3D 等其他 token 压缩场景
  • Diff-Pruning 和 DiP-GO 从模型剪枝角度加速扩散，与本文 token 层面压缩互补
  • ABM（Adaptive Block Merging）变体可进一步加速高分辨率阶段

评分

  • 新颖性: ⭐⭐⭐⭐ 首次提出生成+分类双目标 token 压缩
  • 实验充分度: ⭐⭐⭐⭐ 四个数据集、两种架构全面评估
  • 写作质量: ⭐⭐⭐⭐ 动机清晰，推导连贯
  • 价值: ⭐⭐⭐⭐ 对双用途扩散模型部署有直接指导