1. ViT-5:视觉Transformer的现代化改造
过去几年,视觉Transformer(ViT)已经成为计算机视觉领域的重要架构。ViT-5作为面向2020年代中期的视觉Transformer,通过一系列组件级的现代化改造,显著提升了模型性能和训练稳定性。与传统的ViT相比,ViT-5在保持原有架构拓扑的同时,引入了来自大型语言模型(LLM)的最新优化技术,同时针对视觉任务进行了专门调整。
ViT-5的核心创新在于它不改变Transformer的基本结构,而是通过精心选择的组件改进来提升性能。这种方法使得ViT-5既保留了原始ViT的简洁性和通用性,又获得了现代Transformer架构的优势。特别值得注意的是,ViT-5的改进不是针对特定基准测试的过度优化,而是旨在提升模型在各种视觉任务上的通用表现能力。
2. ViT-5的核心架构改进
2.1 激活缩放(LayerScale)
LayerScale是ViT-5中的一个关键组件,它通过学习一个小的缩放因子来调整每个Transformer块的输出。具体来说,给定一个注意力或MLP块F(·),LayerScale执行以下变换:
x_{l+1} = x_l + F(x_l) ⊙ λ
其中λ∈ℝ^d是一个可学习的缩放向量,通常初始化为很小的值(如10^-4)。这种机制最初在CaiT中被提出,用于稳定深层ViT的优化过程。
有趣的是,ViT-5的研究发现LayerScale与后归一化(post-norm)在功能上存在密切联系。后归一化可以重写为:
x_{l+1} = (x_l + F(x_l)) ⊙ λ_p / Norm
其中Norm是RMS归一化项。这表明LayerScale直接控制块输出的尺度,而后归一化则隐式地缩放块输出和残差连接。实验表明,这两种方法带来的性能改进非常相似(见表1),但LayerScale提供了更大的灵活性和更低的计算开销,因此被选为ViT-5的默认组件。
2.2 归一化策略(RMSNorm)
ViT-5采用了RMSNorm替代传统的LayerNorm,这一选择受到了LLM领域最新进展的启发。RMSNorm去除了LayerNorm中的重新中心化操作,仅保留重新缩放功能。研究发现,归一化层的重新缩放不变性在实际应用中起主导作用,而去除重新中心化操作不仅不会降低性能,反而能通过减少不必要的偏移噪声带来轻微提升。
在ViT-5中,使用RMSNorm代替LayerNorm带来了两个好处:一是略微降低了计算成本,二是带来了小幅性能提升(如在ImageNet上ViT-B的top-1准确率提高了0.2%)。这一发现与LLM领域的观察结果一致,进一步证明了视觉和语言Transformer在架构优化上的相通性。
2.3 门控MLP与过门控问题
现代LLM广泛使用门控MLP架构,特别是SwiGLU(Swish-Gated Linear Unit)替代传统的GeLU激活。然而,ViT-5研究发现,在视觉Transformer中同时使用LayerScale和SwiGLU MLP会导致明显的性能下降(见表2)。
这个问题被归因于"过门控"现象:LayerScale和门控MLP都执行通道级过滤,增加了中间表示的稀疏性;当两者一起使用时,它们的组合效应可能导致激活过度稀疏。从这一角度看,LayerScale可以被视为一种静态门控,而同时使用LayerScale和SwiGLU MLP导致的性能下降可以被理解为过门控的结果。
基于这一发现,ViT-5选择保留原始MLP设计(使用GeLU激活),避免在视觉Transformer中同时使用LayerScale和门控MLP。这一决策特别针对ViT-XL规模(隐藏维度1152,参数4.49亿)以下的模型,对于更大规模的模型,过稀疏问题可能会减轻,但这需要未来进一步研究。
3. 位置编码与寄存器令牌
3.1 混合位置编码(APE + RoPE)
标准ViT使用可学习的绝对位置嵌入(APE),但这种方法在处理动态输入分辨率时存在固有局限。ViT-5引入了2D旋转位置嵌入(RoPE),这是从LLM领域借鉴并适应视觉任务的技术。然而,ViT-5并没有完全放弃APE,而是同时使用APE和2D RoPE。
这种混合方法的动机在于,仅使用相对位置编码会引入不希望的等变性。例如,在仅使用RoPE的情况下,图像在块级别上的翻转会变得完全等变(如图2所示),这意味着两个输入图像会被视为等价。虽然这种行为对图像分类等简单任务影响有限,但对于通用视觉骨干网络来说,绝对空间线索对于更复杂的视觉推理可能至关重要。
实验证明,ViT-5的这种混合位置编码策略显著增强了模型对不同分辨率的适应能力。如图3所示,ViT-5在224到512的广泛输入范围内没有表现出明显的性能下降,而且随着输入尺寸的增加,性能持续提升,显示出对动态分辨率更好的扩展性和泛化能力。
3.2 寄存器令牌(Register Tokens)
寄存器令牌是ViT-5的另一项重要创新。这些额外的可学习令牌被附加到输入块令牌中,最初是为了解决ViT中的伪影问题,但实际作用远不止于此。寄存器令牌提供了灵活的表征空间,支持各种扩展,如图像的元查询和1D令牌化。
ViT-5研究发现,在使用RoPE的ViT中,寄存器令牌也应该分配相对位置嵌入。这是因为当向量旋转操作仅应用于块令牌而寄存器令牌保持不旋转时,寄存器自然表现出与旋转角度较小的块令牌更低的余弦相似度。这种不平衡会扭曲注意力分布,并隐式引入不希望的位置偏差。
为解决这一问题,ViT-5为寄存器令牌配备了独立的2D RoPE,其频率基数明显高于用于块令牌的基数。这种设计在不同通道维度上诱导了寄存器和块令牌的独特旋转行为,有效地解耦了它们的位置相关性并消除了不希望出现的偏差。实验证明,这一修改稳定了寄存器与块令牌的交互,提高了整体表征质量(见表3)。
图4展示了DeiT-III和ViT-5在384×384分辨率下的注意力可视化比较。与先前观察一致,寄存器令牌有效抑制了注意力图中的背景伪影,使类别令牌能够更准确地关注图像的语义相关区域。在ViT-5中,特征图明显更清晰、更集中,这归功于寄存器令牌和相对位置嵌入的共同作用。
4. 自注意力机制的改进
4.1 QK归一化(QK-Norm)
最新的LLM(如Qwen3和Gemma3)开始通过额外归一化查询和键来改造自注意力机制。具体来说,QK归一化机制执行以下操作:
Q' = RMSNorm(Q), K' = RMSNorm(K)
Attention(Q,K,V) = Softmax(Q'K'^T/√d)V
ViT-5发现QK归一化能为ViT带来小幅性能提升,更重要的是能显著增强训练稳定性,减少优化过程中出现的尖锐损失峰值(如图5所示)。为利用这些鲁棒性优势,ViT-5将QK-Norm作为默认组件。
4.2 QKV投影的偏置项
ViT-5在整个模型中使用无偏置的RMSNorm,包括所有前归一化和QK-Norm层。这突出了自注意力更依赖加权投影而非加性偏置的观点。因此,ViT-5移除了QKV投影层中的偏置项以确保结构一致性。这一修改使QK-Norm运行更有效,并带来了明显的性能提升。
5. ViT-5的配置与性能
5.1 模型配置
ViT-5提供了四种规模的模型配置(见表4):
- ViT-5-Small:12层,384隐藏维度,6个头,4个寄存器,2200万参数
- ViT-5-Base:12层,768隐藏维度,12个头,4个寄存器,8700万参数
- ViT-5-Large:24层,1024隐藏维度,16个头,4个寄存器,3.04亿参数
- ViT-5-XL:28层,1152隐藏维度,16个头,4个寄存器,4.49亿参数
这种配置与现有ViT模型的实践保持一致,所有模型都使用4个可学习的寄存器令牌。
5.2 图像分类性能
在ImageNet-1k分类任务上,ViT-5在不同参数规模和计算预算下都实现了最先进的性能(见表5)。值得注意的是,ViT-5的性能优势随着模型规模和输入分辨率的增加而平稳扩展。特别是ViT-5-Large在384×384输入下达到了86.0%的测试准确率,显著超过了ViT(85.4%)和基于CNN的模型(85.5%)的先前最优结果。
这些结果表明,通过系统性地改进ViT的组件可以获得实质性的、可扩展的性能提升。由于ImageNet-1k仍然是评估视觉骨干网络最广泛使用的基准,ViT-5在这一数据集上的持续改进提供了强有力的证据,表明它具有增强的表征学习能力,可以作为替代现有架构的竞争性通用视觉骨干。
5.3 图像生成性能
ViT-5作为扩散Transformer(Diffusion Transformer)的骨干网络在图像生成任务上也表现出色。实验表明,在相同的训练配置下,仅将原始ViT骨干替换为ViT-5就能带来显著的性能提升(见表6)。当训练扩展到700万步时,ViT-5在Fréchet Inception Distance(FID)、Inception Score(IS)以及精确率和召回率指标上始终优于其ViT对应物(见表7)。
图6展示了不同规模模型在不同训练长度下的扩展曲线。使用ViT-5作为骨干的模型始终优于使用标准ViT的模型,并且随着训练的进行表现出平滑稳定的扩展行为。这些结果表明ViT-5能有效地泛化到视觉理解和图像生成任务,突显了组件级架构现代化的广泛影响。
5.4 密集预测性能
在ADE20k语义分割任务上,ViT-5同样表现出色(见表8)。所有模型在512×512分辨率下训练160k次迭代,骨干网络在ImageNet-1k上预训练相同周期。ViT-5-Small、ViT-5-Base和ViT-5-Large分别达到47.5%、49.1%和52.0%的mIoU,而相同参数预算下的DeiT-III对应模型分别为45.2%、48.0%和49.3%。
值得注意的是,性能差距随着模型规模的增加而扩大,表明组件级现代化的好处在更大模型中变得更加明显。这些结果也证明ViT-5能有效迁移到密集预测任务,并在不引入任务特定架构变化的情况下提供持续增益,进一步支持其作为强大通用视觉骨干的角色。
6. 消融研究与设计选择分析
6.1 与现有设计的比较
表9比较了ViT-5与现有视觉Transformer和现代语言模型的架构配置。在视觉导向设计中,如DeiT-III、DINO v2/v3、VisionLLaMA等模型采用了现代Transformer组件的子集,但没有同时采用本研究中确定的所有有效改进。因此,它们的性能始终低于ViT-5,在ImageNet-1k上的top-1准确率差距从0.31%到超过1.15%不等。
研究还发现,直接将语言模型的配置迁移到视觉任务是不够的。尽管LLaMA、Qwen和Gemma风格的设置引入了先进的归一化、门控MLP和相对位置编码,但当适应到ViT时(将RoPE扩展到2D),它们仍然表现不如ViT-5。这突显了为语言模型优化的架构选择不会简单地转化为视觉任务的最佳性能。
相比之下,ViT-5通过系统整合这些组件并考虑视觉特定因素,实现了最佳准确率,证明现有设计选择对视觉Transformer来说是次优的,促使了原则性的、组件级的现代化。
6.2 单个组件的影响
表10报告了不同规模模型上单组件消融的结果。总体而言,移除任何单个组件都会导致与完整ViT-5配置相比一致的准确率下降,确认每个设计选择都对最终性能有积极贡献。虽然不同组件和模型规模下的退化程度不同,但没有单一修改在所有设置中占主导地位,强调了所采用架构改进的互补性。
研究还发现,单个架构组件的影响在不同模型规模上并不一致。某些组件对较小模型表现出更强影响,而其他组件随着模型规模增长变得越来越关键。例如,在没有适当稳定的情况下用SwiGLU替换GeLU会导致小模型性能显著下降,反映了紧凑模型对门控诱导稀疏性的更高敏感性。
相比之下,LayerScale和2D RoPE等组件在更大模型上显示出更实质性的效果,其中更深的网络和更高的表征能力放大了训练稳定性和相对位置建模的重要性。值得注意的是,移除寄存器和QK-Norm对小模型导致相对适度的性能下降,但对基础和大型变体导致一致且更大的下降,表明这些组件随着模型容量增加发挥更重要作用。
6.3 组件细节分析
表11分析了ViT-5对几个关键架构超参数的敏感性。研究发现:
- LayerScale初始化值对性能有适度影响,太小的初始化值会使ViT-5-Base的准确率降低0.15%
- 2D RoPE的频率基数默认设置为1e-5(遵循大多数LLM的1D RoPE实践),表现出最佳性能
- 在所有ViT-5模型中使用了4个可学习寄存器,发现寄存器数量对预测性能只有微不足道的影响,并且对不同设置表现出良好的鲁棒性
7. 实际应用建议与注意事项
基于ViT-5的研究成果和实践经验,以下是为希望在项目中应用ViT-5的研究人员和工程师提供的实用建议:
-
组件选择平衡:虽然ViT-5集成了多种改进组件,但在实际应用中需要根据具体任务需求和计算资源进行权衡。例如,对于计算资源有限的小规模模型,可以优先考虑LayerScale和RMSNorm,而对于大规模模型则可以充分发挥2D RoPE和寄存器令牌的优势。
-
初始化策略:LayerScale的初始化值对模型性能有微妙但重要的影响。建议从1e-4开始尝试,并在验证集上微调。过小的初始化值(如1e-6)可能导致性能下降。
-
位置编码配置:在处理需要绝对位置信息的视觉任务(如目标检测)时,不应完全依赖RoPE,而应保留APE。对于分辨率变化较大的应用,可以适当提高RoPE的频率基数。
-
寄存器令牌使用:虽然4个寄存器令牌在大多数情况下表现良好,但对于特别复杂的视觉任务,可以尝试增加到16个,但要注意这不会显著增加模型参数。
-
训练稳定性监控:即使采用了QK-Norm,在训练大型ViT-5模型时仍应密切监控损失曲线,特别是当batch size较大时。出现不稳定迹象时可以适当降低学习率或增加预热步数。
-
硬件适配:ViT-5的无偏置设计使其特别适合在专用AI加速器上高效实现。在实际部署时,可以利用这一特性进行特定优化。
-
扩展应用:ViT-5作为骨干网络已证明在分类、生成和分割任务上有效,但将其适配到其他视觉任务(如视频理解、多模态学习)时,可能需要调整某些组件。建议从Base配置开始,再逐步扩展。
