GeoSPRINT：用于扩散轨迹推理的几何冗余感知步骤剪枝

# GeoSPRINT：面向扩散轨迹推理的几何冗余感知步进剪枝  
来源：https://arxiv.org/html/2609.02160  
Arpita Joshi  
机构：斯克里普斯研究所  
机构地址：美国加州圣地亚哥拉霍亚托里松北路10550号  
机构邮箱：[email protected] / [email protected]  

###### 摘要  
扩散模型虽然能生成高质量样本，但推理过程代价高昂，因为采样需要多次顺序的神经网络函数评估（NFE）。现有的加速方法要么采用固定的跳步策略，要么基于局部数值误差自适应调整步长，或者需要额外的训练。我们提出GeoSPRINT（轨迹推理中的几何步进剪枝），这是一种无需训练的框架，能够从去噪轨迹的几何结构中构建非均匀采样调度。GeoSPRINT通过潜在空间中的超平面性测试检测几何冗余步骤，并通过QR分解高效实现，同时将得到的冗余度分布转换为采样调度，为轨迹的高曲率区域分配更多步骤。此外，我们引入了轨迹投影分数αtraj（α_{\mathrm{traj}}），这是一种基于残差方差的度量，用于量化轨迹的直线度，并作为重整化流质量的无模型诊断工具。在CIFAR-10（32×32）、LSUN Church（256×256）和Stable Diffusion v1.5（512×512 潜在空间）上，GeoSPRINT在相同NFE预算下均持续优于均匀DDIM（去噪扩散隐式模型）调度。在CIFAR-10上，GeoSPRINT在49-89个NFE下将FID（Fréchet Inception Distance）提升了0.7-1.1，且在NFE≥30时超越了DPM-Solver++（即使使用一阶DDIM求解器）。在LSUN Church上，它在52步时将FID从1.48降至1.26；在Stable Diffusion v1.5上，其FID提升高达1.93。这些结果表明，轨迹几何为推理步骤分配提供了有效的全局信号，调度质量可以显著提升扩散采样的效率，而无需重新训练。  

## 1 引言  
扩散模型[1]、[2]及其通过基于分数的随机微分方程实现的连续时间泛化，已成为图像[3]、视频[4]、音频[5]和分子设计[6]等领域的主流生成建模范式。尽管其生成质量卓越，但反向采样过程的迭代特性仍然是一个计算瓶颈：生成单个样本通常需要数十到数千次的顺序神经网络评估。学术界已开发出丰富的加速技术。无需训练的方法包括DDIM[7]，它将扩散采样重新解释为确定性常微分方程（ODE），并支持均匀跳步；DPM-Solver[8]和DPM-Solver++[9]采用指数积分器，并基于局部截断误差自适应调整步长；UniPC[10]则统一了预测-校正方案。需要训练的方法包括渐进式蒸馏[11]，它通过迭代将所需步数减半；一致性模型[12]、[13]学习直接将轨迹上的任意点映射到终点；以及重整化流[14]，它通过拉直传输路径以减少积分步数。尽管取得了这些进展，仍存在一个差距：现有的无需训练方法缺乏一个基于几何冗余的*全局*、轨迹级步骤重要性判据。DDIM均匀跳步，对轨迹几何一无所知。DPM-Solver基于单个时间步的*局部*误差估计自适应调整步长。自适应非均匀时间步采样[15]侧重于*训练*而非推理。最近的SDM框架[16]分析局部ODE刚性，但仍然是逐点操作的。这些方法都没有提出全局性问题：*在整个轨迹中，哪些步骤真正为从噪声到数据的路径贡献了新的几何信息？* 我们认为，这恰恰是几何数据实例缩减[17]所解决的问题。该研究证明，有序数据点序列可以通过几何冗余性测试（2D中的共线性，3D中的共面性）大幅缩减，同时保留结构的本质形状和方差，这可以通过基于被移除点协方差矩阵特征值的投影分数[18]来量化。关键洞察是，有序数据集中的许多连续点近似地位于相同的线性或平面子空间内，因此可以在不损失信息的情况下被移除。去噪轨迹{zT, zT−1, ..., z0}正是这样一个有序序列。在分数函数变化缓慢的区域，连续的潜在状态描绘出几乎线性的路径——这些步骤在几何上是冗余的。在变化迅速的区域（例如模型解析精细结构时），轨迹急剧弯曲，每一步都携带新的方向信息。这一类比启发了GeoSPRINT。贡献：(1) 我们将几何实例缩减从2-3维PCA空间推广到任意d维潜在空间，通过超平面性测试实现，每步复杂度为O(d⋅k²)。(2) 我们提出了一种新颖的调度构建方法，将对数信噪比（log-SNR）间隔与GeoSPRINT的轨迹曲率密度相结合，生成优于均匀时间步（DDIM）和均匀对数信噪比间隔的非均匀调度。(3) 我们引入了*轨迹投影分数* αtraj（α_{\mathrm{traj}}），这是一种量化轨迹非直线度的残差方差度量；我们证明了对于完美直线的轨迹该分数为零，并展示了从DDPM到DDIM下降了450倍，确立了其作为流匹配质量无需训练诊断工具的地位。(4) 我们证明，使用GeoSPRINT调度的一阶DDIM求解器在CIFAR-10上（NFE≥30）超越了二阶DPM-Solver++，表明*在哪里*步进比*如何*步进更重要。(5) 我们在CIFAR-10（32×32）、LSUN Church（256×256）和Stable Diffusion v1.5（512×512 潜在空间）上提供了基准测试，结果在不同分辨率、模型架构和条件类型下均显示出一致的改进。  

## 2 背景与相关工作  
### 2.1 扩散模型与采样问题  
基于分数的SDE框架。Song等人[2]在连续时间SDE框架下统一了扩散模型。前向过程为dz=f(z,t)dt+g(t)dw，其中f和g定义了漂移和扩散系数。逆向过程为：dz=[f(z,t)−g(t)²∇zlogpt(z)]dt+g(t)d¯w， (1) 其中∇zlogpt(z)是由神经网络sθ(z,t)近似的分数函数（¯w为标准维纳过程）。等效的确定性表述，即概率流ODE为：dzdt=f(z,t)−12g(t)²∇zlogpt(z)。 (2) 采样简化为求解从t=T到t=0的这个ODE，需要离散化为N步，每步涉及一次神经网络函数评估（NFE）。流匹配：Lipman等人[19]和Liu等人[14]提出学习一个速度场vθ(z,t)，通过dz/dt=vθ(z,t)将源分布传输到目标。重整化流[14]通过迭代回流（reflow）拉直这些路径。Wang等人[20]表明，严格的直线性并非必需——一阶ODE一致性就足够了。  

### 2.2 现有加速方法  
均匀跳步。DDIM[7]实现了确定性采样，可均匀跳过时间步，将1000步减少至约50步，但对轨迹几何一无所知。高阶ODE求解器。DPM-Solver[8]推导了扩散ODE线性部分的精确解，并对非线性残差应用指数积分器，在10-20步内实现高质量生成。DPM-Solver++[9]将其扩展到引导采样。两者都使用*局部*截断误差进行自适应步长调整，采用启发式调度（logSNR-均匀、时间-均匀、时间-二次）。需要训练的方法。渐进式蒸馏[11]训练学生模型使用N/2步匹配教师模型，迭代减半直到4步生成。一致性模型[12]、[13]学习直接的轨迹-终点映射。一致性流匹配[21]结合了一致性训练和流匹配。所有方法都需要为每个目标NFE预算进行额外训练。自适应时间步方法。Kim等人[15]提出自适应非均匀时间步采样以加速*训练*。SDM框架[16]分析局部PF-ODE刚性，以在*推理*过程中设置自适应求解器阶数和步长。两者都基于单个时间步的局部属性操作。  

### 2.3 几何实例缩减  
Joshi和Haspel[17]提出了一种数据实例缩减算法，该算法在投影到PCA空间的有序序列上操作。该算法通过共线性（2D）和共面性（3D）测试来检验几何冗余性，移除那些与邻居近似位于相同线性子空间中的点。信息损失通过*投影分数*[18]量化：α=∑j=1pλj(S)∑j=1pλj(Z)， (3) 其中λj(S)和λj(Z)分别是被移除子集S和完整数据集Z的协方差矩阵计算得到的第j个特征值，p是保留的PCA维数。该算法在多样化数据集（分子轨迹、图像、机器学习基准测试）上实现了40-87%的点缩减率，投影分数低至10⁻⁵。关键限制：原始算法在2-3维PCA空间中操作。扩散潜在空间通常是32-512维。第3节直接解决了这个问题。  

## 3 方法：GeoSPRINT  
### 3.1 将几何冗余推广到高维  
设{z0, z1, ..., zN}为ℝd（从t=T到t=0的去噪轨迹上的潜在状态有序序列）中的有序点序列。我们旨在识别并移除近似包含在其邻居张成的仿射子空间中的点。  

###### 定义1（超平面性测试）。  
给定一个由k个*已保留*点组成的窗口W={w1, w2, ..., wk}（其中k≤d）和一个候选点zi，定义：  
MW=[w2−w1 | w3−w1 | −⋯ | wk−w1]∈ℝd×(k−1)。 (4)  
点zi到W张成的仿射子空间的残差距离为：  
ri=‖(I−MWMW+)(zi−w1)‖2， (5)  
其中MW+是Moore-Penrose伪逆。如果ri<τ，则zi近似位于该仿射张成空间中，是一个可移除的候选点。注意，定义1使用最近保留的（不一定连续的）点作为参考窗口，这与算法2的因果滑动窗口逻辑一致。通过QR分解高效计算。我们不直接计算伪逆，而是计算薄QR分解MW=QWRW，从而得到ri=‖(I−QWQW⊤)(zi−w1)‖2。计算MW∈ℝd×(k−1)的QR分解的代价是O(d⋅k²)运算；由于k很小（通常k=2），实际上为O(d)。自适应窗口大小。我们提出一个渐进式层次结构，模仿原始算法中2D→3D的进展：测试按渐进式层次应用：级别1（k=2，共线性）捕获长直段；级别2（k=3，共面性）测试四点是否共面，捕获平面扫描；级别ℓ（k=ℓ）测试是否位于(ℓ−1)维仿射子空间中。这些测试顺序应用：首先是级别1，然后在幸存点上应用级别2，以此类推。在实践中，我们发现k=2（共线性）对于所有测试模型都足够（见附录C.3）。阈值确定。遵循Joshi和Haspel[17]，我们通过轨迹投影分数（第3.3节）引导的二分搜索确定τ，调整直至被移除点的投影分数低于目标（例如10⁻³），同时最大化移除的步数。  

### 3.2 应用于去噪轨迹  
参考轨迹生成。给定一个预训练的扩散或流模型，我们首先使用完整的N步调度生成B条参考轨迹。每条轨迹{zT(b), ..., z0(b)}提供了一个ℝd中的有序序列。  
逐轨迹剪枝。对于每条参考轨迹，我们应用超平面性测试（算法2）来识别并移除几何冗余的时间步，产生一个每条轨迹的保留集合R(b)⊆{t0, t1, ..., tT}。  
通用调度聚合。为了获得一个无需每样本开销的固定调度，我们跨参考轨迹进行聚合：w(t)=1B∑b=1B1[t∈R(b)]。 w(t)=1B∑b=1B