加速TarFlow采样：GS-Jacobi迭代法让AI图像生成提速5倍

在2025年5月，TapTap公司的钱震研究员和浙江大学的刘本同学共同发表了一篇题为《使用GS-Jacobi迭代加速TarFlow采样》的研究论文。这项研究发表在arXiv预印本平台（arXiv:2505.12849v1），为当下热门的图像生成模型TarFlow提供了一种显著提升生成速度的解决方案。

图像生成模型已经在各种场景中得到广泛应用。作为其中的佼佼者，TarFlow模型巧妙地将Transformer架构与规范化流（Normalizing Flow）模型相结合，在多个基准测试中取得了最先进的成果。然而，由于其因果注意力机制需要顺序计算，TarFlow的采样过程极其缓慢。就像一个繁忙的生产线，每个工人必须等前一个工人完全完成工作才能开始自己的任务，这种严格的串行工作方式大大降低了效率。

想象你在组装一个复杂的乐高积木模型，如果你必须严格按照说明书一步一步来，不能并行工作，那么整个过程将会非常耗时。TarFlow的采样过程就面临着类似的挑战。然而，研究团队发现，通过一系列创新的优化策略，可以大大加快TarFlow的采样速度。

在这项研究中，作者首先将TarFlow采样阶段的非线性循环神经网络转换为对角化的非线性系统，然后应用了高斯-赛德尔-雅可比（Gauss-Seidel-Jacobi，简称GS-Jacobi）迭代方法进行求解。你可以把这个过程想象成重新组织乐高积木的组装方式，使得多个部件可以同时组装，而不必严格按照原来的顺序。

通过细致的分析，研究人员发现TarFlow模型中的各个模块具有不同的重要性。就像一支足球队中，有些核心球员对比赛结果起着决定性作用，而其他球员则相对贡献较小。具体来说，TarFlow模型中少数几个模块在图像生成任务中扮演着主要角色，而其他模块的贡献相对较小；某些模块对初始值敏感且容易出现数值溢出，而其他模块则相对稳健。

基于这两个特性，研究团队提出了收敛排名指标（Convergence Ranking Metric，CRM）和初始猜测指标（Initial Guessing Metric，IGM）。CRM用于识别一个TarFlow模块是"简单的"（仅需少量迭代即可收敛）还是"困难的"（需要更多迭代）；IGM则用于评估迭代的初始值是否合适。这就像在学习一门新技能时，有些部分你可能一学就会，而有些部分则需要反复练习；同时，一个好的起点会让整个学习过程更加顺利。

研究团队在四个TarFlow模型上进行了实验，结果表明GS-Jacobi采样可以显著提升采样效率，同时保持生成图像的质量。具体来说，在Img128cond模型上实现了4.53倍的加速，在AFHQ模型上实现了5.32倍的加速，在Img64uncond模型上实现了2.96倍的加速，在Img64cond模型上实现了2.51倍的加速，而且所有这些加速都没有降低FID评分或样本质量。

这项研究的代码和检查点可以通过GitHub获取：https://github.com/encoreus/GS-Jacobi_for_TarFlow。

一、从图像生成模型到TarFlow：技术背景

在图像生成领域，已经有多种方法被提出。从像素递归神经网络（PixelRNN）、生成对抗网络（GANs）到扩散概率模型（DDPM）和稳定扩散（Stable Diffusion），各种模型层出不穷。虽然扩散模型似乎主导了这一领域，但规范化流仍然提供了一些独特的优势。

规范化流就像是一个神奇的翻译器，它可以将简单的随机噪声精确地转换为复杂的真实图像。想象你有一团随机的彩色颗粒，通过一系列特殊的变换，这些颗粒可以逐渐排列成一张清晰的猫咪照片。这种方法的优势在于其精确的可逆性，使其能够进行准确的密度估计，实现高效的单步采样，并支持具有可解释性的潜在空间操作。

最初的规范化流模型如NICE引入了加性耦合层。为了增强非线性能力，RealNVP整合了缩放和平移到非体积保持变换中，形成了仿射耦合层。Glow通过引入可逆的1×1卷积进一步改进了图像生成，而Flow++则包含了注意力机制。这些模型的最大优势在于其逆函数是显式的，雅可比矩阵是下三角形的，避免了一般可逆ResNet框架中的复杂计算。然而，过于简单的结构使得这些流模型的非线性能力有限。

为了改进这一点，规范化流被与自回归模型结合起来。IAF开创了基于前置维度条件的维度方向仿射变换，以改进变分推断。MAF利用MADE创建可逆的自回归映射。NAF用每维单调神经网络替代MAF的仿射变换以增强表达能力。T-NAF通过整合单个自回归Transformer增强了NAF，而Block神经自回归流采用了端到端的自回归单调网络设计。

TarFlow作为本文研究对象，提出了一个基于Transformer的架构，结合了一系列技术来训练高性能的规范化流模型，并在多个领域展示了最先进的结果。然而，TarFlow的采样效率受到了关键瓶颈的限制：每个自回归块内的因果注意力结构形成了一个非线性RNN，这迫使在采样过程中进行严格的顺序计算，大大降低了大规模图像生成任务的计算效率。

线性或非线性系统的并行求解是科学计算中的一个重要问题。对于线性系统，Saad建立了雅可比、高斯-赛德尔、连续过度松弛（SOR）和克里洛夫子空间技术等方法。块雅可比迭代利用GPU并行化求解线性/非线性方程。作为一个特殊情况，当方程采取自回归形式时，许多方法被提出来加速自回归计算。例如，Oord提出了概率密度蒸馏，用于将知识从慢速自回归模型转移到更快的计算中。MintNet开发了一种专门的基于牛顿-拉普森的定点迭代方法来加速自回归反演。

二、GS-Jacobi迭代：巧妙解决TarFlow的速度瓶颈

在传统的TarFlow模型中，一个自回归块的前向和反向计算可以用数学公式表示如下：

前向过程：zt = exp(-s(x<t))(xt - u(x<t)) 反向过程：xt = exp(s(x<t))zt + u(x<t))

其中t从1到T，x<t表示t时刻之前的历史数据，s(x<t)和u(x<t)由因果注意力块生成。在前向方向，所有xt都是已知的，所以所有s(x<t)和u(x<t)可以并行计算。但在反向方向，xt只能在x<t被求解后才能顺序计算。这就像是在玩多米诺骨牌，每一块都必须等前一块倒下才能继续，无法同时推倒所有骨牌。据作者测试，这种串行采样方式在单个A800 GPU上生成100张128×128的图像需要约213秒。

为了加速这一过程，研究团队首先将这个过程转化为矩阵形式，然后将其视为一个非线性系统：

前向：Z = Σ^(-1)(X)(X - μ(X)) 反向：X = Σ(X)Z + μ(X)

对于反向过程，可以将目标X*视为非线性系统g(X) = Σ^(-1)(X)Z + μ(X)的不动点，然后使用非线性雅可比迭代求解：

X^(k+1) = Σ(X^(k))Z + μ(X^(k))

这就像是在解一个复杂的方程，我们不断地进行猜测和修正，直到找到一个足够接近正确答案的解。研究团队提出了一个命题来解释这种非线性系统下雅可比模式迭代的收敛性和误差传播情况。

这种迭代方法涉及两个关键组件：初始值X^(0)和最大迭代次数。正如实验中所示，不同的初始化策略会导致不同的收敛效果，不良的策略甚至会导致模型崩溃。此外，不同的块以不同的速度收敛，有些块快速收敛，而其他块则较慢，这表明我们应该为不同的块采用不同的迭代策略。为了解决这两个问题，作者分别提出了初始猜测指标（IGM）和收敛排名指标（CRM）。

三、巧妙的初始值选择：初始猜测指标（IGM）

常见的初始化选择是将X^(0) = Z，也就是前一个块的输出，基于TarFlow"逐渐"转换图像的直觉。正如实验中所示，从噪声到图像的变化在大多数步骤中是渐进的，Z位于X*的邻域内，这可能是一个很好的初始猜测。

然而，在实践中，研究团队发现在Img64模型的Block0中，将所有X^(0) = Z会导致数值崩溃。这就像是在烹饪一道精致的菜肴时，选择了错误的起点，导致整个烹饪过程都无法挽回。

另一种可行的猜测是X^(0) = Z0 = [z1, 0, ..., 0]'，因为像素值范围从-1到1，中心在0。一个自然的策略是比较Z和Z0，选择较好的一个。由于最严重的膨胀发生在前几次迭代，作者定义了以下"初始猜测指标"：

IGM(X^(0)) = ||Σ(X^(0))Z + μ(X^(0)) - X*||2

用于测量初始值与目标之间的粗略距离，X^(0)从{Z, Z0}中选择。在论文的附录中，不同的范数显示了类似的结果，谱范数稍好一些，因此作者在本文中使用了谱范数。

IGM可以被视为一种模型属性，一旦训练完成就确定了。因此，不需要在每次采样时重复计算，而是可以使用以下步骤：

1. 从训练集中选择一批图像，将其补丁化为大小(B, T, C)，即X*； 2. 将X*通过TarFlow块前向传递，得到Z = Σ^(-1)(X*)(X* - μ(X*))； 3. 使用{Z, Z0}两种初始值计算残差Σ(X^(0))Z + μ(X^(0)) - X*； 4. 计算B维度上残差的均值，计算(T, C)矩阵的范数。

四、智能迭代控制：收敛排名指标（CRM）

当使用雅可比迭代采样时，尽管所有块最终都会收敛，但有些块只需很少的迭代k就能得到不错的解，而其他块则需要接近T-1的迭代。如实验所示，Img128cond的Block6、AFHQ的Block7、Img64cond的Block0和Img64uncond的Block6相比其他块表现更差。为了度量这种差异，作者提出了以下收敛排名指标：

CRM = ||Σ^(-1)(X)X||2 * ||Ws||2 + ||Wu||2

其中Ws和Wu是s(x<t)和u(x<t)的投影输出层的权重矩阵。Ws度量方差的变化；Wu度量均值，Σ^(-1)(X)X度量非体积保持性。

这个指标并不严格测量收敛率，而只代表TarFlow块之间的相对收敛排名，因此被称为排名指标。在实验中，不同的矩阵范数在相对排名上表现相似，因此作者使用了谱范数。

通过CRM，可以知道一个块是能够快速收敛还是缓慢收敛，从而大致确定雅可比迭代的次数。在表格中显示CRM值较大的块在实验中收敛较慢。在实践中，尽管TarFlow中只有很少的块收敛缓慢，但这严重影响了雅可比迭代方法的速度和有效性：对于"困难"的块，较少的迭代会导致生成质量不佳，而更多的迭代虽然改进了模型但同时失去了速度优势。

五、模块化高斯-赛德尔-雅可比迭代：最佳解决方案

对于(B, T, C)张量，雅可比迭代一次更新所有T个单元，而传统的"For"循环一次更新1个单元，串行运行T-1次。自然地，一种中间方法是在一次迭代中并行（使用雅可比）更新一组单元，然后串行到另一组，这就是所谓的高斯-赛德尔-雅可比迭代。

设X := {xt}^T_{t=1}，{Gg}^G_{g=1}是时间步索引1:T的非递减分段，Xg := {xt|t ∈ Gg}^G_{g=1}，X:g := ∪^g_{i=1} Xi，以及{Z, zt}、{Σ, σt}、{μ, ut}的类似定义。

所有关于雅可比模式迭代的分析都适用于GS-雅可比采样的模块。作者指出，GS-雅可比可以有效地改进具有大CRM值的块的解决方案：

1. 由于初始猜测值导致的数值溢出概率大大降低。误差矩阵的大小更小，因此误差累加也减少； 2. 每个子雅可比的收敛将被加速，因为更接近后面的模块将有一个更准确的初始值； 3. 适当的GS-雅可比策略（选择Gg和最大雅可比迭代次数）可以实现既准确又快速的解决方案。

研究团队将困难块分成8个相等的模块并应用GS-雅可比迭代。结果显示，纯雅可比迭代对于困难块需要50到150次才能收敛，而GS-雅可比方法将这个数字降低到大约30次，而且通常只有模块1遇到更困难的收敛轨迹。

因此，一个合适的策略可以利用这种模块化迭代方法。理想情况下，应该为每个GS-雅可比模块计算IGM和CRM，以判断它是否困难。然后，为每个模块分配更多迭代给大CRM，反之亦然。这可以看作是一种自适应策略。

在实践中，等大小分段和相同的雅可比次数通常就足够了。策略可以用[Stack-GS-J-Else]格式表示：Stack表示应该分段的困难块；GS表示相等大小分段的数量，长度为T//GS；J表示每个模块的最大雅可比次数；Else表示其他具有小CRM的块的最大雅可比次数。

要确定哪些块需要堆叠，可以一个接一个地选择具有大CRM的块，直到剩余集合中没有主导块。根据实验结果，作者堆叠了Img128cond中的Block6、AFHQ中的Block7、Img64uncond中的Block0&6以及Img64cond中的Block0&7。

六、实验结果：速度与质量的完美平衡

研究团队训练了四个由Zhai等人给出的模型：用于条件ImageNet 128×128的TARFLOW [4-1024-8-8-N(0, 0.15^2)]；用于AFHQ 256×256的TARFLOW [8-768-8-8-N(0, 0.07^2)]；用于无条件ImageNet 64×64的TARFLOW [2-768-8-8-N(0, 0.05^2)]；用于条件ImageNet 64×64的TARFLOW [4-1024-8-8-N(0, 0.05^2)]。前三个模型的T=1024，最后一个的T=256，所有四个模型都有8个TarFlow块。为方便起见，这些模型分别被称为Img128cond、AFHQ、Img64uncond和Img64cond。

首先，研究团队使用128个训练图像计算了四个模型的IGM。结果显示，Img128cond和AFHQ在两种初始化方法之间没有显著差异，而Img64cond和Img64uncond在将所有块的X^(0)设为Z时会崩溃。这在表格中很明显，因为使用X^(0)=Z时，Img64cond和Img64uncond中Block0的IGM值异常高，而设置X^(0)=Z0可以缓解这一问题。

IGM与迭代过程中可能出现的最大值高度相关。研究发现，Img64cond和Img64uncond对初始值更加敏感。这可能是因为低分辨率图像更容易在像素之间发生突变，导致注意力层参数出现巨大波动。在实践中，GS-雅可比分段可以大大改善数值溢出的问题，因此简单地根据IGM初始化Z或Z0就足够了。

接下来，研究团队计算了四个模型的CRM。结果与实验观察一致，遵循简单规则：CRM越大，收敛所需的雅可比次数越多，反之亦然。

一个重要的特性是，TarFlow模型中只有很少的块具有相对较大的CRM。这可能是因为TarFlow或其他基于规范化流的生成模型是过度确定的，这意味着参数量相对于生成能力是冗余的，许多块并没有显著修改图像，只是精心调整。从实验可视化结果可以看出，许多中间块没有明显的变化，这为GS-雅可比加速提供了可能性。

最后，研究团队调整了超参数cfg（无分类器引导）、lr（去噪学习率）、attntemp（注意力温度），使用"For"迭代采样了50,000张图像，以恢复Zhai等人的FID结果。将其视为目标FID，然后保持超参数一致，使用不同的GS-雅可比策略进行采样，记录FID、相对误差（%）、运行时间（100秒）和加速率。相对误差小于1%的速率被加粗显示。所有采样都在8个具有80G内存的A800 GPU上进行。

在Img128cond上，保持困难块Block6的"For"迭代，对其他块使用少量纯雅可比就足以获得良好的FID，如[6-1024-1-10]，速度提升3.67倍。然后，研究团队固定Block6的总雅可比次数为128和256，尝试不同的[GS-J]对。结果发现，简单策略如[6-8-32-10]可以实现相对误差<1%的结果和惊人的速度提升。

AFHQ的采样结果类似。由于这两个模型都只有一个困难块，加速率表现相似。对于Img64模型，情况有所不同。在Img64uncond中，由于堆叠了Block0和6，加速率不如单一困难块模型高，但仍能提速约3倍。

在Img64uncond中，由于两个困难块的CRM没有绝对差距，研究团队将它们视为相等。对于Img64cond，研究团队首先将Block0和Block7都堆叠到原始"For"循环中，得到2.31的速率。然后，保持Block0不变，尝试不同的Block7策略，速率可以提高到2.42。根据实验结果，研究团队注意到Block0的行为比其他任何块都更困难，因此将Block0分成更多模块，获得了2.51倍的速度提升。

基于所有实验，研究团队得出结论：具有主导CRM的块越少，补丁化后的时间步越长，GS-雅可比采样能够实现的加速就越明显。这与直觉一致。GS-雅可比通过并行迭代一批方程来实现加速，避免了"For"循环中重复串行更新kv缓存。

七、结论与未来展望

本研究全面优化了TarFlow模型的采样过程。通过识别TarFlow块中的非均匀转换模式并提出IGM和CRM，研究团队有效地解决了初始值选择和收敛率差异的问题。GS-雅可比迭代的引入及其深入的误差传播分析为TarFlow采样提供了实用且高效的解决方案。

多个TarFlow模型上的实验结果表明了所提出方法的优越性。GS-雅可比采样在不降低样本质量的情况下，在Img128cond上实现了4.53倍的加速，在AFHQ上实现了5.32倍的加速，在Img64uncond上实现了2.96倍的加速，在Img64cond上实现了2.51倍的加速，这对TarFlow模型的应用具有重要意义。

然而，仍有一些方面可以进一步改进。计算CRM的强假设需要更多的理论验证。此外，当前确定GS-雅可比参数的方法相对简单，未来有望开发出更智能和自适应的策略。

这项研究不仅提供了一种加速TarFlow采样的实用方法，还为理解和优化复杂生成模型的采样过程提供了有价值的见解。随着人工智能图像生成技术的不断发展，这种优化技术将对提高模型的实用性和效率产生重要影响。