Sora出圈，背后DiT也火了！作者NYU谢赛宁官宣全新升级版SiT

新智元报道

编辑：好困

【新智元导读】扩散Transformer（DiT）重磅升级——可扩展插值Transformer（SiT）！虽然核心架构没变，但在质量、处理速度和使用灵活性方面都实现了显著提升。

最近，OpenAI的视频生成模型Sora实火。

根据纽约大学计算机系助理教授谢赛宁的分析，Sora是基于自己和William Peebles共同提出的DiT框架设计而成。论文被ICCV 2023接收。

随后，William Peebles加入了OpenAI，领导了开发Sora的技术团队。

论文地址：https://arxiv.org/abs/2212.09748

与此同时，随着大家对DiT的关注越来越高，谢赛宁团队还官宣了重磅升级之后的SiT架构——Scalable Interpolant Transformer，即可扩展插值Tranformer。

虽然基于相同的骨干，但SiT实现了更好的质量、速度和灵活性。

对此，谢赛宁表示，SiT突破了传统扩散技术的限制，从而可以利用插值法开拓更加广泛的设计领域！

论文地址：https://arxiv.org/abs/2401.08740

可扩展插值Transformer（SiT），是在扩散Transformer（DiT）技术的基础上发展而来的。

SiT框架提供了一种比传统扩散模型更加灵活的方法，使得我们能够更灵活地连接两个不同的数据分布。

这种方法让我们能够从多个角度审视和优化基于动态传输的生成模型的设计，包括选择离散还是连续的时间来学习、确定模型的学习目标、挑选用于连接分布的插值方法，以及使用确定性或概率性的数据采样器。

结果显示，SiT在条件ImageNet 256x256的基准测试中，使用相同的网络架构、参数数量和计算量，实现了在不同模型规模上超越DiT的性能。

SiT通过调整不同的扩散系数（与学习过程独立），达到了2.06的FID-50K分数，展现了其卓越的图像生成能力。

流与扩散

近年来，一种全新的生成模型逐渐成为焦点，其核心思想是将纯粹的噪声ε转化为具有特定分布p(x)的数据x*。

目前，基于扩散的模型是这种转换最常用的框架，而随机插值和基于流的模型则是这个领域的新成员。

新模型将αt和σt的值限制在时间区间[0,1]内，且α0=σ1=1，α1=σ0=0。如此一来，xt就可以精确地在x*和ε之间插值。

这种方法的优点是，让插值函数插值函数的选择更加灵活，因为它们不再受制于前向SDE。

此外，这些模型还使用了更简单的概率流ODE进行推理：

其中，速度v(Xt ,t)是通过流量匹配目标进行估计的：

简单来说，这可以被看作是预测一个粒子在t时刻从某个ε开始移动的速度。

为了更好地理解这些模型，团队将关键组成部分总结在了下表中：

研究证明，在相同的αt和σt条件下，扩散和基于流的方法实际上遵循着相同的时间发展过程。也就是说，无论是基于流的ODE还是基于扩散的ODE和SDE，其对应的pt(x)都是一致的。

论文中，团队不仅揭示了这些模型组件之间的数学等价性，还探讨了它们对模型性能的影响。

并更进一步地解决了其中的「?」，即在缺少明确前向SDE的情况下，基于流的方法也可以通过反向时间SDE进行有效采样。

可扩展插值Tranformer

从上表可以看出，设计这些模型时需要考虑以下四个方面：

- 时间空间（Timespace）：离散或连续的时间间隔；

- 模型预测（Model Prediction）：Ls或Lv的目标；

- 插值器（Interpolant）：αt和σt的选择；

- 采样器（Sampler）：常微分方程（ODE）或随机微分方程（SDE）。

时间空间

首先，将模型从基于离散时间的去噪方式改进为基于连续时间的评分方式后，性能有了小幅提升。

模型预测

团队发现，速度模型与评分模型之间存在一个随时间变化的权重函数：

其中， 。

将这种线性关系应用到Lv中，可以得到：

这一发现与之前的观察一致，即不同的扩散模型预测结果相当于对普通去噪目标进行了不同时间函数的加权。

为了解决数值稳定性的问题，团队采用了与SBDM相同的方法，将训练和采样的范围限制在[ε,1]区间内。

因此，较大的λε能够弥补Ls梯度消失的问题，但这也使得优化Lv变得更加困难。

插值器

团队主要研究了三种不同的插值方法：

下面的例子将展示这些插值方法对简单的一维数据分布产生的影响。

从标准高斯分布开始的一段时间内，VP插值方法在两种情况下都引起了最显著的变化，而GVP和线性插值方法的变化则相对平缓。

直观上来说，VP插值方法的这种急剧变化增加了速度场的Lipschitz常量，从而使其学习变得更加困难。

从标准高斯分布到定义在±1处的伯努利分布

从标准高斯分布到具有两个峰值，分别位于-1和2的高斯混合分布

当从简单的示例转向更复杂的图像生成任务时，可以再次看到不同插值方法之间存在显著的性能差异：

这种差异的一个可能原因如下图所示：当从SBDM-VP改为GVP或线性插值时，路径的长度（也就是传输成本）有所减少。

采样器

在SBDM设置下，速度的反向时间SDE可按以下方式构建：

其中，用g(t)来表示SBDM模型中的扩散系数。

根据之前的讨论，也可以根据g^2(t)=2λtσt这一关系，为GVP和线性插值方法构建类似的SDE模型：

随后，团队进一步提出，扩散系数g(t)可以独立于学习过程进行调整。

实际上，任何非负的函数w(t)（不必是单调的）都可以作为扩散系数。

因此，可以将反向时间的SDE模型进行推广：

除了SBDM系数外，团队还尝试使用w(t)=σt（消除数据分布附近的分数奇异性）和w(t)=sin^2(πt)作为扩散系数，并探讨了它们对速度模型或评分模型的影响。

结果显示，选择最优的扩散系数取决于插值方法和目标，而在实验中，这也很大程度上取决于模型的大小。

经验表明，对于SiT-XL模型，最佳选择是使用线性插值方法和连续时间速度模型，并使用w(t)=σt作为扩散系数的SDE进行采样。

最后，在不同的计算资源下，ODE和SDE积分器的性能可能会有所不同。

如实验所示，相比于SDE，ODE在更少的函数评估次数下能够更快地收敛，而在更大的计算资源下，SDE能够实现更低的FID分数。

无分类器引导

在速度模型中，团队使用了「无分类器引导」的方法。可以看到，这种方法在DiT中所带来的显著性能提升，同样适用于SiT。

参考资料：

https://twitter.com/sainingxie/status/1747863734884745431

https://huggingface.co/papers/2401.08740