自回归模型迎来全能选手！FlexVAR一模型通吃图像生成/修补，推理速度与质量自由调节

解读：AI生成未来

文章链接：https://arxiv.org/pdf/2502.20313 
Git链接：https://github.com/jiaosiyu1999/FlexVAR 

亮点直击

• 生成各种分辨率和宽高比的图像，甚至可以超过训练分辨率；
• 支持图像到图像的任务，例如图像修复、图像细化和图像扩展，而无需微调；
• 享受灵活的推理步骤，允许通过减少步骤加速推理或通过增加步骤提高图像质量。

总结速览

原文来源

本文作者来自北京交通大学、悉尼科技大学、美团、佐治亚理工学院。

解决的问题

• 灵活性不足：现有的视觉自回归模型（如VAR）通常只能生成固定分辨率的图像，缺乏生成不同分辨率和宽高比图像的灵活性。
• 残差预测的限制：现有模型采用残差预测的方式，依赖于固定的步骤设计，限制了图像生成的适应性和灵活性。同时，不同尺度的残差缺乏语义连续性，可能限制模型表示多样图像变化的能力。

提出的方案

• FlexVAR模型：提出了一种新的视觉自回归建模范式FlexVAR，摒弃了残差预测，改为在每个步骤中直接预测真实值（ground-truth），从而确保相邻尺度之间的语义连贯性。
• 多尺度约束的VQVAE：设计了一个可扩展的VQVAE tokenizer，增强了VQVAE对不同隐空间尺度的鲁棒性，使得图像可以在任意分辨率下重建。
• 可扩展的2D位置嵌入：提出了可扩展的2D位置嵌入方法，使用2D可学习查询初始化，使得尺度自回归建模可以扩展到各种分辨率和步骤，包括训练时未使用的分辨率和步骤。

应用的技术

• 自回归建模：利用自回归模型学习图像的概率分布，生成图像。
• VQVAE（向量量化变分自编码器）：用于图像的分割和重建，增强模型对不同尺度的适应性。
• Transformer架构：用于学习多尺度隐空间特征的概率分布，建模下一个尺度的真实值。
• 2D位置嵌入：引入2D可学习查询，扩展模型的分辨率和步骤适应性。

达到的效果

• 生成多样性：FlexVAR能够生成各种分辨率和宽高比的图像，甚至超过训练图像的分辨率。
• 多任务支持：支持多种图像到图像的任务，包括图像细化、图像修复、图像扩展等。
• 自适应步骤：可以根据需要调整自回归步骤，实现更快的推理速度或更高的图像质量。
• 性能提升：在ImageNet 256×256基准测试中，1.0B模型优于其VAR对应模型。在零样本转移生成过程中，13步生成的图像性能进一步提升，FID达到2.08，优于现有的自回归模型和流行的扩散模型。
• 零样本转移：在ImageNet 512×512基准测试中，FlexVAR在零样本转移情况下取得了与全监督训练的VAR 2.3B模型相竞争的结果。

方法

FlexVAR 概述

FlexVAR 是一种灵活的视觉自回归图像生成范式，它通过真实值（ground-truth）预测而非残差预测来实现自回归学习，从而能够在任何步骤独立生成合理的图像。在本文的方法中：(1) 一个可扩展的 VQVAE tokenizer 将输入图像量化为多尺度的 token 并重建图像。(2) 一个 FlexVAR transformer 通过尺度自回归建模进行训练，去除了残差预测。

多尺度量化与图像重建

主流的 VQVAE tokenizer 在单一分辨率下表现良好。然而，当扩展隐空间时，它们通常无法重建图像（如下图 3 所示）。这一观察促使我们探索一种可扩展的 tokenizer，能够将输入图像量化为多尺度的 token 并重建图像。具体来说，所提出的可扩展 tokenizer 首先将图像编码为多尺度的隐空间，然后使用量化器将隐空间特征转换为离散的 token，最后使用解码器从每个尺度的离散 token 中重建原始图像。

无残差的视觉自回归建模

实验

实现细节

最先进的图像生成

在 ImageNet-1K 基准测试中将 FlexVAR 与现有的生成方法进行比较，包括 GAN、扩散模型、随机扫描、光栅扫描和尺度扫描自回归模型。结果如下表 2 和表 3 所示。

ImageNet 256×256 的总体比较。为了确保公平比较，在表 2 中展示了参数规模小于 1B 的模型。我们的 FlexVAR 在所有生成方法中实现了最先进的性能，并且与 VAR 相比表现尤为出色。具体来说，在不同模型规模下，我们分别实现了 FID 提升 -0.45、-0.56 和 -0.12。

零样本推理（更多步骤）。在未训练的情况下使用 13 个步骤进行图像生成，如表 2 的最后一行所示。FlexVAR 可以灵活地采用更多步骤来提高图像质量。通过使用 13 个推理步骤，FlexVAR 进一步将性能提升至 2.08 FID 和 315 IS，展现了强大的灵活性和泛化能力。具体的步骤设计详见附录。

ImageNet 512×512 基准测试的零样本推理。使用 FlexVAR-d24 生成 512×512 图像，并在未训练的情况下在 ImageNet-512 基准测试上进行评估，如前面表 3 所示。尽管 FlexVAR 仅在分辨率 ≤256×256 上训练且仅有 1.0B 参数，但我们的 FlexVAR-d24 表现出了与 VAR 相竞争的性能。

消融实验

本文对 FlexVAR 中的各种设计选择进行了消融实验。由于计算资源有限，在 下表 4、5以及前面表6 中报告了使用短期训练方案（即 40 轮训练，约 70K 次迭代）得到的结果。

组件级消融实验

为了理解每个组件的影响，从标准 VAR 开始，并逐步添加每个设计：

• 基线：VAR 采用残差预测范式，表现良好（第 1 个结果），但其在图像生成方面的灵活性未达到预期）。
• 预测类型：直接将预测类型转换为 GT 是不可行的（第 2 和第 3 个结果）。我们采用了 VAR 和 Llamagen 的 VQVAE 分词器，二者均表现较差。这并不意外，因为当前的分词器缺乏对具有不同隐空间的图像的鲁棒性，而我们在训练过程中强制这些分词器获取多尺度隐空间特征（在前面图3中提供了详细分析）。
• 分词器：可扩展的分词器在训练过程中获取了合理的多尺度隐空间特征，使 FID 下降了 13.87（第 4 个结果）。然而，灵活的图像生成仍未实现。
• 位置嵌入：引入我们的可扩展 PE（最后一个结果）后，图像生成的灵活性显著提升，并进一步将 FID 降低至 3.71。

使用不同 VQVAE 重建图像

在前面图3中，通过缩放 VQVAE 分词器的隐空间特征来重建多尺度图像。现有的 VQVAE 分词器通常不支持在小到大范围内缩放隐空间特征。VAR 的 VQVAE 采用基于残差的训练策略，直接将其应用于非残差图像重建未能达到预期效果（第 1 行）。Llamagen 的 VQVAE 分词器仅在原始隐空间下表现出色，表明其不适用于尺度自回归建模（第 2 行）。

将 FlexVAR 迁移至 Mamba

最新研究 AiM 使用 Mamba 架构进行基于 token 的自回归建模。受此启发，对 FlexVAR 进行了 Mamba 适配，并评估其性能（前面表 5）。在类似的模型参数下，Mamba 相较于 Transformer 模型展现了竞争力，这表明 GT 预测范式可以有效适应 Mamba 这样的线性注意力机制。然而，考虑到该 Mamba 架构并未体现速度优势，我们未将其集成到最终版本中。

Mamba 具有固有的单向注意力机制，使得图像 token 在同一尺度下无法实现全局注意力。为了解决这一问题，我们在不同的 Mamba 层中采用了 8 条扫描路径来捕获全局信息。具体的 Mamba 架构细节见附录。

位置嵌入

在前面表 6 中，对多种步长 PE 和 x-y 坐标 PE 进行了实验。为了使模型在推理步骤上更具鲁棒性，并能够生成任意分辨率的图像，移除了固定长度的步长嵌入（第 2 行的结果），性能仅出现轻微变化。本文采用了一种类似于 ViT的非参数变体，与可学习变体相比，FID 仅相差 0.03。

GT 预测范式分析

收敛速度

本文对比了 VAR 和 FlexVAR 的训练损失，如下图4 所示。FlexVAR 展示了显著更低的损失值和更快的收敛速度，这表明预测真实值（ground-truth）而非残差对训练更加友好。这可能是由于不同尺度的残差缺乏语义连续性，而这种隐式预测方法可能会限制训练的收敛速度。

生成任意分辨率的图像

在下图1 和 图5 中，使用 FlexVAR-d24 生成了不同分辨率的图像。通过控制推理步骤，FlexVAR 能够生成任意分辨率的图像，尽管其仅在分辨率 ≤ 256px 的图像上进行了训练。生成的图像在多个尺度上表现出较强的语义一致性，并且更高分辨率的图像显示出更清晰的细节。

生成任意比例的图像

在 上图1 和 下图6 中，展示了不同长宽比的生成样本。通过在推理过程中控制每一步的长宽比，FlexVAR 能够生成不同长宽比的图像，体现了其在图像生成中的灵活性和可控性。

生成任意推理步长的图像

在下图7中，研究了使用 6 到 16 步生成 256×256 图像的 FID 和 IS，实验涵盖了 3 种不同的模型尺寸（depth 16、20、24）。随着步长的增加，生成图像的质量提升。更大规模的模型（如 FlexVAR-d24）提升更为显著，因为更大的 Transformer 结构能够学习更复杂和精细的图像分布。在训练过程中，我们最多使用 10 步以避免 OOM（内存溢出）问题。然而，在推理阶段，使用 13 步可使 FID 降低 0.13，这表明我们的 FlexVAR 对推理步长具有较强的适应性，可以通过减少步长加速图像生成，或通过增加步长获得更高质量的图像。更多步长设计的细节请见附录。

高分辨率图像细化

在下图8 中，输入低分辨率图像（如 256×256），并让 FlexVAR 输出高分辨率的细化图像。尽管仅在 ≤ 256px 的图像上进行了训练，FlexVAR 仍然能够通过提高输入图像分辨率来细化图像细节，例如示例中狗的眼睛。这展示了 FlexVAR 在图像到图像生成任务中的高度灵活性。

图像修补与扩展

对于图像修补（in-painting）和扩展（out-painting），对掩码外的 GT token 进行教师强制（teacher-force），并让模型仅生成掩码内的 token。同时，还注入了类别标签信息。结果如 下图9 所示。在不修改架构设计或训练方式的情况下，FlexVAR 在这些图像到图像的任务上表现良好。

图像扩展

对于图像扩展任务，生成目标类别的 1:2 长宽比图像，并在中心强制使用 ground-truth token。FlexVAR 在图像扩展方面表现良好，这表明其具有较强的泛化能力。

失败案例

当生成的图像分辨率达到训练分辨率的 3 倍或更高时，FlexVAR 生成的图像会出现明显的波浪状纹理和模糊区域（见下图11）。这种失败可能是由于当前训练数据集结构过于同质化，例如 ImageNet-1K 通常缺乏从粗到细的多尺度目标，因此在生成高分辨率图像的细节时会产生误差。

推测如果使用包含更精细细节的复杂数据集进行训练，模型可能会在更高分辨率的生成任务上表现得更加稳健。

结论

本文提出了 FlexVAR，一种灵活的视觉自回归图像生成范式，实现了无需残差预测的自回归学习。本文设计了可扩展的 VQVAE tokenizer 和 FlexVAR-Transformer。该 ground-truth 预测范式赋予了自回归模型极大的灵活性和可控性，使其能够在训练时未涉及的分辨率、长宽比和推理步长下进行图像生成。此外，它还能零样本迁移至多种图像到图像的生成任务。希望 FlexVAR 能成为一个稳固的基线，并促进视觉自回归建模及相关领域的研究。

局限性

发现，当生成的图像分辨率达到训练图像的3倍或更高时，会出现明显的波浪状纹理。这一问题可能是由于ImageNet-1K训练集结构过于同质化导致的。未来将进一步研究如何在零样本高分辨率图像生成任务中保持稳定性。

本文转自AI生成未来 ，作者：AI生成未来

原文链接：​​https://mp.weixin.qq.com/s/bURLljPt47K7xRJX8s7Ueg​​