TAP: A Token-Adaptive Predictor Framework for Training-Free Diffusion Acceleration

本文提出了无需训练的 TAP 框架，通过利用单层模型探测为每个 token 自适应选择最优预测器，在几乎不增加开销的前提下显著加速了扩散模型的推理过程，同时保持了生成质量。

要点

这篇论文介绍了一种名为 TAP (Token-Adaptive Predictor) 的新方法，它的核心目标是让 AI 画图（生成图像）或生成视频的速度快得多，同时画质几乎不下降，而且不需要重新训练模型。

为了让你轻松理解，我们可以把 AI 生成图像的过程想象成一位画家在画一幅复杂的画。

1. 现在的痛点：画家太累了

目前的 AI 绘画模型（扩散模型）就像一位极其严谨但动作很慢的画家。

• 传统画法：为了画好一张图，画家需要走 50 步（50 个去噪步骤）。每一步，他都要重新审视整张画布，计算每一笔该怎么画，哪怕画布上有些部分（比如蓝天）已经画得很完美了，他也要重新计算一遍。
• 结果：画得很慢，电脑显卡（GPU）累得冒烟，用户等得着急。

2. 以前的加速方法：要么“偷懒”，要么“猜错”

为了加速，以前的方法主要有两种，但都有缺点：

• 方法 A（直接复用/缓存）：就像画家说：“这块蓝天刚才画好了，后面几步我直接照抄刚才的笔触吧。”
  • 缺点：如果画到后面，蓝天里突然要加一只鸟，直接照抄就会出错，画面会模糊或变形。
• 方法 B（全局预测/泰勒展开）：就像画家说：“我根据刚才的笔触，用数学公式猜一下后面几步该怎么画。”
  • 缺点：以前是对所有地方都用同一个公式猜。比如，画“平滑的蓝天”用简单的公式就够了，但画“复杂的毛发”用简单公式就会猜错。如果强行用同一个公式，要么猜不准，要么为了猜准而算得很慢。

3. TAP 的绝招：给每个像素派一个“专属预测员”

TAP 的核心思想是：不要一刀切，要“看人下菜碟”。

它把一张图看作由成千上万个**小像素点（Token）**组成的。TAP 认为：

• 平滑的像素（如蓝天、墙壁）：变化很慢，用简单的预测（低阶公式）就能猜得很准。
• 复杂的像素（如人脸、毛发、文字）：变化剧烈，需要复杂的预测（高阶公式）才能猜准。

TAP 是怎么工作的？（三步走）

第一步：快速“体检” (The Probe)
在每一步画画时，TAP 不会立刻让画家（大模型）算完整个画面。它先让画家只算第一层（就像只看了一眼画布的轮廓），这是一个非常便宜、非常快的操作。

• 比喻：就像你还没开始做整桌菜，先尝一口汤底，就能大概猜出这道菜是咸是淡，需要加多少盐。

第二步：智能“选将” (The Selection)
TAP 手里有一群不同能力的预测员（有的擅长猜简单的，有的擅长猜复杂的）。

• 它利用刚才那个“快速体检”的结果，给每个像素点（Token）算一下：如果用“简单预测员”猜，误差大不大？如果用“复杂预测员”猜，误差大不大？
• 决策：哪个预测员猜得最准，就派谁去负责这个像素点。
  • 蓝天像素 -> 派“简单预测员”（省力气）。
  • 头发像素 -> 派“复杂预测员”（保质量）。

第三步：并行“干活” (The Execution)
一旦选好了，TAP 就直接用预测员的计算结果来代替画家重新计算。因为预测员的计算非常快（只是简单的数学公式），所以整体速度瞬间提升。

4. 为什么 TAP 这么厉害？

1. 不用重新训练 (Training-Free)：
   • 以前的加速方法往往需要让 AI 重新学习很久。TAP 就像给现有的画家配了一个智能助手，不需要画家重新上学，直接就能用。
2. 没有“门槛” (Threshold-Free)：
   • 以前的方法需要人工设定规则（比如“误差超过 0.5 就重新算”）。TAP 不需要人工定规则，它自己比较谁猜得准，自动做决定。
3. 几乎不占内存：
   • 它只记住了很少的信息（第一层的输入和残差），就像只记了个“大概轮廓”，而不是把整张画都存下来。

5. 实际效果如何？

论文里的实验显示：

• 速度：在 FLUX.1 模型上，TAP 能让生成速度提升 6.24 倍（比如原来要 1 分钟，现在只要 10 秒）。
• 画质：在这么快的速度下，画出来的图肉眼几乎看不出区别，甚至有时候因为预测更准，细节反而更好。
• 通用性：无论是画图片还是生成视频，这个方法都管用。

总结

TAP 就像是一个聪明的“画室经理”。
它不再让画家笨拙地重复劳动，也不再让画家用同一种方法去猜所有事情。它通过快速检查，给画布上的每一块区域分配最合适的“猜图助手”。

• 简单的地方，用简单的助手（快）。
• 复杂的地方，用厉害的助手（准）。

最终结果是：画得飞快，而且画得依然漂亮。

技术摘要

1. 研究背景与问题 (Problem)

核心痛点：
扩散模型（Diffusion Models, DMs）虽然在图像和视频生成方面取得了卓越成果，但其推理速度缓慢。主要原因在于扩散采样是一个顺序去噪过程，每一步都需要对庞大的模型进行完整的正向传播（Full-model forward pass），导致计算成本极高。

现有方法的局限性：
现有的加速方法主要分为两类，但都存在明显缺陷：

1. 减少采样步数（如 DDIM, DPM-Solver）： 虽然减少了步数，但在极端加速下会牺牲生成质量，且无法解决单步计算成本高的问题。
2. 特征缓存与预测（Feature Caching & Prediction）：
   • 直接复用（Reuse-based）： 直接复用上一时刻的特征。当时间步跨度增大时，误差迅速累积，导致质量下降。
   • 全局预测（Global Prediction）： 如 TaylorSeer，使用单一的全局预测策略（如固定阶数的泰勒展开）应用于所有 Token 和所有时间步。
   • 关键缺陷： 忽略了Token 级别的异质性（Token-level Heterogeneity）。不同的 Token（如平滑背景 vs. 边缘/运动物体）随时间演变的动态特性不同。简单的背景可能只需低阶预测，而复杂区域需要高阶预测。全局策略要么对简单 Token 过度计算，要么对复杂 Token 预测不准，导致误差累积和严重的质量退化。此外，许多自适应方法依赖人工调整阈值，缺乏鲁棒性。

---

2. 方法论 (Methodology)

作者提出了 TAP (Token-Adaptive Predictor)，这是一个免训练（Training-Free）、探针驱动（Probe-driven） 的框架。其核心思想是：在每一个采样步，为每个 Token 动态选择最合适的预测器。

核心组件：

1. 紧凑的预测器家族 (Taylor Predictor Family)：

   • 构建了一个包含多种候选预测器的集合。
   • 主要基于泰勒展开（Taylor Expansions），通过变化展开阶数（Order, $m$） 和 预测视界/距离（Horizon/Distance, $k_p$） 来覆盖不同的 Token 动态。
   • 低阶预测器对突变更鲁棒，高阶预测器对平滑动态拟合更好。不同的预测距离可以适应不同的收敛半径。

2. 探针机制 (Probe Mechanism)：

   • 洞察： 模型输入层的扰动与输出误差高度相关。
   • 操作： 在每个 $N$ 步窗口的第一步，进行完整的模型计算，并缓存第一层的调制输入（Modulated First-Layer Input, $h_t$） 和 全局残差（Residual, $r_t$）。
   • 作用： 这个第一层的输入作为一个轻量级的“探针”，用于评估后续预测器的质量，而无需进行昂贵的完整模型推理。

3. “探针 - 选择”策略 (Probe-then-Select)：

   • 并行预测： 利用缓存的特征，并行运行所有候选预测器，预测出每个 Token 的调制输入。
   • 代理损失计算 (Proxy Loss)： 计算每个预测器预测的调制输入与真实缓存的调制输入之间的代理损失（如余弦距离）。
   • Token 级选择： 对于每个 Token，选择代理损失最小的预测器（$p^* = \arg\min_p L_p$）。
   • 输出合成： 使用选定的预测器计算该 Token 的残差，合成最终输出，跳过该 Token 的后续深层计算。

流程概览：

1. 第 1 步（窗口开始）： 完整运行模型，缓存 $h_t$ 和 $r_t$。
2. 第 2 到 $N$ 步：
   • 利用缓存数据，并行计算所有候选预测器的输出。
   • 计算每个 Token 的代理损失。
   • 为每个 Token 选择最佳预测器。
   • 生成预测残差并合成输出，跳过深层网络计算。

---

3. 关键贡献 (Key Contributions)

1. Token 自适应预测框架：

   • 首次提出在推理阶段为每个 Token 动态选择最优预测器，而非使用全局固定策略。
   • 证明了仅通过单层（第一层）的完整评估即可作为有效的代理，指导预测器选择，无需额外训练。

2. 多样化的泰勒预测器家族：

   • 识别出不同 Token 在不同泰勒阶数和预测距离下表现最佳。
   • 设计了一个包含不同阶数（0-2 阶）和不同预测距离的紧凑预测器池，能够灵活适应从平滑到突变的各类 Token 动态。

3. 无阈值、免训练的自适应设计：

   • 基于预测器之间的相对代理误差进行选择，完全消除了对人工调整阈值（Thresholds）的依赖，提高了方法的通用性和鲁棒性。
   • 架构灵活，可兼容其他预测方法（如多项式回归、Hermite 插值等）。

4. 极低的开销与显著的加速：

   • 显存开销： 仅增加约 0.1 GB (0.3% 模型大小)，因为只缓存第一层输入和残差，而非所有层特征。
   • 计算开销： 增加的计算量极小（点积和小多项式运算），相比全局预测基线仅增加约 0.015% FLOPs。

---

4. 实验结果 (Results)

作者在多个 SOTA 扩散架构（FLUX.1-dev, Qwen-Image, HunyuanVideo）和任务（文生图、文生视频）上进行了广泛评估。

• FLUX.1-dev (文生图)：
  • 在 6.24 倍加速（$N=8$）下，TAP 实现了 0.99 的 ImageReward 和 31.19 的 CLIP Score，相比原始模型（0.95/30.63）甚至略有提升，且感知质量（PSNR, SSIM）无明显损失。
  • 相比之下，直接复用方法（FORA, TeaCache）和全局预测方法（TaylorSeer）在同等加速比下质量显著下降。
• Qwen-Image：
  • 在 3.57 倍加速下，ImageReward 达到 1.23（优于 TaylorSeer 的 1.18），PSNR 提升约 2.1 dB。
• HunyuanVideo (文生视频)：
  • 在 4.98 倍加速下，VBench 得分达到 65.46，仅比未加速基线下降 1.7%，实现了近乎无损的视觉质量。
• 消融实验：
  • 证明了混合不同阶数和距离的预测器池比单一全局预测器效果更好。
  • 证明了“调制输入”作为探针比原始输入更有效。
  • 证明了该方法兼容非线性预测器（如 Hermite 多项式），进一步提升了性能。

对比基线： TAP 在精度 - 效率前沿（Accuracy-Efficiency Frontier）上显著优于现有的缓存方法（TeaCache, FORA）和预测方法（TaylorSeer, SpeCa）。

---

5. 意义与影响 (Significance)

1. 重新定义了扩散加速范式： 从“全局统一策略”转向“细粒度 Token 自适应策略”，揭示了扩散模型内部 Token 动态的异质性，为高效推理提供了新视角。
2. 真正的免训练与即插即用： 无需微调模型，无需调整超参数阈值，即可直接应用于现有的扩散模型（包括蒸馏模型），极大地降低了部署门槛。
3. 极致的效率与质量平衡： 在实现数倍推理加速的同时，几乎不牺牲生成质量，甚至在某些指标上有所提升，解决了扩散模型落地应用中的速度瓶颈。
4. 低资源消耗： 相比其他需要缓存多层特征的方法，TAP 的显存占用极低，使其在资源受限的设备上更具可行性。

总结： TAP 通过一种巧妙的“探针 - 选择”机制，利用轻量级的第一层信息动态调度计算资源，成功实现了扩散模型的高效、高质量推理，是扩散模型加速领域的一项突破性工作。