Blog Tags GitHub
Back to blog
diffusion-model diffusion-transformer efficient-attention image-generation benchmark

ProCache: Constraint-Aware Feature Caching with Selective Computation for Diffusion Transformer Acceleration

ProCache提出约束感知的特征缓存与选择性计算方法，加速扩散Transformer模型的推理过程。

ProCache: 基于约束感知特征缓存与选择性计算的扩散Transformer加速

原文链接: arXiv:2512.17298

代码仓库: GitHub - macovaseas/ProCache

会议: AAAI 2026 (Poster)

一、论文概述

1.1 研究背景

扩散Transformer（Diffusion Transformers, DiTs）在图像和视频生成领域取得了最先进的性能，但其高昂的计算成本严重制约了实时部署。特征缓存（Feature Caching）作为一种无需训练的加速方案，通过利用去噪步骤间的时间冗余来复用已计算的特征，但现有方法存在两个关键局限：

  1. 均匀缓存间隔与非均匀时间动态的错配：现有方法采用固定的均匀缓存间隔（如每N步全量计算），无法适应DiT去噪过程中非均匀的特征演化规律。
  2. 误差累积问题：随着缓存复用间隔增大，过时特征的误差呈指数级衰减，导致严重的质量退化。

1.2 核心贡献

贡献描述
分析洞察揭示了DiT特征在去噪过程中的非均匀演化规律——早期阶段稳定，后期阶段高度动态；误差传播集中在深层网络
约束感知缓存模式搜索提出离线约束采样机制，生成针对模型时间特性定制的非均匀激活调度模式
选择性计算模块在缓存步骤中仅对深层块和高重要性token进行选择性计算，以最小开销缓解误差累积
实验验证在PixArt-alpha和DiT上实现最高1.96x和2.90x加速，质量退化可忽略不计

二、核心思想

2.1 关键观察

特征误差演化 图1：DiT块中相对L1误差在扩散步骤中的演化。误差随步骤递增，深层块（如Block 25-28）的误差幅度显著高于浅层块（如Block 1-4），揭示了误差累积在网络深度上的非均匀性。

输出误差 图3：DiT-XL/2在整个扩散过程中当前步骤与前一步骤的输出L1误差。早期和中期阶段特征差异较小，后期步骤急剧增长，呈近似指数趋势。

ProCache的建立在两个关键观察之上：

  1. 时间维度的非均匀性：DiT输出在去噪早期和中期变化缓慢，后期急剧变化（呈近似指数趋势），均匀缓存策略无法适应这一特性。
  2. 深度维度的非均匀性：误差主要在网络深层累积，浅层表现出较强的时间稳定性。

2.2 设计思路

基于上述观察，ProCache采用两阶段策略：

  • 离线阶段：通过约束采样搜索最优的非均匀缓存模式，在稳定阶段最大化特征复用，在关键阶段确保频繁计算。
  • 在线阶段：在缓存步骤中插入轻量级的部分计算，仅刷新深层块中的关键token，以最小开销维持语义一致性。

三、技术架构

方法流程图 图2：ProCache整体流程。(1) ProCache在三个原则性约束下探索有效缓存模式，并通过轻量级离线采样基于质量指标（如FID）选择最优策略。(2) 然后在连续缓存步骤中插入部分计算，选择性更新深层中的高重要性token。

3.1 问题形式化

设 G\mathcal{G} 为预训练的DiT模型，包含 LL 层：G=g1∘g2∘⋯∘gL\mathcal{G} = g_1 \circ g_2 \circ \cdots \circ g_L。每层包含三个组件：

gl=FlSA∘FlCA∘FlMLPg_l = \mathcal{F}_l^{SA} \circ \mathcal{F}_l^{CA} \circ \mathcal{F}_l^{MLP}

所有操作统一为：F(x)=x+AdaLN∘f(x)\mathcal{F}(\mathbf{x}) = \mathbf{x} + \text{AdaLN} \circ f(\mathbf{x})

缓存模式表示为二进制序列 s=[s1,s2,…,sT]∈{0,1}T\mathbf{s} = [s_1, s_2, \ldots, s_T] \in \{0, 1\}^T，其中：

  • st=1s_t = 1：在步骤 tt 执行计算
  • st=0s_t = 0：从缓存复用特征

3.2 约束感知缓存模式搜索

三个核心约束

约束数学表达含义
预算约束M=∑t=1Tst≤BM = \sum_{t=1}^T s_t \leq B激活步骤总数不超过预定义预算 B≪TB \ll T
单调约束vi+1≤vi,∀i=1,…,M−2v_{i+1} \leq v_i, \forall i = 1, \ldots, M-2复用间隔随时间非递增，后期阶段计算更频繁
有界约束vmin≤vi≤vmaxv^{min} \leq v_i \leq v^{max}每个间隔在可行范围内，防止过长导致误差累积或过短影响效率

其中 vi=ti+1−ti−1v_i = t_{i+1} - t_i - 1 定义为第 ii 个复用间隔，表示两个计算步骤之间连续缓存特征的时间步数。

搜索算法

约束搜索空间定义为：

C={s∈{0,1}T∣s 满足约束(2)(3)(4)}\mathcal{C} = \{\mathbf{s} \in \{0,1\}^T \mid \mathbf{s} \text{ 满足约束(2)(3)(4)}\}

搜索流程（Algorithm A）：

  1. 初始化空候选集 C←∅\mathcal{C} \leftarrow \emptyset
  2. 当 ∣C∣<K|\mathcal{C}| < K 且尝试次数 m<Mm < M 时：
    • 生成随机二进制序列 s∈{0,1}T\mathbf{s} \in \{0,1\}^T
    • 检查预算约束 → 不满足则丢弃
    • 提取激活索引 A\mathcal{A} 并验证有界约束 → 不满足则丢弃
    • 通过所有约束检查则加入候选集
  3. 对候选集在小规模代表性数据集上评估（如FID-1k），选择最优 KK 个模式

计算效率：在AMD Ryzen 7 5800H处理器上，K=5K=5 时搜索仅需约0.001秒；整个搜索过程（含评估）在单GPU上不超过1小时。

3.3 选择性计算模块

块误差分布 图4：DiT-XL/2中不同块在步骤20处的特征相对L1误差。误差主要在深层累积。

时间步选择

在缓存步骤中，选择性计算每隔一个位置进行更新。设 ℓ(t)\ell(t) 为包含 tt 的最大连续零块的起始时间步：

Tcmpt={t∣st=0, t−ℓ(t)+1 为偶数}\mathcal{T}^{cmpt} = \{t \mid s_t = 0, \, t - \ell(t) + 1 \text{ 为偶数}\}

对于连续零块，部分计算注入模式示例：

  • [0,0]→[0,α][0, 0] \rightarrow [0, \alpha]
  • [0,0,0,0]→[0,α,0,α][0, 0, 0, 0] \rightarrow [0, \alpha, 0, \alpha]

其中 α\alpha 表示仅执行深层更新的时间步。

层选择：聚焦深层块

Ucmpt={l∣(L−D)≤l≤L}\mathcal{U}^{cmpt} = \{l \mid (L-D) \leq l \leq L\}

其中 D=r×LD = r \times L，rr 为计算比例。实验表明选择最深的25-50%层即可有效控制误差。

Token选择：优先计算重要token

Token重要性基于注意力模块输出的 ℓ2\ell_2 范数定义：

T(xi)=∥vi∥2T(\mathbf{x}_i) = \|\mathbf{v}_i\|_2

其中 vi=ValueProj(xi)\mathbf{v}_i = \text{ValueProj}(x_i)。选择top-p%p\% 最重要的token：

Icmpt={i∣rank(T(xi))≤p%×N}\mathcal{I}^{cmpt} = \{i \mid \text{rank}(T(\mathbf{x}_i)) \leq p\% \times N\}

注意：在自注意力模块中不采用token选择策略（因为每个token都参与所有其他token的计算），仅在交叉注意力和FFN模块中应用。

更新规则

xi(l)={f(l)(xi(l−1)),if l∈Ucmpt and i∈IcmptCache(l)(xi),otherwise\mathbf{x}_i^{(l)} = \begin{cases} f^{(l)}(\mathbf{x}_i^{(l-1)}), & \text{if } l \in \mathcal{U}^{cmpt} \text{ and } i \in \mathcal{I}^{cmpt} \\ \text{Cache}^{(l)}(\mathbf{x}_i), & \text{otherwise} \end{cases}

额外开销：选择性计算仅引入约3%的额外延迟（在缓存步骤的25%层和7-30% token上计算）。

四、核心创新

4.1 创新点总结

创新点描述优势
约束感知模式搜索通过三个原则性约束（预算、单调、有界）限制搜索空间，结合离线评估选择最优模式避免均匀缓存的时间动态错配，搜索开销可忽略（<0.001s）
深层选择性计算仅在最深25-50%层执行部分计算利用误差在深层累积的特性，以最小开销缓解误差传播
重要性token选择基于注意力value向量的 ℓ2\ell_2 范数选择top-p%p\% token兼容FlashAttention等优化实现，无需访问注意力分数
零块注入策略在连续缓存块中每隔一步插入部分计算防止长时间无更新导致的不可逆误差漂移
完全无训练整个流程无需任何模型微调或重新训练即插即用，部署成本低

4.2 与现有方法的对比

方法缓存策略误差缓解Token级优化训练需求
DeepCache均匀间隔，U-Net低分辨率复用周期性全量更新无无
FORA均匀间隔，注意力特征复用无无无
Δ\Delta-DiT均匀间隔，MLP表示共享无无无
ToCa均匀间隔，token级缓存无基于重要性分数无
L2C动态计算路由训练学习跳过决策层级需要训练
ProCache非均匀约束搜索选择性深层+token计算基于value范数无

五、实验结果

5.1 实验设置

设置项详情
基线模型DiT-XL/2, PixArt-alpha, FLUX.1-dev, FLUX.1-schnell
任务类条件图像生成（ImageNet）、文生图（COCO、PartiPrompts）
分辨率256x256（DiT, PixArt-alpha）, 1024x1024（FLUX）
评估指标FID, sFID, Precision, Recall, Inception Score, CLIP Score, Image Reward
硬件2x NVIDIA RTX 3090 (DiT, PixArt), 1x NVIDIA A800 (FLUX)
实现框架Python 3.9, PyTorch 2.4.0

超参数配置

模型K (采样预算)B (激活预算)p (token比例)vminv^{min}vmaxv^{max}r (计算层比例)
DiT-XL/25177%2575%
PixArt-alpha5730%2350%
FLUX.1-dev51730%2350%
FLUX.1-schnell--30%--50%

5.2 主要结果

DiT-XL/2 类条件生成（ImageNet）

方法延迟(s)FLOPs(T)加速比FID↓sFID↓Precision↑Recall↑IS↑
DDPM-250 steps49.564118.68-2.314.980.820.58243.42
DDIM-50 steps4.54923.741.00x2.434.400.800.59241.25
DDIM-25 steps2.26311.872.00x3.184.740.790.58232.01
DDIM-20 steps1.8629.492.50x3.815.150.780.58221.43
Δ\Delta-DiT (N=3\mathcal{N}=3)2.57216.461.47x3.755.700.770.54207.57
FORA (N=3\mathcal{N}=3)2.1918.592.76x3.886.430.790.54229.02
ToCa (N=3\mathcal{N}=3)2.08710.232.32x3.045.140.790.56230.70
ToCa (N=4\mathcal{N}=4)2.0638.732.72x3.645.550.780.56223.25
ProCache (Ours)1.7258.182.90x2.964.930.800.57232.85

关键发现：ProCache在2.90x加速下，FID为2.96（最优）、sFID为4.93（最优），相比FORA提升约30%，同时加速比提高5%。

PixArt-alpha 文生图（MS-COCO2017）

方法延迟(s)FLOPs(T)加速比FID↓CLIP↑
PixArt-alpha2.14211.181.00x28.1216.29
50% steps1.0445.592.00x37.6215.81
Δ\Delta-DiT1.7247.681.54x28.9116.41
FORA (N=2\mathcal{N}=2)1.5755.661.98x29.6316.40
FORA (N=3\mathcal{N}=3)1.3016.052.79x29.8416.42
ToCa1.4736.331.77x28.0216.43
ProCache1.2155.701.96x27.6616.45

关键发现：ProCache在约2.0x加速下，取得最低FID-30k（27.66）和最高CLIP分数（16.45）。

FLUX 文生图（PartiPrompts）

模型方法延迟(s)FLOPs(T)加速比Image Reward↑
FLUX.1-devEuler-50 steps33.853719.501.00x1.202
68% steps23.022529.261.47x1.200
FORA20.822483.321.51x1.196
ToCa (N=2\mathcal{N}=2)19.882458.061.51x1.202
ProCache18.732415.251.54x1.207
FLUX.1-schnellLCM-4 steps2.882277.881.00x1.133
75% steps2.162208.411.33x1.132
FORA2.365225.601.23x1.129
ToCa (N=2\mathcal{N}=2)1.890181.301.53x1.134
ProCache (N=2\mathcal{N}=2)1.817177.261.56x1.138

5.3 定性结果

FLUX生成样本 图5：1024x1024分辨率下1.56x加速比的图像生成样本。

DiT生成对比 图6：原始采样器（上，1.00x速度）与ProCache（下，2.90x加速）的512x512图像生成定性对比。

5.4 高加速比性能

采样FID分析 图7：DiT-XL/2在ImageNet-1k(a)和ImageNet-50k(b)上的鲁棒性与效率-质量权衡分析。

ProCache在高加速比下仍保持鲁棒性能：质量退化降低56.2%，在超过4.53x加速比的场景下仍表现稳定。

5.5 消融实验

组件消融

配置加速比FID↓sFID↓
Default (B=13B=13)2.93x4.758.43
+ Searched Pattern2.93x3.155.12
+ Selective Computation2.90x3.285.95
ProCache (Full)2.90x2.944.93

发现：

  • 搜索缓存模式在无额外计算开销下显著提升FID和sFID
  • 选择性计算以微小速度损失（2.90x vs 2.93x）带来可观的质量改进

采样预算K的消融

KFID↓IS↑
545.07181.69
1045.02182.50
1544.96181.63

发现：K从5增加到15时FID略有改善，但即使K=5（最低预算）也优于现有方法。采用K=5作为默认值以最大化计算效率。

计算块比例r的消融

rFLOPs(T)FID↓IS↑
50%9.13845.32184.18
65%8.91145.33184.09
75%8.59445.31183.61
90%8.34445.38182.46

发现：75%的块取得最低FID和竞争力的IS。主实验中使用25%（即计算最深75%的块）以最大化加速比。

六、相关工作

6.1 DiT加速方法分类

方法类别代表工作特点局限性
剪枝Structural-Pruning, ToMe, SiTo, TinyDiffusion移除冗余层/token需要额外后训练
量化PostTrain, TemporalDynamic, HQDiT降低数值精度需要校准数据
知识蒸馏SnapFusion, FastSampling压缩模型能力需要教师模型和训练
采样步数减少DDIM, DPM-Solver++, Rectified Flow, DC-Solver减少迭代次数与特征缓存正交
特征缓存DeepCache, FORA, Δ\Delta-DiT, ToCa, PAB复用时间冗余特征均匀间隔，误差累积
ProCache（本文）约束搜索+选择性计算非均匀缓存+轻量修正-

6.2 特征缓存方法详细对比

方法目标架构缓存策略Token优化误差控制
DeepCacheU-Net低分辨率特征复用无周期性更新
Faster DiffusionU-Net跳过编码器计算无无
FORADiT注意力特征复用无无
Δ\Delta-DiTDiTMLP表示共享无无
PABDiT优化注意力头计算无无
L2CDiT动态计算路由层级训练学习
ToCaDiT均匀间隔+token选择基于注意力分数无
ProCacheDiT非均匀约束搜索基于value范数选择性深层计算

七、总结

7.1 主要结论

ProCache是一个完全无训练的动态特征缓存框架，专门针对DiT推理加速设计。其核心创新在于：

  1. 揭示了DiT特征演化的非均匀性：误差在时间和深度两个维度上都呈现非均匀分布，为设计自适应缓存策略提供了理论基础。

  2. 约束感知缓存模式搜索：通过三个原则性约束（预算、单调、有界）将搜索空间限制在合理范围内，结合离线轻量评估，高效发现针对特定模型定制的最优缓存模式。

  3. 选择性计算策略：利用深层误差累积的特性，仅在最深层的高重要性token上执行部分计算，以约3%的额外延迟有效缓解误差传播。

7.2 性能亮点

  • DiT-XL/2：2.90x加速，FID 2.96（最优），sFID 4.93（最优）
  • PixArt-alpha：1.96x加速，FID 27.66（最优），CLIP 16.45（最优）
  • FLUX.1-dev：1.54x加速，Image Reward 1.207（最优）
  • FLUX.1-schnell：1.56x加速，Image Reward 1.138（最优）

7.3 局限性与未来方向

  1. 模型特异性：每个模型需要离线搜索其专属缓存模式，虽然开销小（<1小时），但无法直接跨模型迁移。
  2. 与采样加速的正交性：ProCache专注于特征缓存加速，与采样步数减少方法（如DPM-Solver++）正交，未来可探索二者的联合优化。
  3. 扩展至视频生成：论文主要验证了图像生成场景，视频生成等更长序列场景的适用性有待进一步探索。

八、参考资源

8.1 核心链接

资源链接
arXiv论文https://arxiv.org/abs/2512.17298
HTML版本https://arxiv.org/html/2512.17298v1
PDF下载https://arxiv.org/pdf/2512.17298
GitHub代码https://github.com/macovaseas/ProCache

8.2 关键引用

  • DiT: Peebles, W., & Xie, S. (2023). Scalable diffusion models with transformers. ICCV.
  • PixArt-alpha: Chen, J., et al. (2024). PixArt-α\alpha: Fast training of diffusion transformer for photorealistic text-to-image synthesis. ICLR.
  • FLUX: Black Forest Labs. (2024). FLUX.
  • DeepCache: Ma, X., et al. (2024). DeepCache: Accelerating diffusion models for free. CVPR.
  • FORA: Li, H., et al. (2024). Faster diffusion via temporal attention decomposition. arXiv.
  • ToCa: Wu, X., et al. (2024). ToCa: Token-wise caching for diffusion transformers. arXiv.
  • Δ\Delta-DiT: Yang, S., et al. (2024). Δ\Delta-DiT: Speeding up diffusion transformers with delta feature caching. arXiv.
  • DDIM: Song, J., et al. (2021). Denoising diffusion implicit models. ICLR.
  • DPM-Solver++: Lu, C., et al. (2022). DPM-Solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv.
  • Rectified Flow: Liu, X., et al. (2022). Flow straight and fast: Learning to generate and transfer data with rectified flow. ICLR.

8.3 图片资源

所有关键图片已下载至 docs/figures/proc-cache/ 目录：

文件名对应论文图表说明
feature_error.jpgFigure 1DiT块中相对L1误差在扩散步骤中的演化
method_pipeline.jpgFigure 2ProCache整体流程图
output_error.jpgFigure 3DiT-XL/2输出L1误差随扩散过程的变化
block_error.jpgFigure 4不同块在步骤20处的特征相对L1误差
vis_flux.jpgFigure 5FLUX 1024x1024生成样本
vis_dit.jpgFigure 6DiT 512x512定性对比
sampling_fid.jpgFigure 7鲁棒性与效率-质量权衡分析
tradeoff.jpg-加速比与质量权衡分析

分析文档生成时间: 2026-05-30