Which Tokens to Merge? Diffusion Dynamics for Efficient Image Generation¶

会议: NeurIPS 2026
arXiv: 2610.22370
OpenReview: https://openreview.net/forum?id=cDWOMIJrbJ
代码: https://github.com/SeungJuCha/DaTMe
领域: 图像生成
关键词: Token Merging, Diffusion Dynamics, Tweedie Estimation, Diffusion Transformer, Efficient Inference

一句话总结¶

提出无需训练的动力学感知 Token 合并框架（DaTMe），基于 Tweedie 去噪估计构建时序演化与空间高频双轴显著图，自适应分离约束锚点与安全合并源，解决 CFG 剪枝忽略非文本区域与时变特性的痛点，在 PixArt-$\alpha$ 与 FLUX 上实现全图细节保真与显著加速。

研究背景与动机¶

扩散模型与 Diffusion Transformer（DiT）在图像与视频生成领域展现出强大的表征能力与可扩展性，但在高分辨率多步迭代去噪过程中，海量空间 Token 在每个时间步反复经过自注意力计算，造成了巨大的计算开销与延迟瓶颈。作为一种免训练的推断加速手段，Token 合并（Token Merging）通过聚类相似 Token、融合表征来削减后续计算量。早期方案如 ToMeSD 主要依赖特征相似度进行无约束合并，这往往导致前景或关键内容 Token 被错误吸收，造成生成图像产生严重结构畸变。

为保护关键信息，近期代表性工作引入无分类器引导（CFG）作为 Token 重要性判据，通过条件预测与无条件预测之间的扰动差异来筛选文本引导的关键 Token 并予以保留。然而，以 CFG 引导为重要性指标存在两大固有盲区：其一，CFG 无法捕捉时间维度的动态演化特性。扩散生成是一个由粗到细的渐进过程，前期更新集中于前景轮廓，中后期细节注入则向非提示词区域持续扩散，而 CFG 响应在整个去噪历程中高度静态化，始终锚定在文本显式提到的前景上，导致在去噪中后期正在积极更新的背景和外围 Token 被过早合并；其二，CFG 对提示词无关的结构信息缺乏敏感度。图像中的背景物体、物体边缘、高频纹理等支撑性几何结构并未在文本中显式提及，在 CFG 信号下始终获得极低评分，极易沦为合并源，从而造成背景物体边界溶解、结构崩塌。

这两个盲区的本质矛盾在于：Token 的真实重要性不应由静态文本相关性单维决定，而应取决于扩散动力学本身——即该 Token 在当前步是否仍处于未收敛的演化阶段，以及其在像素或空间流形上承载了多少结构与高频信息。核心 idea：本文提出动力学感知 Token 合并（DaTMe），基于 Tweedie 去噪估计在纯净图像空间构建时序演化与空间高频双轴重要性图，将 Token 自适应划分为受保护的约束锚点与低动态均质的安全合并源，并结合网格均匀补全与局部最近邻匹配，实现全图结构无损的免训练高效生成。

方法详解¶

整体框架¶

DaTMe 面向扩散 Transformer 及 Flow Matching 架构设计，在去噪迭代的指定注意力层内执行免训练的动态 Token 合并。整个流水线首先通过 Tweedie 公式将带噪状态映射至预测干净图像空间，分别沿时间维度计算步间残差以获得时序演化图，沿空间维度计算下采样-上采样高斯金字塔残差以获得空间结构图；接着在两张图的交集上完成动力学感知角色划分；随后通过网格空间重采样补全目标合并预算；最后执行局部受限的空间邻近匹配合并，完成特征融合与加速。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["当前步输入状态 xt 与网络预测 (噪声 / 速度)"] --> Tw["Tweedie 纯净图像空间估计 x̂0(t)"]
    Tw --> S1["时序演化分析：纯净空间步间残差<br/>M_temp = |x̂0(t) - x̂0(t-1)|"]
    Tw --> S2["空间结构滤波：多尺度金字塔残差<br/>M_spat = |x̂0(t) - Up(Down(x̂0(t)))|"]
    S1 & S2 --> RA["动力学感知角色划分<br/>Top 交集设为约束集 S_con，Bot 交集设为保守源集 S_src"]
    RA --> SC["空间网格均匀采样补全<br/>划分 P 个网格，从余量中均衡补齐至目标合并数 n_tgt"]
    SC --> LM["局部受限空间最近邻匹配<br/>排除约束集，将最终合并源映射至空间欧氏最近目标"]
    LM --> Out["输出压缩后的 Token 特征 F'(t) 送入后续层"]

关键设计¶

1. 基于纯净空间的 Tweedie 时序与空间双轴重要性度量：剥离预设噪声计划干扰并兼顾几何结构

针对传统 CFG 判据无法反映全图时变收敛性与结构重要性的问题，DaTMe 构建了无需额外反向传播与注意力重计算的双轴度量体系。在时序演化维度上，直接对原始 Score 差分 $|s(t) - s(t-1)|$ 会面临严重数学缺陷：通过代数分解可知，Score 斜率可严格解耦为两项： $$s(t) - s(t-1) = \frac{\epsilon_\theta(t-1) - \epsilon_\theta(t)}{\sigma_{t-1}} + \epsilon_\theta(t)\left(\frac{1}{\sigma_{t-1}} - \frac{1}{\sigma_t}\right)$$ 其中第一项是反映模型预测实质变化的差异项，第二项则是单纯由全局预设噪声计划决定的尺度项。在去噪后期 $\sigma_t \to 0$ 时，$\left|\frac{1}{\sigma_{t-1}} - \frac{1}{\sigma_t}\right|$ 急剧膨胀，尺度项彻底掩盖了代表语义演化的真实差异项。为了彻底消除这一虚假膨胀，DaTMe 利用 Tweedie 公式将状态投射至预测纯净空间 $\hat{x}_0(t) = \frac{x_t - \sigma_t \epsilon_\theta(t)}{\alpha_t}$，其步间残差严格推导为： $$\hat{x}_0(t) - \hat{x}_0(t-1) = \frac{\sigma_{t-1}}{\alpha_{t-1}}\big(\epsilon_\theta(t-1) - \epsilon_\theta(t)\big)$$ 该差分项完全剔除了纯噪声计划项，且因子 $\frac{\sigma_{t-1}}{\alpha_{t-1}}$ 在去噪末期自然衰减至 0，从而构筑了极具物理意义的时序动态图 $\mathcal{M}_{\mathrm{temp}}^{(t)} = |\hat{x}_0(t) - \hat{x}_0(t-1)|$。在 Flow Matching 架构下，速度预测 $v_\theta(t)$ 在欧拉步下的 Tweedie 残差同样具备类似收敛衰减形式 $(t-\Delta t)(v_\theta(t-\Delta t) - v_\theta(t))$。在空间结构维度上，为捕获非文本描述的边界、边缘与局部纹理，DaTMe 对纯净预测图施加空间尺度因子为 $s$ 的下采样与双线性插值上采样，构建高通滤波空间图： $$\mathcal{M}_{\mathrm{spat}}^{(t)} = \left|\hat{x}_0(t) - \mathrm{UpSample}_{\times s}\big(\mathrm{DownSample}_{\times s}(\hat{x}_0(t))\big)\right|$$ 两轴联合作用，在全图范围内清晰定位出哪些 Token 仍在激烈更新、哪些 Token 承载了关键物理轮廓，为后续决策提供了严密的动力学支撑。

2. 动力学感知角色划分：严密解耦约束锚点与安全合并源

在获得双轴显著图后，给定目标压缩率 $r$ 和总 Token 数量 $N$，目标合并 Token 数为 $n_{\mathrm{tgt}} = \lfloor rN \rfloor$。为避免硬编码阈值在不同去噪阶段带来的泛化失效，DaTMe 采用相对分位数双轴交集策略自适应分配 Token 角色。一方面，只有在时序演化与空间几何双维度上均处于最底层的 Token，才被判定为既已收敛又高度平滑的冗余信息，构成保守合并源集合： $$\mathcal{S}_{\mathrm{src}} = \mathrm{BotFrac}_r\big(\mathcal{M}_{\mathrm{temp}}^{(t)}\big) \cap \mathrm{BotFrac}_r\big(\mathcal{M}_{\mathrm{spat}}^{(t)}\big)$$ 另一方面，为了防止核心物理边缘和仍在发生剧烈构图变化的 Token 被周边平坦区域同化溶解，定义关键约束锚点集合： $$\mathcal{S}_{\mathrm{con}} = \mathrm{TopFrac}_{1-r}\big(\mathcal{M}_{\mathrm{temp}}^{(t)}\big) \cap \mathrm{TopFrac}_{1-r}\big(\mathcal{M}_{\mathrm{spat}}^{(t)}\big)$$ 约束集 $\mathcal{S}_{\mathrm{con}}$ 在整个合并流程中被完全冻结，既不作为合并源被消除，也不接受其他 Token 融合，从而以结构锚点的形式稳固了物体的轮廓骨架与非文本背景的完整性。

3. 空间网格均匀采样补全与局部受限邻域匹配：防范空间坍缩与跨域失真

由于保守源集 $\mathcal{S}_{\mathrm{src}}$ 采用严格双重交集筛选，其实际 Token 数 $n_{\mathrm{src}} = |\mathcal{S}_{\mathrm{src}}|$ 通常严格小于预算需求 $n_{\mathrm{tgt}}$。若直接按单一数值全局贪心补全，极易导致某一局域背景（如天空或大块纯色底色）内的 Token 被过度抽空，引发局域特征塌陷。DaTMe 设计了网格空间均匀补全机制：将未分配的候选集合 $\mathcal{U} = \mathcal{S}_{\mathrm{tot}} \setminus (\mathcal{S}_{\mathrm{con}} \cup \mathcal{S}_{\mathrm{src}})$ 映射回二维空间平面，划分为 $P$ 个互不重叠的均匀网格单元，从每个网格单元内等额均匀采样 $\frac{n_{\mathrm{tgt}} - n_{\mathrm{src}}}{P}$ 个 Token 作为补充源 $\mathcal{S}_{\mathrm{src}}^{\mathrm{add}}$，组装出最终合并源集合 $\mathcal{S}_{\mathrm{src}}^* = \mathcal{S}_{\mathrm{src}} \cup \mathcal{S}_{\mathrm{src}}^{\mathrm{add}}$。在路由阶段，鉴于 DiT 架构中自注意力机制缺少天然的空间归纳偏置，若单纯基于特征向量余弦相似度进行全局二分图匹配，极易将空间相隔遥远却纹理相近的孤立 Token 错误合并，破坏空间几何拓扑。因此，DaTMe 将未分配 Token 划为目标池 $\mathcal{D} = \mathcal{S}_{\mathrm{tot}} \setminus (\mathcal{S}_{\mathrm{con}} \cup \mathcal{S}_{\mathrm{src}}^*)$，并严格依据 Token 的二维物理坐标欧氏距离执行空间邻域单射： $$d(j) = \arg\min_{k \in \mathcal{D}} \|\mathbf{p}_j - \mathbf{p}_k\|_2, \quad \forall j \in \mathcal{S}_{\mathrm{src}}^*$$ 源 Token 与其最近邻目标加权合并后与保留的约束集 $\mathcal{S}_{\mathrm{con}}$ 合流，形成最终尺寸为 $N - n_{\mathrm{tgt}}$ 的精简 Token 特征集，从根源上规避了长程乱序合并导致的拓扑畸变。

实验关键数据¶

主实验¶

主实验在 PixArt-$\alpha$（基于 COCO 2014 验证集采样 10,000 个图文对）和 FLUX.1-dev 上全面展开，与代表性免训练加速方法 ToMeSD、CFG-based Token Merging 以及 ToMA 进行不同压缩率下的对比评测。测试硬件基于单张 RTX PRO 6000 Blackwell 显卡。

表 1：PixArt-$\alpha$ 上的文本对齐度与效率评测（主对比表）

模型与方法	压缩比 $r$	CLIP-T $\uparrow$	HPSv2 $\uparrow$	FID $\downarrow$	延迟 (s) $\downarrow$	计算量 (TFLOPs) $\downarrow$
PixArt-$\alpha$ (Base)	0.0	31.2485	29.2780	42.1001	1.9417	183.45
ToMeSD	0.3	30.5735	24.1699	61.5229	1.6917	172.51
CFG-based	0.3	30.8801	27.7375	44.6070	1.7458	172.51
DaTMe (Ours)	0.3	31.1375	28.5591	42.4474	1.8030	175.63
ToMeSD	0.5	29.5178	21.3334	94.8901	1.5918	165.20
CFG-based	0.5	30.4645	25.8345	52.2912	1.6445	165.20
DaTMe (Ours)	0.5	31.0142	27.5174	43.4155	1.6884	170.41
ToMeSD	0.7	28.2805	18.9797	123.0139	1.5356	157.90
CFG-based	0.7	29.4543	22.6156	70.7062	1.5761	157.90
DaTMe (Ours)	0.7	30.6259	25.2393	47.0519	1.6426	165.19

表 2：PixArt-$\alpha$ 上的图像结构与细节保真度评测

方法	压缩比 $r$	SSIM $\uparrow$	PSNR $\uparrow$	LPIPS $\downarrow$	美学评分 (Aes.) $\uparrow$
ToMeSD	0.3	0.7306	18.9748	0.2433	5.9906
CFG-based	0.3	0.7137	17.8157	0.2369	6.1521
DaTMe (Ours)	0.3	0.8112	20.5136	0.1501	6.2014
ToMeSD	0.5	0.6737	17.7161	0.3151	5.7739
CFG-based	0.5	0.6331	16.2768	0.3254	6.1135
DaTMe (Ours)	0.5	0.7300	18.4277	0.2258	6.1695
ToMeSD	0.7	0.6301	16.8537	0.3682	5.5809
CFG-based	0.7	0.5695	15.2003	0.4129	5.9559
DaTMe (Ours)	0.7	0.6581	16.9048	0.3078	6.0544

表 3：FLUX.1-dev 上的跨架构泛化评测（Flow Matching 架构）

方法	压缩比 $r$	CLIP-T $\uparrow$	HPSv2 $\uparrow$	FID $\downarrow$	SSIM $\uparrow$	PSNR $\uparrow$	LPIPS $\downarrow$	延迟 (s) $\downarrow$
FLUX Base	0.0	31.2274	29.6927	30.6141	1.0000	$\infty$	0.0000	14.601
ToMeSD$^\dagger$	0.5	31.1852	28.3595	30.8491	0.6166	16.9059	0.2932	13.278
ToMA	0.5	31.2036	27.7566	30.9883	0.7009	20.0375	0.2305	13.288
DaTMe (Ours)	0.5	31.2621	28.9266	31.3423	0.7087	20.2021	0.2176	13.271

消融实验¶

表 4：双轴重要性度量与 Score Slope 替换消融（PixArt-$\alpha$, $r=0.3$）

重要性信号配置	CLIP-T $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	FID $\downarrow$	说明
仅 CFG 判据 [43]	30.8801	0.7137	17.8157	44.6070	缺乏非文本与动态感知，结构掉点显著
仅时序轴 $\mathcal{M}_{\mathrm{temp}}$	31.0355	0.8109	20.5021	42.5269	捕捉动态收敛但遗漏静止高频边界
仅空间轴 $\mathcal{M}_{\mathrm{spat}}$	31.1318	0.8110	20.5029	42.5736	保护边缘高频但缺乏时间阶段适应性
双轴 (Score Slope 替代)	31.1166	0.8028	20.3944	42.9554	晚期受尺度项虚假膨胀干扰，精度退化
完整双轴 Tweedie (DaTMe)	31.1375	0.8112	20.5136	42.4474	理论最优解，指标与保真度全面领先

表 5：约束锚点集合 $\mathcal{S}_{\mathrm{con}}$ 选择策略消融（PixArt-$\alpha$, $r=0.3$）

约束集构建准则	CLIP-T $\uparrow$	SSIM $\uparrow$	PSNR $\uparrow$	LPIPS $\downarrow$	FID $\downarrow$	结论分析
无约束集 (None)	31.1419	0.8037	20.3736	0.1601	43.7692	缺少结构保护，FID 恶化近 1.3 点
$\mathcal{M}_{\mathrm{temp}}\uparrow \cap \mathcal{M}_{\mathrm{spat}}\downarrow$	31.1250	0.8005	20.3459	0.1641	43.8062	锁定平坦区域，未能保护物理轮廓
$\mathcal{M}_{\mathrm{temp}}\downarrow \cap \mathcal{M}_{\mathrm{spat}}\uparrow$	31.1372	0.8076	20.5027	0.1569	43.6605	保护静态边缘，但忽视了演化未定区域
$\mathcal{M}_{\mathrm{temp}}\uparrow \cap \mathcal{M}_{\mathrm{spat}}\uparrow$ (Ours)	31.1375	0.8112	20.5136	0.1501	42.4474	兼顾演化活跃性与结构复杂度，效果最佳

关键发现¶

极度耐受高压缩率：在高压缩率（$r=0.7$）下，基线方法 FID 剧烈崩塌（ToMeSD 达 123.01，CFG 达 70.71），而 DaTMe 依然坚挺在 47.05，SSIM 高达 0.6581，证明动力学感知有效杜绝了灾难性特征损毁。
Tweedie 去噪估计彻底修复 Score 斜率缺陷：理论与实证均表明，原始 Score Slope 在去噪中后期（归一化时间 $\tau \in [0.4, 0.6]$）与 Pseudo-GT 变化幅度的相关性暴跌，而 Tweedie 纯净图像空间差分剥离了噪声计划的 $1/\sigma_t$ 放大效应，全阶段维持极低的时序自相关性与高度真实的相关性。
多模型架构通用性：DaTMe 无需重新微调，不仅在标准 DiT（PixArt-$\alpha$）上大幅领先，在不具备 CFG 机制的现代 Flow Matching 大模型 FLUX.1-dev 以及经典的卷积 U-Net（SD 1.5、SDXL）上均表现出卓越的零样本加速与细节保持能力。

亮点与洞察¶

解耦噪声计划与模型演化：通过严格的数学推导揭示了原始 Score 斜率在去噪后期被尺度项严重污染的根本机理，证明 Tweedie 纯净空间步间残差能够纯净隔离模型更新并自然衰减，是极其优雅且具普适性的理论设计。
提示词盲区破局：跳出了传统以文本 Cross-Attention 或 CFG 强度作为唯一重要性先验的局限思维，从生成动力学与图像自身高频金字塔的本质出发，有效守护了非提示词背景与未建模物体的轮廓完整性。
可复用的免训练加速即插即用管线：空间网格分块均衡采样与欧氏空间局部近邻单射的设计，为缺乏卷积归纳偏置的现代 Transformer 视觉生成模型提供了防塌陷的标准工程范式。

局限与展望¶

极端复杂场景下的信息丢失：在包含极其密集的小物体群或极端交错背景的复杂构图中，若施加 $r \ge 0.7$ 的超高压缩率，即使是双轴度量也可能无法在有限 Token 预算下完整保留所有细微纹理。
层级动力学异质性未解耦：当前方法主要均匀施加于主干 Transformer 的固定层段，尚未结合 DiT 各层对于布局、纹理、全局色调等不同表征层级的分工进行层自适应（Layer-adaptive）的动态压缩比调节。
未来改进方向：可进一步探索结合层级敏感度分析的动态按层剪枝，并将 Tweedie 动力学度量延伸至视频扩散模型（如 Sora 类架构）的时空 3D Token 合并中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [揭示 Score Slope 晚期尺度项污染机理并创新性引入 Tweedie 纯净空间差分作为时序动力学判据，视角深刻。]
实验充分度: ⭐⭐⭐⭐⭐ [横跨 PixArt-$\alpha$、FLUX.1-dev、SD 1.5、SDXL 四大模型，覆盖完整指标对比、严密消融及人工偏好实验。]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严丝合缝，动力学演化与几何空间配准逻辑清晰有力。]
价值: ⭐⭐⭐⭐⭐ [即插即用、无需微调，对推动高分辨率大模型实时交互推理部署具有极高的实用与学术价值。]

模型与方法	压缩比 \(r\)	CLIP-T \(\uparrow\)	HPSv2 \(\uparrow\)	FID \(\downarrow\)	延迟 (s) \(\downarrow\)	计算量 (TFLOPs) \(\downarrow\)
PixArt-\(\alpha\) (Base)	0.0	31.2485	29.2780	42.1001	1.9417	183.45
ToMeSD	0.3	30.5735	24.1699	61.5229	1.6917	172.51
CFG-based	0.3	30.8801	27.7375	44.6070	1.7458	172.51
DaTMe (Ours)	0.3	31.1375	28.5591	42.4474	1.8030	175.63
ToMeSD	0.5	29.5178	21.3334	94.8901	1.5918	165.20
CFG-based	0.5	30.4645	25.8345	52.2912	1.6445	165.20
DaTMe (Ours)	0.5	31.0142	27.5174	43.4155	1.6884	170.41
ToMeSD	0.7	28.2805	18.9797	123.0139	1.5356	157.90
CFG-based	0.7	29.4543	22.6156	70.7062	1.5761	157.90
DaTMe (Ours)	0.7	30.6259	25.2393	47.0519	1.6426	165.19