灵活视图变换器：为扩散模型设计的革新图像生成架构

在自然界中，图像的分辨率是无限的，而现有的图像生成模型在跨任意分辨率泛化方面存在困难。虽然扩散变换器（DiT）在特定分辨率范围内表现出色，但在处理不同分辨率的图像时却力不从心。为了克服这一限制，来自上海人工智能实验室的研究团队及其合作者提出了灵活视图变换器（Flexible Vision Transformer，简称FiT），这是一种专为生成任意分辨率和纵横比的图像而设计的变换器架构。与传统方法不同，FiT将图像视为动态大小的令牌序列，这种视角使得FiT在训练和推理阶段都能灵活适应不同的纵横比，从而促进分辨率的泛化，并消除了由图像裁剪引入的偏见。

FiT-XL/2模型在不同分辨率下生成的图像样本

方法

论文首先介绍了一些关键的预备概念和技术，特别强调了1-D Rotary Positional Embedding（1-D旋转位置嵌入，简称RoPE）的概念。RoPE是一种新颖的位置编码方式，它通过将绝对位置编码和相对位置编码结合在一起，为大型语言模型（LLMs）提供了一定程度的序列长度外推能力。

RoPE通过在复数向量空间中对关键向量和查询向量应用偏置，实现了位置编码。这种方法利用了一个旋转频率矩阵，该矩阵定义了一系列的旋转角度，这些角度与位置编码的维度有关。在实际空间中，这种旋转矩阵可以被看作是一个由余弦和正弦函数组成的矩阵，它能够根据位置信息调整关键向量和查询向量。

其次注意力分数的计算方式也很重要，即如何结合使用RoPE的查询向量和关键向量来得到注意力分数。这涉及到利用RoPE编码后的查询向量和关键向量之间的复数实部内积，以确定模型在不同位置编码下的注意力权重。

最后是NTK-aware Interpolation（NTK感知插值）和YaRN（Yet another RoPE extensioN，另一种RoPE扩展）插值技术。这些技术主要用于处理大型语言模型在面对测试时上下文长度超过训练时最大长度的情况。通过调整RoPE的旋转基，NTK-aware Interpolation能够适应更长的上下文。而YaRN技术则引入了一个额外的超参数来调整旋转频率，从而更有效地扩展上下文窗口。

为了处理不同分辨率的图像，FiT采用了一个灵活的预处理步骤。与传统的将所有图像调整到固定分辨率的方法不同，FiT在预处理阶段避免了裁剪或过度缩放图像。FiT仅将高分辨率图像调整到一个预设的最大分辨率限制，确保图像的原始纵横比不变，从而避免了引入不必要的失真或信息损失。

接下来，FiT的训练流程包括将图像编码为潜在代码，然后将这些潜在代码分割成固定大小的潜在令牌。由于不同图像的潜在令牌序列长度可能不同，FiT通过填充令牌将所有序列统一到一个最大长度，这样就可以将不同长度的序列打包到同一个批次中进行并行处理。在训练过程中，模型只计算去噪输出令牌的损失，而忽略填充令牌。

在推理阶段，FiT首先根据目标图像的分辨率定义一个位置图，然后从高斯分布中采样噪声令牌作为输入。通过一系列的去噪步骤，模型逐步改善这些噪声令牌，最终根据位置图将去噪后的令牌重塑并解码成最终的图像。

FiT的训练和推理流程还特别考虑了GPU硬件的特性，确保数据以统一形状的批次进行处理，以优化并行计算效率。这种灵活的训练和推理流程使得FiT能够生成各种分辨率和纵横比的高质量图像，同时保持了图像的细节和完整性。

为了评估FiT生成的图像质量，研究者们采用了多种评估指标，如Frechet Inception Distance（FID）、Inception Score（IS）、以及改进的精确度和召回率等。这些指标能够全面评估生成图像的质量和多样性，确保FiT在不同分辨率下都能产生高保真度的图像输出。通过这些评估，研究者们能够验证FiT在训练分布内外的分辨率上都具有出色的图像生成能力。

DiT和FiT在处理数据时的流程差异。(a)部分展示了DiT的流程，包括固定分辨率的处理方式；(b)部分展示了FiT的流程，突出了其灵活处理不同分辨率的能力

FiT模型视图：
(a) 灵活的训练流程，展示了如何将图像编码、分块、填充和计算损失。(b) 灵活的推理流程，展示了如何定义生成图像的位置、去噪和重塑。(c) FiT块，展示了模型的基本构成元素，包括2D位置嵌入、掩码多头自注意力（Masked MHSA）、SwiGLU等

FiT模型的架构基于DiT（Diffusion Transformer）进行了改进，特别是针对分辨率外推的局限性进行了优化。为了适应不同大小的图像，FiT采用了2D Rotary Positional Embedding（2D旋转位置嵌入，简称2D RoPE），这是从1D RoPE演变而来的技术，它通过在两个维度上分别应用旋转频率，增强了模型对空间位置信息的编码能力。

FiT模型中传统的多层感知器（MLP）被替换为Swish-Gated Linear Unit（SwiGLU），这是一种新型的激活函数，它结合了Sigmoid线性单元（SiLU）和Hadamard乘积，能够提供更有效的非线性变换，有助于提高模型的表达能力。

为了更有效地处理填充令牌，FiT使用了Masked Multi-Head Self-Attention（Masked MHSA）机制。这种机制通过掩码操作，区分了实际的令牌和填充令牌，确保在自注意力计算过程中，只有实际的令牌之间进行交互，而忽略填充令牌，从而避免了不必要的计算和潜在的错误信息传播。

FiT模型还引入了一种灵活的训练策略，允许模型在训练过程中动态调整序列长度，以适应不同分辨率的图像。这种策略通过将图像编码为不同长度的潜在令牌序列，并使用填充操作将它们统一到最大序列长度，使得模型能够处理任意长度的输入。

FiT模型的设计还包括了对网络架构的细致调整，如层数、隐藏尺寸和注意力头数的配置，以及对训练过程中的正则化和优化策略的选择。这些细节共同构成了FiT模型的核心，使其能够在保持高效率的同时，生成高质量的图像。

FiT模型在设计时考虑到了在不同分辨率下生成图像的能力，特别是在训练数据中未出现的分辨率。为了实现这一点，研究者们采用了一种无需额外训练的方法来扩展模型的分辨率适应性。

FiT模型在训练期间能够处理各种分辨率和纵横比的图像，但为了在推理时生成更高或更低分辨率的图像，需要一种有效的外推技术。论文提出了几种不同的外推方法，包括但不限于：

• Vanilla NTK和YaRN实现：直接将大型语言模型中的外推技术应用于2D RoPE，通过调整旋转基来适应不同的分辨率。
• VisionNTK和VisionYaRN：这是两种针对视觉任务改进的外推方法，它们利用了2D RoPE的解耦特性，分别对高度和宽度的旋转频率进行独立调整，以适应不同的纵横比。

这些方法通过修改模型的位置编码来适应新的分辨率，而不需要对模型权重进行任何调整或额外的训练。这样，FiT模型能够灵活地在不同的分辨率下生成图像，即使是在训练时未曾遇到的分辨率。

实验

研究者们遵循DiT-B和DiT-XL的设置，为基本模型FiT-B和xlarge模型FiT-XL设置了相同的层数、隐藏尺寸和注意力头数。研究者们采用了与DiT相同的预训练VAE编码器来编码和解码图像/潜在令牌。

研究者们通过一系列消融实验来检验和优化FiT模型的各个组成部分。这些实验包括对位置编码方式、前馈网络（FFN）中的激活函数，以及多头自注意力（MHSA）机制的调整。通过在特定的训练步骤后评估不同模型变体在多个分辨率上的性能，研究者们确定了哪些架构设计对于提升模型在不同分辨率图像生成任务上的表现最为关键。结果表明，引入2D旋转位置编码（2D RoPE）和SwiGLU激活函数，以及使用Masked MHSA代替传统的MHSA，能够有效提高模型对分辨率变化的适应性和生成图像的质量。

从DiT-B/2到FiT-B/2模型的消融研究结果

研究者们评估和改进了FiT模型在生成训练分布之外分辨率图像的能力。他们探索了多种分辨率外推技术，包括基于RoPE的插值方法，如NTK和YaRN，以及为FiT专门设计的VisionNTK和VisionYaRN方法。这些方法允许模型在不经过额外训练的情况下，有效地生成不同分辨率的图像，从而显著提高了模型的灵活性和应用范围。通过这些技术，FiT能够在保持图像质量的同时，处理更广泛的分辨率，包括那些在训练时未遇到的分辨率。

在分布外分辨率下，不同模型在ImageNet数据集上生成图像的性能比较。它包括了使用不同插值方法的FiT-B/2和DiT-B/2模型

研究者们训练了最高Gflops的模型FiT-XL/2，并在1.8M步数下进行了实验，以评估FiT在三种训练分布内的分辨率下的表现：256×256、160×320和128×384。他们将FiT与其他最先进的类条件生成模型进行了比较，包括BigGAN、StyleGAN-XL、MaskGIT、CDM、U-ViT、ADM、LDM、MDT和DiT。

实验结果表明，FiT-XL/2在所有评估的分辨率上都展现出了卓越的性能，超越了其他所有模型，特别是在160×320和128×384分辨率上，FiT-XL/2显著降低了FID得分，显示出其在生成高质量图像方面的显著优势。FiT-XL/2在256×256分辨率上尽管训练步数较少，但与其他经过更长时间训练的模型相比，仍显示出了竞争力。

在ImageNet数据集上，使用类条件图像生成模型在分布内和分布外分辨率生成图像的性能比较。表格中列出了不同模型在不同分辨率下的性能指标，如FID、sFID、IS、Precision和Recall

究者们评估了FiT-XL/2在三种训练分布外的分辨率下的表现：320×320、224×448和160×480，并与U-ViT、ADM、LDM-4、MDT和DiT等其他最先进的类条件生成模型进行了比较。他们发现FiT-XL/2在所有评估的分辨率和纵横比设置下都实现了最佳性能。

通过这些实验，研究者们展示了FiT模型在不同分辨率和纵横比下生成高质量图像的能力，证明了其在图像生成任务中的优越性和灵活性。

论文链接：https://arxiv.org/abs/2402.12376