梯度爆炸终结者：ViT-22B训练稳定性秘诀+ViTUnet图像生成首秀，全面力压ViT

本文主要解决了什么问题

1. 1. ViT-22B模型在局部环境中的训练和反应问题：由于ViT-22B在庞大的私有数据集上训练且参数未公开，导致在局部或站点级别数据上难以应用该模型。
2. 2. ViT-22B模型训练过程中的不稳定性问题：作者观察到训练过程中出现梯度爆炸性增加的现象，影响模型稳定性。
3. 3. ViT-22B模型在图像生成任务中的应用问题：图像生成任务在ViT-22B中尚未尝试，需要探索其适用性。
4. 4. ViT和ViT-22B在相同条件下的客观比较问题：之前研究中两种模型的预训练数据集大小不同，无法进行公平比较。
5. 5. 基于ViT的图像生成架构缺失问题：现有基于ViT的图像重建结构相对有限，需要提出新架构填补这一空白。

本文的核心创新是什么

1. 1. 提出了解决ViT-22B训练不稳定性的方法：通过向并行网络引入归一化解决梯度爆炸问题。
2. 2. 提出了ViTUnet架构：一种基于ViT的新图像到图像翻译架构，结合了Transformer和CNN残差结构，用于图像生成任务。
3. 3. 在相同条件下从头训练并比较ViT和ViT-22B：使用相同的优化器、训练周期和数据增强技术，对两种模型进行客观性能比较。
4. 4. 探索了ViT-22B在图像生成任务中的应用：这是ViT-22B研究中尚未涉及的领域，扩展了其应用范围。

结果相较于以前的方法有哪些提升

1. 1. 在相同参数规模下，从头训练的ViT-22B模型整体性能优于ViT：在相同条件下从零开始训练两种模型，确认ViT-22B性能超越传统ViT架构。
2. 2. 成功稳定了ViT-22B的训练过程：通过引入归一化解决了梯度爆炸问题，使模型可以在局部环境中稳定训练。
3. 3. 在部分图像生成任务中取得更好结果：在Monet2Photo任务中，ViT-22B主干网络模型以更好的FID分数表现优于ViT主干网络模型。
4. 4. 提出了有效的基于ViT的图像生成架构：ViTUnet结构允许ViT的注意力机制与调整大小相结合，通过残差连接最小化信息损失，实现图像重建。

局限性总结

1. 1. ViT-22B在图像生成任务中并不总是表现更优：在苹果和橙子的转换任务中，ViT-22B-S主干网络在重建形状和颜色方面完全失败。
2. 2. 生成的图像存在块状形状问题：这似乎源于ViT的注意力过程，而非模型大小，即使使用更大规模的ViT-B作为主干网络也未能解决。
3. 3. 评估指标可能不够全面：尽管生成的苹果和橙子图像质量较差，ViT-22B主干网络模型仍记录了最高FID分数，表明FID分数可能不是生成模型评估的绝对标准。
4. 4. 研究仍处于探索阶段：作者表示，在ViT-22B中未探索的实验结果得到验证后，该结构才可能在各种任务中证明其全面益处。

导读

已有尝试在视觉模型中构建类似于LLM的大规模结构，例如ViT-22B。尽管这项研究提供了大量分析和见解，但作者对其实际应用的理解仍不完整。因此，作者考察了该模型结构在局部环境中的反应和训练方式。此外，作者还指出了训练过程中的不稳定性，并对模型进行了一些修改以增强其稳定性。从头训练的ViT-22B模型，在相同参数规模下，整体性能优于ViT。此外，作者探索了图像生成任务，这一任务在ViT-22B中尚未尝试。作者提出了一个基于ViT的图像生成架构，并研究了在图像生成任务中，ViT和ViT-22B哪种结构更为合适。

1 引言

最大规模大语言模型（LLMs）[5, 27]的出现为作者如何处理深度学习模型提供了新的视角。它通过首先设计一个庞大的基础模型，然后将其应用于下游任务，而非创建更小的特定任务模型[22, 29]，展现了卓越的性能。这一趋势在视觉研究中也日益显著。[31]提出了一种扩展ViT模型的方法，展示了为视觉领域构建基础模型的可行性。然而，这一尝试（规模为15B）与拥有540B[5]的LLM模型相比仍然较小。因此，ViT-22B[6]通过架构和配方修改将可学习参数数量提升至22B。该模型不仅在多个可定量评估的下游任务中达到了最先进或相当的性能，而且在定性上相较于几种视觉模型也展现了更优越的性能。

按照这种方式，作者寻求与“所有深度学习模型的大小越大越好”这一信念相符的发现。然而，一些研究警告不要认为基础模型适用于所有任务，强调需要仔细考虑[3]。从语言学的角度进行了分析，其中对模型可靠性和发展潜力的评论在不同领域的专家之间有所不同。由于视觉基础模型与语言基础模型相比进展较慢，因此难以确定具体问题。然而，可以确定的是ViT-22B是在一个庞大的私有数据集上训练的[26, 19]，并且其训练参数未公开，这使得作者难以从中获益。

这意味着如果作者有局部或站点 Level 的数据，作者需要自行训练ViT-22B模型。此外，作者尚未了解模型的进步是否完全源于其规模，还是其结构变化中存在其他原因。现有研究缺乏统一模型规模和预训练数据集（即JFT）的实验。因此，作者局部训练ViT-22B结构并观察训练过程，以了解其反应方式。在模型训练过程中，作者观察到梯度爆炸性增加的现象。这似乎源于没有归一化或调节的并行线性结构。作为证据，向该并行网络引入归一化解决了这一问题。在相同条件下从零开始训练ViT和ViT-22B的结果，指导作者选择合适的结构。

此外，作者分析了图像生成任务，该任务在ViT-22B研究中尚未尝试。虽然图像生成是计算机视觉领域广泛研究的课题[16, 24, 4, 30]，但基于ViT架构的研究有限[15, 13, 9]。作者引入了一种名为ViTUnet的新结构，并使用ViT和ViT-22B作为 Backbone 网络，观察这些模型在图像生成任务中的表现。ViTUnet结构涉及与Unet[25]类似的过程，通过残差连接和编码-解码步骤在低维空间中映射和重建图像。作者强调了基于ViT的图像重建结构的缺失，并定性和定量地展示了作者提出的模型性能。

2 相关工作

2.1 视觉Transformer

Transformer架构[28]在语言翻译[10]和问答任务[5]中取得了显著的性能提升。基于注意力机制的操作忠实地模拟了语言能力，即推理序列中词语之间的关系。另一方面，[7]提出图像也可以被表述为序列。将图像按顺序切割成块并展平，会得到一个块序列。作者将这种结构称为视觉Transformer（ViT），当在包含数百万私有图像的数据集上进行预训练时，它在各种视觉任务中取得了最先进（SOTA）的结果[26]。注意力机制缺乏传统视觉模型（如CNN）中存在的局部特征归纳偏置，因此需要更广泛的训练数据才能达到其效果[28, 8]。[17]提出了Swin Transformer，这是ViT的一种变体，通过调整块大小来模拟CNN中不同感受野的效果。类似于更深的CNN层会增加感受野，Swin Transformer的上层探索更广泛区域的信息。[20]引入了一种具有减少学习参数的开创性ViT结构。虽然传统的ViT通过显著增加模型容量来解决归纳偏置的缺失问题，但这一 Proposal 仅通过张量展开和自注意力来实现。尽管取得了这些进展，但视觉基础模型的探索仍然不足。

2.2 ViT-22B

ViT-22B [6]针对视觉基础模型缺失的问题，引入了一种新的ViT架构及参数设置。它采用并行线性网络替代 FFN ，并对键和 Query 引入层归一化以防止梯度爆炸。训练22B参数需要特殊的并行化策略。本研究从两个新颖的视角应对这一挑战。根据其实验结果，大规模ViT结构在大多数下游任务中表现出良好性能。然而，对该结构的深入分析以及在实际局部规模应用中的信息仍然不足。

3 在局部训练 ViT-22B

3.1 ViT-22B的训练方案

3.2 从零开始分类

作者对从头开始训练的模型进行了客观的性能比较，具体为ViT-22B和ViT。由于预训练数据集大小的差异，[6]中的实验结果无法进行客观比较。特别是对于ViT-22B，在4B数据集上进行预训练提供了一个显著有利的起点，相比之下，ViT是在300M数据集上训练的。作者使用AdamW优化器[18]对模型进行了200个epoch的训练，并在总epoch数的30%、60%和90%处引入了多步学习率调度器。仅应用了如翻转等基本的空间增强。

4 图像生成

作者探讨了一个在ViT-22B研究中未被涉及的主题：生成模型。尽管ViT在几乎所有视觉领域都得到了广泛的研究和实践应用，但它并不常用于生成模型。这是因为ViT的自注意力机制本身无法捕捉图像的表示[23, 2]。因此，只有少数研究尝试基于ViT进行图像生成。[13]提出了一种利用ViT结构从噪声z生成高分辨率图像的方法。该方法涉及将噪声生成的图像进行上采样以创建高分辨率图像，同时引入一个分辨率辅助任务作为损失项来生成高质量图像。[15]与[14]类似，首先将风格映射到某个嵌入空间，然后将其输入到ViT结构中以生成图像。[9]使用一个称为码本的特征包来保留CNN的归纳偏置优势。由于这些研究要么从噪声生成图像，要么需要参考图像，因此它们无法应用于图像到图像的转换。因此，作者提出了一种架构，该架构能够利用ViT和ViT-22B结构实现图像到图像的转换。

4.1 基于ViT的新图像到图像翻译架构

作者提出了一种基于ViT的生成式架构用于图像翻译：ViTUnet。该结构借鉴了在图像生成中广泛使用的Unet[25]。类似于模仿基于CNN模型观察不同分辨率的Swin[17]架构，作者的结构通过接收不同大小的输入并将它们融合来工作。图3展示了ViTUnet的完整流程。输入被分割成大小为p的块，并嵌入块维度。为此嵌入，作者添加了可学习的位置嵌入以提供空间信息。该模型由编码和解码过程组成，解码过程结合了Transformer和CNN残差结构。在编码过程中，嵌入的块在通过一个包含深度/6次迭代的Transformer结构后进行下采样，然后被输入到下一个包含深度/6次迭代的Transformer结构。随后，在解码过程中，输入被放大两倍并输入到一个包含深度/6次迭代的Transformer结构。在整个解码过程中，所有输入都包括残差连接，在增加输入块数量的方向上进行合并。通过 Reshape 以对齐图像的宽度和通道，并通过CNN残差重建原始图像的形状。这种结构允许ViT的注意力机制与调整大小相结合，通过残差最小化信息损失，并实现图像重建。

4.2 生成结果

作者采用了Apple2Orange和Photo2Monet进行图像到图像的转换，遵循了文献[32]中的训练策略。

是一个具有

感受野的CNN块生成对抗网络[12]。图4展示了当ViTUnet使用ViT-S作为主干网络时和使用ViT-22B-S作为主干网络时的生成结果。这两种主干网络在Monet2Photo任务中表现出定性的较差性能。然而，从照片到莫奈风格的转换产生了相对合理的图像。在苹果和橙子的转换中，只有ViT主干网络生成了可理解的图像。ViT-22B-S主干网络在重建水果之外的形状和颜色方面完全失败。最后，作者采用ViT-B作为主干网络，增加了训练参数数量，并生成了图像。图5显示了生成的图像，为苹果和橙子提供了整体风格转换，其他色调的变化最小。然而，它未能解决块状形状的问题，这似乎源于ViT的注意力过程，而非模型大小。表2展示了图像生成实验的Fréchet inception距离（FID），其中最大的模型获得了最佳的FID分数。尽管生成的苹果和橙子图像质量较差，ViT-22B主干网络模型记录了最高的FID分数。与定性评估相比，ViT-22B-S主干网络模型在Monet2Photo任务中以更好的FID分数表现优于ViT-S主干网络模型。这表明FID分数可能不是生成模型评估的绝对标准。

5 结论

在本研究中，作者加深了对局部训练的类似LLM的视觉模型ViT-22B的理解。局部训练该模型时发现其稳定性不稳定，作者提出了几种方法来稳定它。在相同条件下，作者确认ViT-22B的性能超越了传统的ViT架构。此外，作者引入了一种图像生成模型，这是ViT结构中较少探索的方面。这表明ViT结构也可以作为图像生成的有价值工具。然而，作者不能断言ViT-22B结构在图像生成方面始终表现更优。基于这些结果，可以预期，在ViT-22B中未探索的实验结果得到验证后，该结构可能在各种任务中证明其益处。作者期望即使在作者的尝试之后也能继续验证，为建立视觉应用的标准模型铺平道路。

参考

[1]. Deeper Inside Deep ViT