CVPR 2026 Highlight 丨 用“几何感知”把扩散 Transformer 采样做成免训练加速器

前言

扩散模型已经成为图像生成和视频生成中最重要的一类基础模型。从 DiT 到 FLUX，再到 HunyuanVideo，越来越多的高质量生成系统都依赖扩散 Transformer 来逐步完成从噪声到图像、从噪声到视频的生成过程。

但一个非常现实的问题是：扩散模型很慢。

这种慢并不是因为单次前向传播特别复杂，而是因为生成过程通常需要很多轮连续去噪。每一步都要调用一次大模型，几十步甚至上百步串行执行下来，推理延迟和计算成本都会变得很高。对于交互式图像生成、视频生成、移动端部署或大规模内容生产来说，这种成本很难忽视。

一个很自然的问题于是出现了：

既然扩散模型慢在采样步数多，那能不能直接少采几步？

直觉上，这似乎是最直接的加速方式。但在实际系统中，粗暴减少采样步数往往会带来明显的质量下降：图像结构变形、语义不稳、细节丢失，视频中还可能出现闪烁和运动不连续。

也就是说，扩散模型加速并不是简单地"少算几步"。真正困难的是：

如何在大步采样时，仍然让生成轨迹沿着模型学到的正确方向演化？

我们的论文 GeoRK2: Geometry-Guided Runge-Kutta Integration for Diffusion Transformer Acceleration 正是围绕这个问题展开。目前文章已被CVPR2026收录为HIGHLIGHT。GeoRK2 是一个免训练、可插拔的扩散 Transformer 加速框架。它把数值积分中的二阶 Runge-Kutta 方法和深度特征空间中的几何结构结合起来，使扩散采样在更少计算下仍然保持稳定的生成质量。

简单来说，GeoRK2 的核心思想是：

扩散采样不是在平直空间里走直线，而是在模型学到的弯曲特征流形上前进。加速采样时，必须尊重这个几何结构。

图片 1 - 多种方法在多个提示词下生成图像的可视化结果。在高加速比下，FORA 和 TaylorSeer 等方法会出现不同程度的图像质量下降，而 GeoRK2 仍能保持更优的性能。

动机：为什么扩散模型一加速就容易"跑偏"？

传统扩散采样器通常把去噪过程看成一个普通的数值积分问题。给定当前噪声状态，模型预测下一步应该往哪里走，然后采样器沿着这个方向推进。

在 DDIM、DPM-Solver 以及很多后续加速方法中，一个隐含假设是：

扩散状态的演化发生在一个近似平直的欧氏空间中。

如果这个假设成立，那么大步采样就可以被理解为在平直空间中做更大的跳跃。只要数值积分公式足够高阶，采样轨迹就应该能够保持稳定。

但扩散 Transformer 的内部特征并不是这样工作的。

我们在 DiT-XL/2 和 FLUX.1-dev 的中间激活上观察到一个非常明显的现象：模型的特征变化高度集中在少数主方向上。换句话说，虽然特征向量位于一个高维空间中，但真正承载主要变化的方向其实很少。

论文中的激活谱分析显示，前 64 个主方向已经能够解释超过 99% 的特征方差。这说明扩散 Transformer 的去噪过程更像是在一个低维、弯曲的特征流形上运动，而不是在整个高维欧氏空间中自由移动。

图片 2 - 主导方向解释的方差。针对 (a) DiT-XL/2 和 (b) FLUX.1-dev 的热力图显示，前 64 个主成分方向解释了超过 99% 的方差。

这带来一个关键问题：

当采样步长变大时，如果仍然使用普通欧氏空间中的直线外推，轨迹就可能逐渐偏离模型真正学到的特征流形。我们把这种现象称为：

manifold drift，流形漂移。

直观来说，模型知道一张图像应该如何从噪声中逐步成形，但这个过程发生在它内部学到的特征几何结构上。普通加速方法只看到"往前走"，却没有看到"应该沿着哪条弯曲的路走"。当步子迈得很大时，轨迹就容易切过流形曲率，而不是贴着流形前进。

图片 3 - (a) GeoRK2 遵循内在流形几何结构，相比忽略曲率的欧氏空间预测方法，能够生成更平滑且更忠实的特征轨迹。 (b) 在大步长条件下，GeoRK2 仍保持稳定，而与几何无关的预测器会因误差快速累积而性能受损。

最终表现出来的就是：

• 图像结构开始扭曲；
• 语义条件逐渐变弱；
• 局部纹理和全局布局不一致；
• 视频生成中出现帧间闪烁和运动不连续。

因此，扩散 Transformer 加速的核心矛盾可以概括为：

我们想减少采样步数，但不能让采样轨迹偏离模型学到的特征流形。

GeoRK2 的出发点，就是把这个被忽略的几何问题显式建模出来。

方法：GeoRK2 = 几何感知预测 +流形校正 + 自适应稳定

图片 4 - GeoRK2 通过将黎曼积分嵌入 Transformer 推理流程，实现具有几何感知能力的扩散采样。（左）预测阶段利用缓存的激活值，在较大时间间隔上外推潜在动态；校正阶段则在局部构建的几何流形上细化轨迹。（右）在每个 Transformer 块内，GeoRK2 引入轻量级的预测-校正模块，用曲率感知的二阶积分替代现有采样器（如 DDIM、DPM-Solver）的默认数值更新，从而在无需重新训练的情况下实现稳定且加速的生成。

GeoRK2 的方法设计可以用一句话概括：

在 RK2 二阶积分的基础上，用模型内部激活估计局部特征几何，并用这个几何结构约束大步采样。

从整体上看，GeoRK2 包含三个核心模块：

1. 几何感知的 RK2 预测；

2. 低秩度量预条件校正；

3. 自适应稳定机制。

下面分别来看。

1. 几何感知 RK2：不是直接外推，而是沿主特征方向外推

Runge-Kutta 方法是经典数值积分中的一类高阶方法。普通的一阶方法只看当前位置的速度，而 RK2 会额外估计一个中点，从而更准确地预测下一步状态。

在扩散采样中，这意味着：

不要只根据当前一步的方向前进，而是先估计中间位置，再用中间位置的方向来决定大步更新。

但 GeoRK2 并不是直接把标准 RK2 套到扩散 Transformer 上。原因在于，高维特征空间中并不是所有方向都同样可靠。很多方向可能只是噪声，或者并不位于模型真正使用的特征流形上。

因此，GeoRK2 在做中点预测时，会先把更新投影到主导特征子空间中：

h_mid = Projection(h_t + step / 2 * v_t)

这里的 Projection 不是任意设计的，而是来自模型中间激活的低秩主方向。也就是说，GeoRK2 会先问模型：

当前这个阶段，哪些特征方向才是真正重要的？

然后只在这些方向上进行大步预测。

这样做的直观效果是：

• 保留模型最有信心的主要变化方向；
• 抑制偏离流形的噪声方向；
• 避免大步采样时轨迹漂到不可靠区域。

可以把它理解成：普通 RK2 是"往前走得更准"，而 GeoRK2 是"沿着模型认为可靠的路往前走"。

2. 低秩几何校正：用激活协方差估计模型内部的局部度量

只做投影还不够。因为扩散模型的去噪过程并不是一条直线，而是在不同噪声阶段经历不同曲率的轨迹。特别是在高噪声向低噪声过渡时，特征空间的几何结构会发生明显变化。

为了解决这个问题，GeoRK2 会从模型的中间激活中构造局部协方差矩阵：

G_t = covariance(H_t) + epsilon * I

这个矩阵可以理解为当前特征空间的一个局部"几何度量"。大特征值方向表示模型在这些方向上变化强、结构敏感，更新时需要更加谨慎；小特征值方向则表示变化较平坦，可以更放心地加速。

但直接使用完整协方差矩阵代价很高。对于扩散 Transformer 的中间层来说，完整矩阵求逆会带来明显计算开销。

GeoRK2 的关键工程设计是：只保留低秩主方向。

由于论文中观察到 top-64 主方向已经解释超过 99% 的方差，GeoRK2 使用截断 SVD 来近似这个几何度量。这样既保留了最重要的特征几何，又把复杂度控制在可接受范围内。

在得到低秩几何度量之后，GeoRK2 会对预测结果做一次 metric-preconditioned correction：

geometry correction = - lambda * G_inverse * prediction_error

直观来说，这一步是在问：

当前预测结果和模型真正希望的去噪方向之间还有多少偏差？这个偏差应该按照特征流形的几何结构如何修正？

这就是 GeoRK2 中的"几何校正"。

它不是重新训练模型，也不是额外学习一个校正网络，而是直接利用预训练模型内部已有的激活统计来完成校正。

为了进一步降低开销，GeoRK2 还使用 Woodbury identity 来高效计算低秩矩阵逆，并且每隔若干步更新一次 metric，而不是每一步都完整重算。

3. 自适应稳定：遇到剧烈阶段切换时及时刹车

扩散采样过程并不是平稳的。

在早期高噪声阶段，模型主要决定全局布局和大致语义；在后期低噪声阶段，模型开始细化纹理、边缘和局部细节。不同阶段之间的过渡往往会带来速度和加速度的突变。

如果加速方法在这些位置仍然激进外推，就可能出现不稳定。

因此，GeoRK2 加入了一个非常简洁的自适应稳定机制。它会监测当前采样轨迹的加速度方差：

ifVar(a_t) suddenly increases:    
    use conservative fallback
else:    
    use geometry-aware RK2 prediction

当检测到状态变化过于剧烈时，GeoRK2 会暂时退回到更保守的两点外推，避免错误被快速放大。最后还会使用 momentum mixing 对输出进行平滑：

h_out = rho * corrected_prediction + (1 - rho) * h_t

这一步的作用不是让方法变复杂，而是让几何积分在真实模型中更加稳健。

论文中也观察到，这个 rollback 机制主要在高加速设置下的关键时间段触发，例如 FLUX 中从概念布局转向细节 refinement 的阶段。它带来的额外开销很小，但可以显著降低发散风险。

工程上是否可落地？

GeoRK2 的一个重要特点是：它不需要重新训练模型，也不需要修改模型结构。

它更像是一个轻量级 PyTorch wrapper，在推理时拦截中间激活，估计局部几何结构，然后替换或增强原本的采样更新。

从复杂度上看，GeoRK2 的额外开销主要来自三部分：

• 主特征方向投影；
• 低秩 metric correction；
• 周期性截断 SVD。

论文中的 profiling 显示，在 DiT-XL/2 上，当截断秩取 64 时：

• projection step 约 0.014 TFLOPs；
• metric-preconditioned inversion 约 0.022 TFLOPs；
• amortized truncated SVD 每步约 0.011 TFLOPs；
• 总额外 FLOPs 约 5.1%；
• 实测每步 wall-clock overhead 约 3.8%。

考虑到 GeoRK2 能够带来 4-5 倍级别的整体采样加速，这个额外开销是很小的。

更重要的是，它的超参数并不敏感。论文中使用同一组配置：

lambda = 0.1
rho = 0.85
beta = 0.9
rank = 64

就可以在 DiT-S/2、DiT-B/2、DiT-XL/2 等模型上稳定工作。这说明 GeoRK2 不是依赖精细调参的特定技巧，而是利用了扩散 Transformer 特征几何中的普遍结构。

实验结果：4-5 倍加速下保持高质量生成

为了验证 GeoRK2 的效果，论文在三个代表性场景上进行了实验：

• ImageNet-256 上的 class-conditional image generation，使用 DiT-XL/2；
• DrawBench 上的 text-to-image generation，使用 FLUX.1-dev；
• VBench 上的 text-to-video generation，使用 HunyuanVideo。

对比方法包括传统采样器 DDIM、DPM++，以及多种扩散 Transformer 加速方法，例如 FORA、TaylorSeer、ToCa、SmoothCache、TeaCache、DBCache 等。

1. ImageNet-256：高加速下仍保持较低 FID

在 ImageNet-256 + DiT-XL/2 上，GeoRK2 在多个加速档位下都取得了较好的速度-质量权衡。

例如：

• GeoRK2(N=2) 将延迟从 8.38s 降到 4.42s，达到 1.95x 加速，同时 FID 为 2.41；
• GeoRK2(N=3) 在 2.70x 加速下取得 FID 2.67；
• GeoRK2(N=8) 在 4.92x 加速下仍保持 FID 3.32。

相比之下，很多方法在类似高加速设置下 FID 会明显上升。这说明在 aggressive acceleration 下，几何感知积分能够更好地保持采样轨迹稳定。

表格 1 - 使用 DiT-XL/2 在 ImageNet-256 上的定量比较。结果为 5 次运行的平均值。Speed表示相对于 DDIM-50 的加速倍数。

2. FLUX.1-dev：复杂文本条件下保持语义一致性

在 FLUX.1-dev + DrawBench 上，GeoRK2 的优势尤其体现在语义一致性上。

在 NFE=50 设置下：

• GeoRK2(N=5) 达到 3.52x latency speedup，ImageReward 为 0.9889，CLIP Score 为 34.963；
• GeoRK2(N=7) 达到 4.06x latency speedup，仍保持 ImageReward 0.9792；
• GeoRK2(N=8) 达到 4.39x latency speedup，CLIP Score 仍保持在 33.476。

这一点很重要。因为文本到图像生成不仅要求图像清晰，还要求模型正确理解 prompt。部分缓存类方法虽然也能加速，但在复杂文本条件下更容易损失语义一致性。

GeoRK2 通过几何约束保持中间特征轨迹稳定，因此在大步采样时更不容易偏离文本条件所对应的生成方向。

表格 2 - 在 FLUX.1-dev 上的比较。结果为 5 次运行的平均值。Speed 表示相对于 50 步参考方法的加速倍数。

3. HunyuanVideo：视频生成中的时序稳定性

视频生成比图像生成更难加速。因为模型不仅要保证单帧质量，还要保证帧与帧之间的运动连续性。

在 HunyuanVideo + VBench 上，GeoRK2(N=8) 达到：

• latency 从 323.89s 降到 69.44s；
• latency speedup 为 4.66x；
• FLOPs speedup 为 6.77x；
• VBench Score 为 80.73。

相比其他加速方法，GeoRK2 在高加速下更好地保持了时间一致性。论文中也指出，一些方法虽然单帧质量尚可，但会出现可见 flickering，而 GeoRK2 的几何度量估计能够更自然地适配时空注意力特征。

表格 3 - 在 HunyuanVideo 上基于 VBench 的比较。结果为 5 次运行的平均值。Speed 表示相对于 50 步基线方法的加速倍数。

消融实验：每个组件都在发挥作用

为了验证 GeoRK2 的设计是否必要，论文还做了系统消融。

在 DiT-XL/2 上：

• 完整 GeoRK2 的 FID 为 2.31；
• 去掉 geometry correction 后，FID 上升到 3.02；
• 去掉 RK2、改用 Euler 后，FID 上升到 2.87；
• 去掉欧氏预测相关设计后，FID 上升到 3.41。

这说明 GeoRK2 的优势不是来自某一个单独 trick，而是来自三个部分的协同：

二阶预测负责更准确地大步前进，几何校正负责把轨迹拉回流形，自适应稳定负责避免阶段切换时发散。

DiT-XL/2 上的受控消融实验（N=3）。FID degradation 表示相对于完整 GeoRK2 的 FID 相对增加幅度。

表格 4 - DiT-XL/2 上的受控消融实验（N=3）。FID degradation 表示相对于完整 GeoRK2 的 FID 相对增加幅度。

论文还分析了截断秩的影响。结果显示，从较低 rank 增加到 64 时质量明显提升，但继续增加到 128 后收益变小。这与前面的谱分析一致：主导特征方向已经覆盖了大部分有效几何信息。

图片 5 - 截断秩对特征重建的影响。散点图比较了在 1/64、1/32、1/16、1/8 比例下的原始特征与低秩特征，结果显示较高秩时与恒等线的对齐更紧密，并在 1/64--1/32 处趋于饱和，这与 FID 饱和以及捕获 99\% 谱能量的现象一致。

总结

GeoRK2 的核心贡献，并不是简单提出一个新的采样器，而是把扩散 Transformer 加速中的一个关键问题讲清楚：

高倍采样加速失败的原因，往往不是数值阶数不够，而是采样轨迹偏离了模型内部学到的弯曲特征流形。

从这个视角出发，GeoRK2 做了三件事情。

第一，从现象上，它指出了 aggressive acceleration 下的 manifold drift 问题：普通欧氏空间外推会让去噪轨迹逐渐偏离模型真实的特征几何。

第二，从方法上，它把二阶 Runge-Kutta 积分和 Riemannian feature geometry 结合起来，用激活协方差估计局部度量，并通过低秩投影和 metric correction 保持采样轨迹稳定。

第三，从工程上，它保持了非常克制的设计：不重新训练、不修改模型结构，只作为 plug-and-play 推理模块接入现有扩散 Transformer，并在 DiT、FLUX、HunyuanVideo 上实现 4-5x 级别加速。

对于正在做图像生成、视频生成或扩散模型推理优化的团队来说，这项工作的启发是：

扩散模型加速不只是减少采样步数，而是要让每一次大步更新都沿着模型学到的几何结构前进。

当我们把特征流形的几何信息纳入采样器设计时，大模型生成就可以在更低成本下保持更稳定的质量。

-End-

作者丨五角场打工王、争气尾流