ViT-5: Vision Transformers for The Mid-2020s—— 面向2020年代中期的视觉Transformer

这篇题为《ViT-5:面向2020年代中期的视觉Transformer》的论文,核心研究内容是对经典的视觉Transformer架构进行系统性的现代化升级。其目标并非颠覆性重构,而是通过"组件级"的精细优化,在保留其基础结构的同时,显著提升性能、稳定性和泛化能力。

以下是对该研究内容的完整专业总结:

一、研究背景与动机

  1. 现状诊断: 自2020年ViT问世以来,其核心架构变化甚微。而同期,大语言模型通过引入RMSNorm、RoPE、门控机制等一系列组件级改进,取得了巨大成功。

  2. 核心问题: 研究者提出疑问:视觉Transformer的潜力是否被充分挖掘?能否将语言模型中验证有效的架构演进,系统地迁移并适配到视觉领域?

  3. 研究方法: 论文采取了一种"模块化"的研究路径,逐一审视并改进ViT的每个关键组件,而非设计全新的网络拓扑结构。

二、核心技术贡献:ViT-5的架构创新

ViT-5(代表过去五年架构演进的集成)在保留标准Attention-FFN结构的前提下,进行了以下七大核心组件的改进与重新设计:

  1. 激活缩放(Activation Scaling):

    • 改进: 引入可学习的LayerScale(层缩放),对每个残差块的输出进行缩放。

    • 发现: 论文发现LayerScale在功能上与LLM中常用的后归一化(Post-Norm)高度相似,但LayerScale计算开销更小、灵活性更高,因此被采用。

  2. 归一化(Normalization):

    • 改进: 用均方根归一化全面取代层归一化。

    • 动机: 借鉴LLaMA等LLM的成功经验,移除层归一化中的中心化操作,仅保留缩放不变性,在略微降低计算成本的同时带来了微小的性能增益(+0.2%)。

  3. 门控MLP(Gated MLP):

    • 重要发现(反例): 尽管SwiGLU在LLM中非常流行,但论文发现将SwiGLU与LayerScale结合会导致"过度门控"(Over-gating)问题,使激活过于稀疏,反而降低性能。

    • 结论: 在中规模模型中,应避免两者同时使用。因此ViT-5保留了原始的GeLU激活函数。

  4. 位置编码(Positional Encoding):

    • 创新: 联合使用绝对位置编码和2D旋转位置编码。

    • 动机: 仅用RoPE(相对位置)会导致对图像的全局翻转产生不变性(图2),这对于需要绝对空间线索的复杂视觉任务是致命的。两者结合既获得了RoPE的动态分辨率处理能力,又保留了APE的绝对空间信息。

    • 效果: 显著提升了模型对动态输入分辨率的鲁棒性和可扩展性(图3)。

  5. 寄存器令牌(Register Tokens):

    • 改进: 引入额外的可学习"寄存器"令牌,并对其进行特殊处理。

    • 关键发现: 在使用RoPE时,必须为寄存器令牌也施加位置嵌入,且需采用比图像块更高的频率基数,以消除其与图像块间的位置偏差。

    • 效果: 有效抑制了注意力图中的背景伪影,使类别令牌能更精确地关注语义区域(图4)。

  6. QK-归一化(QK-Normalization):

    • 改进: 在自注意力计算前,对查询和键分别应用RMSNorm。

    • 效果: 带来了微小的性能提升,但更重要的是显著增强了训练稳定性,有效防止了训练过程中损失函数的剧烈尖峰(图5)。

  7. 偏置项(Bias Terms):

    • 改进: 移除QKV投影层中的偏置项。

    • 动机: 确保与RMSNorm的结构一致性,使QK-Norm能更有效地发挥作用。

三、实验验证与核心结果

论文在多个任务上对ViT-5进行了全面评估,证明了其作为新一代主干的优越性:

  1. 图像分类:

    • 数据集: ImageNet-1k。

    • 结果: ViT-5各尺寸模型均超越当前最先进的DeiT-III和ConvNeXt。例如,ViT-5-Base达到84.2%,优于DeiT-III-Base的83.8%;ViT-5-Large在384分辨率下达到86.0%,显著刷新性能记录。

  2. 图像生成:

    • 框架: 嵌入SiT扩散模型。

    • 结果: 在同等训练条件下,使用ViT-5的模型FID(弗雷歇初始距离)指标全面优于使用原始ViT的SiT/DiT。在700万步的长训练中,ViT-5-XL取得了1.84 FID的优异成绩,远优于基线模型的2.06。

  3. 密集预测(语义分割):

    • 数据集: ADE20k。

    • 结果: ViT-5在UperNet框架下,各尺寸模型均大幅超越DeiT-III,且性能差距随模型规模增大而扩大(ViT-5-Large达到52.0% mIoU,DeiT-III-Large为49.3%)。

  4. 消融与对比研究:

    • 验证有效性: 逐组件消融实验证明,每个改进都贡献了正向收益。

    • 对比现有设计: 无论是从现有视觉模型(如DeiT-III)还是从LLM(如LLaMA)直接移植的配置,其性能均落后于ViT-5。这证明专为视觉任务量身定制的组件组合至关重要。

四、结论与意义

  1. 核心发现: 视觉Transformer仍有巨大的优化空间。通过系统性的、原则性的组件级现代化改造,可以在不改变宏观架构的情况下释放出显著的性能潜力。

  2. 研究路径: ViT的演进可以像LLM一样,遵循"最佳实践驱动"的模块化发展路径,通过精良的组件选择来换取性能和稳定性,而非依赖越来越复杂的宏观设计。

  3. 实际价值: ViT-5作为一个即插即用的升级方案,为2020年代中期的视觉和多模态系统提供了一个更强大、更鲁棒的基础模型,有望促进该领域的进一步发展。

这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:

项目地址在这里,如下所示:

模型地址如下:

Model Input Resolution Params Top-1 (ImageNet-1K) HF Link
ViT-5-Small 224 22M 82.2% Download
ViT-5-Base 224 87M 84.2% Download
ViT-5-Base 384 87M 85.4% Download
ViT-5-Large 224 304M 84.9% Download
ViT-5-Large 384 304M 86.0% Available soon

环境依赖安装操作如下:

python 复制代码
# Install PyTorch (CUDA 12.4)
pip install torch==2.4.1 torchvision --index-url https://download.pytorch.org/whl/cu124

# Install core dependencies
pip install timm==0.4.12 numpy==1.26.4 wandb einops

# Install NVIDIA Apex (required for fused optimizers)
git clone https://github.com/NVIDIA/apex
cd apex
APEX_CPP_EXT=1 APEX_CUDA_EXT=1 pip install -v --no-build-isolation .

# (Optional) Install Flash Attention for faster training
pip install flash-attn==2.6.3 --no-build-isolation

训练微调命令如下:

python 复制代码
# ImageNet Pretraining (8 GPUs example)
# ViT-5-Small
torchrun --nproc_per_node 8 main.py \
  --model vit5_small --input-size 224 --data-path YOUR_IMAGENET_PATH --output_dir DIR_TO_SAVE_LOG_AND_CKPT \
  --batch 256 --accum_iter 1 --lr 4e-3 --weight-decay 0.05 --epochs 800 --opt fusedlamb --unscale-lr \
  --mixup .8 --cutmix 1.0 --color-jitter 0.3 --drop-path 0.05 --reprob 0.0 --smoothing 0.0 --ThreeAugment \
  --repeated-aug --bce-loss --warmup-epochs 5 --eval-crop-ratio 1.0 --dist-eval --disable_wandb

# ViT-5-Base
torchrun --nproc_per_node 8 main.py \
  --model vit5_base --input-size 192 --data-path YOUR_IMAGENET_PATH --output_dir DIR_TO_SAVE_LOG_AND_CKPT \
  --batch 256 --accum_iter 1 --lr 3e-3 --weight-decay 0.05 --epochs 800 --opt fusedlamb --unscale-lr \
  --mixup .8 --cutmix 1.0 --color-jitter 0.3 --drop-path 0.2 --reprob 0.0 --smoothing 0.0 --ThreeAugment \
  --repeated-aug --bce-loss --warmup-epochs 5 --eval-crop-ratio 1.0 --dist-eval --disable_wandb

# ViT-5-Large
torchrun --nproc_per_node 8 main.py \
  --model vit5_large --input-size 192 --data-path YOUR_IMAGENET_PATH --output_dir DIR_TO_SAVE_LOG_AND_CKPT \
  --batch 256 --accum_iter 1 --lr 3e-3 --weight-decay 0.05 --epochs 400 --opt fusedlamb --unscale-lr \
  --mixup .8 --cutmix 1.0 --color-jitter 0.3 --drop-path 0.35 --reprob 0.0 --smoothing 0.0 --ThreeAugment \
  --repeated-aug --bce-loss --warmup-epochs 5 --eval-crop-ratio 1.0 --dist-eval --disable_wandb

# Fine-tuning from Pretrained Checkpoint
# ViT-5-Small
torchrun --nproc_per_node 8 main.py \
  --model vit5_small --finetune PATH_TO_YOUR_CKPT --data-path YOUR_DATASET_PATH --output_dir DIR_TO_SAVE_LOG_AND_CKPT \
  --batch 64 --lr 1e-5 --weight-decay 0.1 --epochs 20 --unscale-lr --aa rand-m9-mstd0.5-inc1 --drop-path 0.05 \
  --reprob 0.0 --smoothing 0.1 --no-repeated-aug --dist-eval --load_ema --eval-crop-ratio 1.0 --disable_wandb

# ViT-5-Base
torchrun --nproc_per_node 8 main.py \
  --model vit5_base --finetune PATH_TO_YOUR_CKPT --data-path YOUR_DATASET_PATH --output_dir DIR_TO_SAVE_LOG_AND_CKPT \
  --batch 64 --lr 1e-5 --weight-decay 0.1 --epochs 20 --unscale-lr --aa rand-m9-mstd0.5-inc1 --drop-path 0.25 \
  --reprob 0.0 --smoothing 0.1 --no-repeated-aug --dist-eval --load_ema --eval-crop-ratio 1.0 --disable_wandb

# ViT-5-Large
torchrun --nproc_per_node 8 main.py \
  --model vit5_large --finetune PATH_TO_YOUR_CKPT --data-path YOUR_DATASET_PATH --output_dir DIR_TO_SAVE_LOG_AND_CKPT \
  --batch 64 --lr 1e-5 --weight-decay 0.1 --epochs 20 --unscale-lr --aa rand-m9-mstd0.5-inc1 --drop-path 0.5 \
  --reprob 0.0 --smoothing 0.1 --no-repeated-aug --dist-eval --load_ema --eval-crop-ratio 1.0 --disable_wandb

摘要

本文通过借鉴过去五年中架构方面的进展,对视觉Transformer主干网络进行了系统性的现代化研究。在保留经典的注意力-前馈网络结构的同时,我们对归一化、激活函数、位置编码、门控机制和可学习令牌等组件进行了逐项改进。这些更新构成了新一代的视觉Transformer,我们称之为ViT-5。大量实验表明,在理解和生成基准测试中,ViT-5始终优于最先进的普通视觉Transformer。在ImageNet-1k分类任务上,ViT-5-Base在相当的计算量下达到了84.2%的top-1准确率,超过了DeiT-III-Base的83.8%。ViT-5作为生成建模的主干网络也表现更佳:当将其嵌入SiT扩散框架时,它实现了1.84的FID,而使用普通ViT主干时为2.06。除了主要指标外,ViT-5还表现出改进的表示学习能力和有利的空间推理行为,并能可靠地迁移到各种任务中。其设计与当代基础模型的实践相一致,ViT-5为2020年代中期的视觉主干网络提供了一个简单的即插即用升级方案,可替代普通ViT。

1. 引言

自2020年底问世以来,视觉Transformer(ViT)深刻地重塑了视觉编码范式。其与语言Transformer在架构上的高度一致性催生了大量用于多模态理解和生成的成功视觉-语言系统。在过去大约五年间,语言模型在原始Transformer设计基础上经历了一系列系统性的改进。这些改进包括采用更先进的激活函数和归一化层、从绝对位置编码向相对位置编码的转变,以及改进的注意力归一化和门控机制。总体而言,这些结构上的演进显著提升了表示能力和训练稳定性,并成为现代基础模型快速进步的主要驱动力。

然而,与语言模型架构快速持续的改进相比,ViT的核心设计自诞生以来基本保持不变。例如,最先进的普通ViT主干DeiT-III仅引入了称为LayerScale的最小修改,而其他部分则保留了原始的ViT架构。类似地,近期的大规模视觉-语言模型如SigLIP-2和Qwen3-VL在视觉编码方面仍然依赖基本相同的普通ViT设计。这种架构演进上的相对停滞引发了一个重要问题:ViT的表示潜力是否仍未得到充分优化?鉴于语言模型在结构改进方面取得的显著成果,一个自然的假设是,这些跨越归一化、激活函数、位置编码和注意力机制的进展,能否系统地迁移到视觉模型中,以释放进一步的性能和效率提升。

在这项工作中,我们对优化ViT的架构设计进行了系统性的研究。我们的研究围绕原始的普通ViT公式展开,有意识地保留其基本的注意力-前馈网络主干,同时专注于识别各个架构组件的最有效变体。我们的目标不是提出一个彻底的重新设计,而是理解如何以原则性和模块化的方式将现代化的设计选择融入ViT。我们明确分析并实证验证了自ViT问世以来出现的结构改进的影响,例如LayerScale、旋转位置编码(RoPE)、QK归一化和寄存器令牌,其中许多已在视觉或语言Transformer中成为标准。通过大量实验,我们得出了两个关键观察。首先,当前的ViT架构仍未得到充分优化,改进其核心组件可以在不同任务上持续带来显著的性能提升。其次,现有的架构改进并非严格正交:简单地将所有现代组件组合在一起并不一定能带来最佳性能,有效的现代化需要精心设计。

受这些发现的启发,我们引入了一种全面改进原始设计的下一代视觉Transformer架构。我们将这个新模型命名为ViT-5,它象征着对过去五年主要架构演进的整合。我们在图1中展示了ViT-5的架构概览。具体来说,ViT-5引入了一套有助于稳定ViT的组件,如LayerScale、RMSNorm和QK-Norm,以及增强空间推理能力的模块,如RoPE和寄存器令牌,同时刻意避免使用SwiGLU激活函数------尽管它在现代LLM中很受欢迎,但在视觉模型中可能导致过度门控问题(定义见第3.3节)。我们在一系列视觉任务上的实验证明了ViT-5强大的表示能力和良好的泛化性。在标准的ImageNet-1k分类任务上,base大小的ViT-5达到了84.2%的top-1准确率,优于之前最先进的普通ViT基线DeiT-III的83.8%。在图像生成方面,基于ViT-5构建的扩散模型实现了1.84的FID,在计算成本几乎相同的情况下,显著优于SiT基线的2.06 FID。此外,密集预测任务的结果以及定量评估进一步突出了ViT-5在空间建模和表示学习方面的优势。

这项研究提高了ViT主干的表示能力,并缩小了视觉模型与现代语言模型之间的架构差距。我们希望ViT-5能够促进多模态系统更高效的构建,并激发开发能够在不同模态间无缝泛化的统一Transformer架构。

2. 相关工作

现代视觉主干。视觉Transformer的引入标志着视觉表示学习从卷积主干向基于Transformer的架构的转变,在视觉预训练、多模态建模和图像生成方面带来了一系列进展。继ViT架构之后,DeiT系列引入了更高效的训练策略,使得普通视觉Transformer在相当的计算预算下能够与当代CNN主干(如ConvNeXt)的性能相匹配。除了普通ViT,大量工作还探索了层级视觉Transformer和ViT-CNN混合架构,这些架构结合了更强的归纳偏置,通常在中等规模的视觉识别基准上表现更佳。本文的目标并非在ImageNet或COCO等特定基准上最大化绝对性能。相反,我们专注于普通视觉Transformer的组件级架构优化,旨在保持其在各种模型大小和不同视觉任务中的强大泛化能力和可扩展性。

Transformer的架构进展。Transformer架构经历了一系列逐元素的改进,这主要受到大型语言模型快速发展的推动。例如,早期的LLaMA模型引入了旋转位置嵌入和RMS归一化,并用基于SwiGLU的前馈网络替换了标准MLP。最近,Gemma3进一步在自注意力中对查询和键进行了显式归一化,而Qwen3则移除了QKV投影中的偏置项以提高训练稳定性和效率。类似的趋势也出现在视觉模型中。CaiT引入了LayerScale以稳定深度视觉Transformer的优化,后续工作表明,引入寄存器令牌可以有效缓解Transformer中的激活伪影。重要的是,这些改进都围绕着原始的Transformer公式,保留了经典的注意力-前馈网络推理结构,而不改变整体模型拓扑结构。如何有效利用这类组件级的架构改进是我们研究的核心焦点。

3. 模型

3.1. 激活缩放

在对ViT的早期改进中,人们发现对每个块的输出引入一个可学习的缩放因子可以显著提高深度模型的训练稳定性和性能。形式上,给定一个注意力或MLP块F(·),它应用以下变换:

x_{l+1} = x_l + F(x_l) ⊙ λ, (1)

其中λ ∈ R^d是一个可学习的缩放向量,通常初始化为一个较小的值(例如10^{-4})。这种机制通常被称为LayerScale,并已成为许多现代ViT架构(如DINO v3)的默认组件。

虽然LayerScale在语言模型中尚未被广泛采用,但我们观察到LayerScale与后归一化之间存在内在联系,后者已被认为是稳定深度LLM训练的一项重要技术。形式上,后RMSNorm可以重写为

x_{l+1} = (x_l + F(x_l)) ⊙ λ_p / Norm (2)

其中

Norm = RMS(x_l + F(x_l)) (3)

且λ_p ∈ R^d是一个缩放向量。我们可以发现,LayerScale直接控制块输出的尺度,而后归一化则隐式地缩放块输出和残差连接的尺度。在我们的实验中,LayerScale和后归一化带来了非常相似的性能改进(见表1)。鉴于LayerScale提供了更大的灵活性和更低的计算开销,我们在ViT-5中将其作为默认组件。我们希望关于LayerScale和后归一化之间功能关系的这一观察,也能为未来的LLM架构设计提供有用的见解。

3.2. 归一化

自LLaMA、PaLM和Gopher等有影响力的语言模型以来,LLM架构中的事实标准已从层归一化(LayerNorm)大幅转向均方根归一化(RMSNorm)。这一转变背后的一个普遍观察是,归一化层的重缩放不变性主导了其实际效果。移除LayerNorm中的中心化操作不会降低性能,甚至可以通过减少不必要的偏移噪声带来轻微的改进。我们在ViT中也观察到了类似的现象。在我们的实验中,用RMSNorm替换LayerNorm略微降低了计算成本,并带来了适度的性能提升(例如,在ImageNet上对ViT-B的top-1准确率提升了+0.2%)。基于这些发现,我们在ViT-5中全面使用RMSNorm,替换了原始架构中的所有LayerNorm层。

3.3. 门控MLP

类似于从LayerNorm到RMSNorm的转变,现代LLM广泛采用了门控MLP架构,其中传统的GeLU激活函数被SwiGLU取代。尽管如此,在本研究中,我们观察到将SwiGLU MLP与LayerScale结合使用会导致ViT的性能显著下降(见表2)。我们将此问题归因于LayerScale和门控MLP都有效地执行了通道级滤波,这增加了中间表示的稀疏性;当一起使用时,它们的组合效应可能导致过于稀疏的激活。从这个角度来看,LayerScale可以被视为一种静态门控形式,而联合使用LayerScale和SwiGLU MLP时观察到的性能下降可以理解为过度门控的情况。

这里我们得出一个初步结论:至少在ViT-XL规模(隐藏维度1152,参数4.49亿)下,应避免将LayerScale和门控MLP结合使用。对于具有更高隐藏维度或更多参数的更大规模模型,过度稀疏的问题可能会得到缓解,但我们将该体制的系统性研究留待未来工作。因此,ViT-5使用带有GeLU激活函数的原始MLP设计。

3.4. 位置编码

标准的ViT采用可学习的绝对位置编码(APE),这已被证明在复杂的视觉推理任务中缺乏显式的相对位置建模能力,并且在处理动态输入分辨率时存在固有限制。基于这些发现,我们将旋转位置嵌入(RoPE)扩展到2D设置,并将其融入我们的模型中。重要的是,我们没有丢弃绝对位置编码;相反,ViT-5联合采用了APE和2D RoPE。这样做的动机是,仅使用相对位置编码可能会引入不希望的不变性。例如,在仅使用RoPE的公式下,图像的块级翻转将变得完全不变,如图2所示,两个输入图像被视为等价。虽然这种行为对图像分类等简单任务影响有限,但对于一个通用的视觉主干来说,它可能会带来潜在的限制,因为在更复杂的视觉推理中,绝对空间线索可能至关重要。

图3比较了ViT-5和DeiT-III处理动态输入分辨率的能力。两个模型都在224×224分辨率下训练,然后在不同测试分辨率下进行评估,无需微调。结果表明,仅依赖绝对位置编码的DeiT-III仅在接近训练分辨率时表现最佳,其准确率随着输入尺寸的增加而迅速下降。相比之下,ViT-5表现出显著更强的分辨率鲁棒性。在从224到512的广泛输入范围内,ViT-5没有表现出明显的性能下降。此外,在一个较宽的分辨率跨度内(例如,对于ViT-5-L,从128到384),增加输入尺寸持续带来性能提升,这表明其改进的可扩展性和对动态分辨率的泛化能力。

3.5. 寄存器令牌

"视觉Transformer需要寄存器"这一近期发现表明,ViT中出现的伪影可以通过附加在输入块令牌之后的额外可学习令牌来有效解决。事实上,寄存器的实际作用超出了其最初的动机。这些可学习令牌提供了一个灵活的表示空间,支持如图像的元查询和1D令牌化等扩展。为了利用这些特性,我们将寄存器作为ViT-5的默认组件引入。

我们观察到,在配备RoPE的ViT中,寄存器令牌也应该被赋予相对位置嵌入。这是因为当向量旋转操作仅应用于块令牌而寄存器令牌保持不变时,寄存器自然与经历较小旋转角度的块表现出更低的余弦相似度。这种不平衡会扭曲注意力分布,并隐式地引入不希望的位置偏差。为了解决这个问题,我们为寄存器令牌配备了一个独立的2D RoPE,其频率基数显著高于用于块令牌的频率基数。这种设计使得寄存器和块令牌在不同通道维度上产生不同的旋转行为,有效地解耦了它们的位置相关性,并消除了不希望出现的偏差。通过实验,这种修改稳定了寄存器-块的交互,并提高了整体表示质量,详细的性能比较见表3。

我们还在图4中对DeiT-III和ViT-5的注意力激活进行了定性比较。与之前的观察一致,寄存器令牌有效地抑制了注意力图中的背景伪影,使类别令牌能够更准确地关注图像中语义上有意义的区域。在ViT-5中,我们观察到明显更清晰、更集中的特征图,我们将其归因于寄存器令牌和相对位置嵌入的共同作用。这两个组件在增强视觉Transformer的空间建模能力方面都发挥着重要作用。更多可视化结果见附录。

3.6. QK-归一化

最新的LLM,如Qwen3和Gemma3,已经开始通过应用额外的归一化来改造自注意力机制,即对查询和键进行归一化。形式上,这种QK-归一化机制定义为:

我们发现QK-归一化为ViT带来了适度的性能提升,但更重要的是,它能显著增强训练稳定性,并减少优化过程中出现的尖锐损失尖峰(对比见图5)。为了利用这些鲁棒性优势,我们将QK-Norm作为ViT-5的默认组件引入。

3.7. QKV的偏置项

我们在整个ViT-5模型中使用无偏置的RMSNorm,包括所有预归一化和QK-Norm层。这突显了自注意力更依赖于加权投影而非加性偏置的观点。因此,我们移除了QKV投影层中的偏置项,以保持结构一致性。这种修改使得QK-Norm能够更有效地发挥作用,并带来了显著的性能提升,详见第4.4节。

4. 实验

4.1. 图像分类

遵循先前实践,我们采用DeiT-III的配方并进行小幅修改,在ImageNet-1k上从头开始训练所有模型。详细的实现见附录。模型配置总结在表4中。如表5所示,在可比的参数量和计算预算下,ViT-5始终实现了最先进的性能。值得注意的是,其性能优势随着模型尺寸和输入分辨率的增加而平稳扩展。特别是,ViT-5-L在384x384输入下达到了86.0%的测试准确率,显著超过了之前的ViT最先进水平(85.4%)和基于CNN的模型(85.5%)。这些结果表明,通过系统地改进ViT的组件,可以实现显著且可扩展的性能提升。由于ImageNet-1k仍然是评估视觉主干最广泛使用的基准,在该数据集上观察到的一致改进有力地证明了ViT-5提供了增强的表示学习能力,并能作为一个有竞争力的、通用的视觉主干来替代现有架构。

4.2. 图像生成

我们通过将ViT-5作为扩散Transformer的主干进行训练来评估其迁移能力。我们遵循与SiT相同的配置,用ViT-5替换普通的ViT主干,并在ImageNet-256上进行训练和评估。实现细节见附录。

如表6所示,在相同的训练配置下,仅用ViT-5替换原始ViT主干,就在多个模型尺寸上显著优于DiT和SiT。在表7中,我们进一步将训练扩展到700万步,观察到ViT-5在FID、Inception Score以及精确率和召回率指标上始终优于其ViT对应模型。图6展示了不同模型尺寸在不同训练时长下的扩散模型缩放曲线。我们观察到,使用ViT-5作为主干的模型始终优于使用普通ViT的模型,并且随着训练的进行表现出平滑且稳定的缩放行为。综上所述,这些结果表明ViT-5在视觉理解和图像生成任务中都能有效泛化,凸显了我们组件级架构现代化的广泛影响。通过系统地改进核心Transformer组件,ViT-5成为了适用于多种视觉任务的强大且通用的主干。

4.3. 密集预测

我们进一步使用UperNet框架在ADE20k上评估ViT-5的语义分割性能。所有模型均在512×512分辨率下训练160k次迭代,主干网络在ImageNet-1k上预训练的轮数相同。所有其他训练设置在不同方法间保持一致,技术细节可在附录中找到。如表8所示,ViT-5在所有模型规模上始终优于DeiT-III。具体来说,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能够有效地迁移到密集预测任务,并在不引入特定任务架构更改的情况下带来一致的增益,进一步支持其作为一个强大且通用的视觉主干的地位。

4.4. 消融研究

在这项工作中,我们研究了ViT中的七个架构组件及其相应的设计选择。由于详尽枚举所有可能的组合代价过高,我们聚焦于两个互补的消融设置。首先,我们将ViT-5与现有视觉模型或语言模型中普遍采用的Transformer配置进行比较,以评估先前工作中探索和广泛使用的设计选择对于视觉任务是否最优。其次,从完整的ViT-5模型开始,我们逐个消融七个组件中的每一个,以分离并量化它们各自对整体性能的贡献。

与现有设计的比较。表9将ViT-5与源自流行视觉Transformer和现代语言模型的架构配置进行了比较。为了在更大规模和更高分辨率下充分评估模型容量,本次消融使用了384×384分辨率的ViT-Large。总结来说,在面向视觉的设计中,诸如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来说是未充分优化的,并激发了原则性的、组件级的现代化研究的动机。

表9. 现有的设计选择对于ViT来说仍未充分优化。我们将ViT-5与流行的视觉模型(如DeiT-III和DINO v2/v3)的设置进行比较。我们还尝试将从语言模型(如LLaMA、Qwen和Gemma系列)适配的配置迁移到ViT中(RoPE扩展到2D)。然而,与ViT-5相比,所有这些情况都表现出明显的性能差距。*表示使用后归一化而非显式LayerScale的模型。Acc. 是ViT-Large在384分辨率下的ImageNet-1k top-1准确率。

表10. 单组件变化的消融研究。我们报告了多种模型大小在ImageNet-1k上的top-1准确率。Δ表示相对于完整ViT-5模型的平均准确率差异。*表示模型受到过度门控问题的影响。

单个组件的影响。表10报告了在不同模型大小的ViT-5上进行单组件消融的结果。总体而言,与完整的ViT-5配置相比,移除任何单个组件都会导致准确率持续下降,证实了每个设计选择都对最终性能有积极贡献。虽然性能下降的幅度因组件和模型大小而异,但没有一个修改在所有设置中占主导地位,这突显了所采用的架构改进的互补性。

我们还观察到,各个架构组件的影响在不同模型规模上并不均匀。某些组件对较小模型的影响更强,而其他组件随着模型规模的增长变得日益关键。例如,在没有适当稳定的情况下,将GeLU替换为SwiGLU会导致Small模型性能显著下降,反映了紧凑模型对门控引起的稀疏性更高的敏感性。相比之下,LayerScale和2D RoPE等组件在更大的模型上表现出更显著的效果,其中更深的网络和更高的表示能力放大了训练稳定性和相对位置建模的重要性。值得注意的是,移除Registers和QK-Norm对Small模型的影响相对较小,但对Base和Large变体导致了一致且更大的下降,表明这些组件随着模型容量的增加发挥着更重要的作用。

组件细节。表11分析了ViT-5对几个关键架构超参数的敏感性。我们观察到,LayerScale的初始化值对性能有适度影响,对于base大小的ViT-5模型,初始化值过小会使准确率下降0.15%。ViT-5默认遵循大多数LLM的1D RoPE实践,将频率基数设置为1e-5,此处表现出最佳性能。在本工作的所有ViT-5模型中,我们使用了四个可学习的寄存器。我们观察到寄存器的数量对预测性能影响甚微,且对不同设置表现出良好的鲁棒性。

5. 结论

在本文中,我们对视觉Transformer进行了系统性的现代化改造,并表明在不重新设计整体架构的情况下,仔细的组件级更新可以释放巨大的潜力。通过更新归一化和激活函数的选择,加强用于空间推理的位置编码,并引入轻量级门控和可学习令牌,ViT-5成为一种实用的下一代ViT,它更易于优化,在各种视觉场景中更加鲁棒。大量的实验验证了在图像分类、生成和语义分割上的一致改进。这些增益可以跨任务迁移,表明其好处并非针对特定基准,而是反映了改进的表示和空间归纳偏置。更广泛地说,我们的结果表明,ViT的演进可以遵循与现代大型语言模型主干相同的模块化、最佳实践驱动的轨迹,其中稳定性和性能来自于原则性的组件选择,而非日益复杂的宏观架构。我们希望ViT-5能为2020年代中期的视觉和多模态系统提供一个强大、兼容的基础,并鼓励社区将主干现代化视为一个一流的、可复现的设计过程。

表11. 架构超参数的消融。我们改变LayerScale初始化值、2D RoPE的频率基数以及可学习寄存器的数量,并报告ViT-5-B在ImageNet-1k上的top-1准确率(%)。突出显示默认设置。

附录

A.1. 图像生成示例

图7展示了我们模型生成的定性示例。从视觉上看,生成的图像具有较高的感知质量,全局结构连贯,物体形状保持良好,纹理细节精细。结果看起来视觉上令人愉悦且自然,没有明显的伪影或空间不一致性。我们将这种强大的视觉质量不仅归因于模型改进的表示能力,还归因于其增强的空间感知能力。更好的空间建模使模型能够更准确地捕捉长程依赖关系和局部细节,从而产生更一致的布局和更清晰的纹理。这些特性共同使模型能够生成具有更丰富细节和更高视觉保真度的图像。

A.2. 架构设计原则

我们的设计遵循一套高层次原则,旨在通过组件级升级而非引入新的架构范式来实现视觉Transformer的现代化。主要目标是在保持ViT强大泛化能力的同时提高表示能力。详细原则包括以下几个方面:首先,我们保留原始的普通ViT结构,以确保广泛的任务适用性和与多模态系统的易于集成性。为此,我们避免在网络中进行任何形式的空间下采样。除了初始的块嵌入步骤外,所有中间特征图在整个模型中保持固定的空间分辨率。其次,我们使用自注意力作为唯一的令牌混合机制。虽然先前的工作表明,引入卷积可以引入更强的归纳偏置并提高标准基准的性能,但我们有意放弃了这些潜在增益,以换取架构的简洁性和通用建模能力。第三,我们不专注于改进块嵌入层,而是采用标准的、带有线性投影的非重叠块嵌入。这一选择的动机是,许多现代面向生成的视觉模型不再依赖块嵌入进行令牌化,而是使用专用的VAE或学习到的令牌化器。因此,对块嵌入的改进主要有利于像素空间的理解任务,与我们最大化通用性的目标不完全一致。

A.3. 更多注意力可视化

图8可视化了DeiT-III和ViT-5在相同实验设置下局部令牌的注意力图,与图4不同之处在于,我们从一个选定的局部令牌而非类别令牌出发进行关注。两个模型都在ImageNet-1K上预训练了相同的时间,并在384×384分辨率下进行了微调;注意力可视化是在相同分辨率下生成的。与DeiT-III相比,ViT-5展现出明显更清晰、更有结构的注意力模式。ViT-5中局部令牌的注意力在空间上更连贯,与语义上有意义的区域对齐得更好,而背景伪影则被大量抑制。相比之下,DeiT-III通常表现出更分散、更嘈杂的激活。结合正文中的类别令牌注意力可视化,这些结果定性地表明,ViT-5在空间建模能力上取得了实质性的改进,超越了数值上的性能增益。更清晰的注意力图表明了更精确的令牌交互和更可靠的空间推理,突显了ViT-5中引入的组件级现代化的有效性。

A.4. 技术细节

ImageNet分类。对于ImageNet-1k分类,ViT-5采用两阶段训练流程,包括大规模预训练和随后的微调。总体而言,我们的训练方案在很大程度上遵循DeiT-III的方案,但为适应ViT-5引入的架构变化进行了少量修改。在预训练阶段(表12),模型使用LAMB优化器从头开始训练,采用大批量大小和余弦学习率调度。我们采用DeiT风格训练中常用的标准数据增强,包括随机裁剪、水平翻转、Mixup和CutMix,同时禁用标签平滑和dropout。根据模型规模应用随机深度,并启用梯度裁剪以保持训练稳定性。在微调阶段(表13),我们切换到AdamW优化器,使用较小的学习率进行短时间训练。微调在224和384分辨率下进行(如适用),遵循常见实践。与预训练相比,微调使用简化的增强策略,启用标签平滑并调整随机深度率。这些设置确保了公平且强大的训练基线,同时允许ViT-5架构改进的效果独立于激进的方案变化进行评估。

图像生成。对于图像生成,我们遵循SiT的训练协议,并用ViT-5替换普通的ViT主干。除非另有说明,所有超参数和优化设置与SiT保持一致,以确保公平比较。特别地,模型使用与原始SiT论文中相同的优化器、批量大小、学习率调度和数据预处理进行训练。我们考虑两种训练机制。首先,我们进行短时间训练,在不同模型大小上最多进行80万步,以评估效率和早期缩放行为。其次,为了评估长期可扩展性,我们对ViT-5-XL进行了700万步的大规模训练。这个长训练实验采用1.5的无分类器指导尺度,并使用随机微分方程公式,遵循SiT的设置。

语义分割。对于语义分割,我们使用UperNet框架在ADE20K上评估ViT-5。我们的训练协议在很大程度上遵循先前基于ViT的分割工作,并在不同主干之间保持一致以确保公平比较。所有模型均从ImageNet-1k预训练检查点初始化,并在512×512分辨率下训练160k次迭代。我们使用密集预测任务的标准数据增强策略,包括随机调整大小、裁剪和水平翻转。使用AdamW优化器和余弦学习率调度进行优化,并根据模型规模应用随机深度。当用ViT-5替换普通ViT主干时,没有引入任何特定于任务的架构修改。

相关推荐
lijianhua_97121 小时前
国内某顶级大学内部用的ai自动生成论文的提示词
人工智能
EDPJ1 小时前
当图像与文本 “各说各话” —— CLIP 中的模态鸿沟与对象偏向
深度学习·计算机视觉
蔡俊锋1 小时前
用AI实现乐高式大型可插拔系统的技术方案
人工智能·ai工程·ai原子能力·ai乐高工程
自然语1 小时前
人工智能之数字生命 认知架构白皮书 第7章
人工智能·架构
大熊背1 小时前
利用ISP离线模式进行分块LSC校正的方法
人工智能·算法·机器学习
eastyuxiao2 小时前
如何在不同的机器上运行多个OpenClaw实例?
人工智能·git·架构·github·php
诸葛务农2 小时前
AGI 主要技术路径及核心技术:归一融合及未来之路5
大数据·人工智能
光影少年2 小时前
AI Agent智能体开发
人工智能·aigc·ai编程
Database_Cool_2 小时前
OpenClaw-Observability:基于 DuckDB 构建 OpenClaw 的全链路可观测体系
数据库·阿里云·ai
ai生成式引擎优化技术2 小时前
TSPR-WEB-LLM-HIC (TWLH四元结构)AI生成式引擎(GEO)技术白皮书
人工智能