Reading

AIM系列--大视觉模型的自回归预训练

Apr 16, 2026
155 views
Self-Supervised
#Autoregressive #CV

AIM v1

论文名称：Scalable Pre-training of Large Autoregressive Image Models (Arxiv 24.01)

论文地址：arxiv.org/pdf/2401.08541.pdf

代码地址：github.com/apple/ml-aim

简介

自回归预训练的大视觉模型会不会像 LLM 一样有缩放性质？

大语言模型 (LLM) 的革命性发展使得与任务无关的预训练成为自然语言处理任务的主流。大语言模型可以解决复杂的推理任务，遵从人类的指令，并且成为人工智能助手。LLM 成功的一个非常关键的因素是：随着模型容量和数据量的扩增，带来的模型能力的持续提升。

为什么大语言模型具有缩放性质？作者提出以下2点原因：

  1. 即使这些模型只使用最简单的目标函数进行 Next Token Prediction 的自回归预训练，它们也可以在复杂的上下文中学习到复杂的范式。
  2. 大语言模型的缩放性质是一些工作在 Transformer 架构中发现的，这也暗示了自回归预训练与 Transformer 架构之间的协同关系。

本文探索的目标是：自回归预训练和 Transformer 可否适用于视觉模型，即 LLM 的缩放性能是否依然在视觉领域里奏效。

预训练数据集

AIM 预训练数据集使用的是 DFN。该数据集是基于 12.8B 的图文对 DataComp。

  1. 从 Common Crawl 过滤得到。
  2. 数据预处理环节去除了 NSFW 内容，模糊人脸，并通过对评估集进行重复数据删除来减少污染。
  3. 使用 data filtering network 根据图像与其对应描述的分数进行排名。
  4. 保留前 15% 的样本，从 DataComp 12.8B 数据集中提取了 2B 图像的子集，称为 DFN2B。

注意，除了隐私和安全过滤器之外，此过程不包括任何基于图像内容的额外管理。

由于 AIM 模型的预训练过程不需要文本，因此 AIM 模型可以使用非配对的更大的图像数据集进行训练，也可以使用图文对应分数不那么好的其他数据来训练。

受 LLM 预训练常见做法的启发，作者在预训练期间对维基百科和书籍等高质量数据源进行过采样。作者以 \(p=0.8\) 的概率从 DFN-2B 中采样图像，并以 \(p=0.2\) 的概率从 ImageNet-1K 中采样图片，最终得到的这个数据集称之为 DFN-2B+。

自回归的训练目标

输入图片首先按照 ViT 的方式划分成 \(K\) 个互不重叠的 Patch，作者假设这些 Patch 序列的顺序是固定的，除非另有说明，否则默认使用 raster order (行主排序)。

给定上述顺序，图像的概率可以分解为 Patch 条件概率的乘积：

\[\begin{equation} P(x) = \prod_{k=1}^K P(x_k~|~x_{<k}), \end{equation} \tag{1}\]

其中，\(x_{<k}\)  表示前 \(k-1\) 个图像 Patch 的集合，前 \(k-1\) 个图像 Patch 被用来预测第 \(k\) 个图像 Patch。与 LLM 模型不同的是， \(K\) 值是固定的，不是变化的，因此不需要截断上下文长度。训练的目标函数是 negative log-likelihood (NLL) 损失函数：

\[\begin{equation*} \sum_{x\in\mathcal{X}} \sum_{k=1}^K -\log P(x_k~|~x_{<k}). \end{equation*} \tag{2}\]

在无限多数量的图像上最小化这个目标在理论上等同于学习真实的底层图像分布。

训练的目标函数是归一化的 Pixel-level 的回归损失函数。

\[\begin{equation} \min_{\theta} \frac{1}{K}\sum_{k=1}^K \|\hat{x}_k(\theta)-x_k\|_2^2. \end{equation} \tag{3}\]

式中，\(\hat{x}_k(\theta) \) 是参数为 \(\theta\) 的模型的预测结果， \(x_k\) 是对应的图片真实值。

除此之外，作者还考虑了一种使用离线 tokenizer 的 cross-entropy 损失函数。但是消融实验表明这种损失函数得到的特征不如归一化的 Pixel-level 的回归损失函数。

模型架构

预训练时期 (Prefix Transformer，和 NLP 中的 T5 模型类似)

AIM 模型使用 ViT + Prefix Mask Self-Attention，不同容量的模型的配置如下图所示。

PixPin_2026-04-17_11-31-59.png

AIM 模型架构如下图所示。

PixPin_2026-04-17_11-32-27.png

在预训练时期，对 Self-Attention 进行 Casual Mask。本来在 Self-Attention 中，对于 Embedding 的计算方式是：

\[\begin{equation} y_i = \sum_{k=1}^K a_{ik} v_i, \end{equation} \tag{4}\]

式中，\(a_{ik} \)  是 attention weight， \(v_i \)是 Self-Attention 的 Value 值。

现在，Casual Mask 使得 \(k>i \)时的 \(a_{ik}=0\) 且 \(\sum_{k=1}^K a_{ik} = 1 \)。

  • 技术细节1 (Prefix causal attention)：

但是，作者在这里认为 Casual Mask 会影响下游任务的质量，因为这里与下游任务中ViT模型的标准用法不同，下游任务中通常采用双向自注意力。无论在下游适应期间是否保留因果掩码，这种差异都会导致性能下降。为了解决这个问题，作者像 T5一样使用 Prefix Attention。(注意：这里就不再是真正的 Casual Transformer 了，准确讲应该叫 Prefix Transformer)。

随机均匀采样前缀长度 (Prefix length)  \(S\)。前面的 \(S\) 个 Patch ( \(S\in [1, K-1]\) ) 采用 Bi-directional Self-attention，不计算损失函数。对于所有 \(K-S\) 的 Patch 使用  。后面的  个 Patch 采用 Casual Self-attention，计算损失函数。

这种修改有助于模型在没有 Casual Mask 的情况下工作，使其能够在下游适应期间被去除。这种方法提高了模型在下游任务中的性能。

AIM 模型在下游任务推理时是 Bi-directional Transformer，和 ViT 一致。

具体来说，前缀注意力的具体执行过程可以分为以下几个清晰的步骤：

    1. 动态划分“前缀”与“预测区”

假设一张输入的图像被切分成了 \(K\)个图像块（Patches）。在每次训练迭代时，模型会做一次随机抽样：

      • 从均匀分布 \(U\{1, 2, ..., K-1\}\) 中随机抽取一个整数 \(S\)。
      • 这个 \(S\) 就是当前的前缀长度（Prefix Length）。
    1. 构建混合注意力掩码 (Attention Mask)

根据抽样得到的 \(S\)，模型会将这 \(K\) 个图像块分为两组，并施加不同的注意力规则：

      • 前缀区（第 \(1\) 到 \(S\) 个 Patch）：这部分图像块之间是完全双向可见的。它们可以自由地相互关注，就像标准的 ViT 或 BERT 一样，从而提取出极其丰富的局部上下文特征。
      • 自回归区（第 \(S+1\) 到 \(K\) 个 Patch）：这部分图像块受到因果掩码（Causal Mask）的限制。第 \(k\) 个 Patch 只能看到前缀区的所有内容，以及它自己和它之前的 Patch（即 \(1\) 到 \(k\)），无法看到 \(k+1\) 及之后的内容。
    1. 针对性的损失计算

因为前缀区的 Patch 已经“互相看光了”，如果让模型去预测它们，任务就太简单了。因此，像素回归损失（\(L_{img}\)）仅仅在非前缀区（即 \(k > S\) 的图像块）上计算。模型必须努力根据双向编码的前缀特征，去自回归地预测剩余的图像块。

PixPin_2026-04-17_11-40-22.png

引入前缀注意力不仅是一个工程技巧，它在特征学习上带来了两个巨大的收益：

    • 无缝衔接下游推理：因为在预训练时，前缀长度 \(S\) 是随机变化的（有时很短，有时几乎覆盖全图），模型已经完全习惯了处理双向可见的图像块。当我们在下游任务（如 ImageNet 分类）中冻结这个视觉编码器时，我们只需将整张图片作为前缀（\(S = K\)）输入，模型就能完美适应，无需任何额外的结构调整。
    • 逼迫模型提取高密度信息：论文中的消融实验表明，使用前缀注意力比完全使用双向注意力效果更好。研究团队推测，这是因为前缀注意力迫使模型从“残缺的局部图像（前缀）”中提取出最具信息量的上下文（Maximally informative contexts），因为后续的视觉和文本 Token 都极度依赖这些前缀特征来进行预测。
  • 技术细节2 (对 Prediction head 的改进)：

在预训练期间采用 MLP head 是自监督预训练的一种常见做法。MLP head 一般在迁移到下游任务时被丢掉。MLP head 的作用是防止主体模型在预训练目标的作用下变得过于专业化，从而提高它们对于下游任务的适配性。作者选择了一种简单的 MLP head 的设计思路：使用 \(N\) 个块作为 MLP head，独立地处理每个 Patch。作者观察到这样的设计在性能和成本之间取得了很好的权衡。

  • 下游任务迁移学习时期 (Bi-directional Transformer，和 CV 中的 ViT 模型类似)

在下游任务期间，所有模型权重都固定，只训练一个分类头。

在分类头中，AIM 模型还进行了一个 attention pooling 的操作。给定一系列的 Patch 的特征 \(P = \{p_{i} \mid 1 \leq i \leq K\}\) ，计算一个全局描述符 \(\hat{p}\) ：

\[\begin{equation} \hat{p_h} = \sum^{K}_{i=1} \frac{\text{exp}(q_h^{T} W^{k}_h p_{i})}{\sum_{j=1}^K\text{exp}(q_h^{T} W^{k}_h p_{j})} W^{v}_h p_{i}, \end{equation} \tag{5}\]

\(h = \{1,...,H\}\) 是 attention 的 head 数，\(q_h\)  是可学习的 query 向量。pooled 之后的特征变为\( \hat{p} = [p_1,..., p_H], \hat{p} \in R^d\) 。这就相当于分类头是一个 Attention 的操作。

AIM 模型的缩放性质结果

如下图所示是随着预训练过程的进行，validation loss 的值与下游任务 ImageNet-1K 精度的变化。而且，随着模型容量的增加，validation loss 和 ImageNet-1K 的性能都会改善，这一结果与 LLM 中观察到的趋势一致。

PixPin_2026-04-17_11-45-54.png

下图表示在 1M 图像的小数据集 (IN-1k) 或更大的 2B 图像集 (即 DFN-2B+) 上进行预训练时，验证集的损失。由于验证集也是 IN-1K，所以以它自己作为预训练数据集得到的 validation loss 比较低也就不足为奇了。然而，validation loss 在训练结束时会恶化，这表明对训练数据的过拟合。

PixPin_2026-04-17_11-46-58.png

当预训练数据集换成更大的 DFN-2B 时，没有出现过拟合的迹象。当相同的数据集用少量 IN-1K 数据增强，即预训练的数据集变为 DFN-2B+ 时，最终超过 IN-1K 预训练的性能。

如下图所示，作者研究了将预训练的长度从 500k 增加到 1.2M iteration 的影响，即训练过程中见到的数据从 2B 到 5B 大小。可以观察到，通过增加模型容量，或者增加预训练的数据集数量，都可以降低 validation loss，从而提高 AIM 模型的性能。值得一提的是，当总的 FLOPs 一致时，大一点的模型 + 小一点的数据集训练出的模型效果与小一点的模型 + 大一点的数据集训练出的模型效果接近。

PixPin_2026-04-17_11-47-28.png

AIM v2

这篇文章介绍了一种新的大规模视觉编码器预训练方法，称为AIMV2。其主要动机是扩展近期在视觉模型自回归预训练方面的进展，将其应用于多模态环境，即图像和文本结合的场景。AIMV2旨在通过简单的预训练过程，实现可扩展性，并在多种下游任务中表现出色。具体而言，AIMV2通过与多模态解码器配对，自回归生成原始图像块和文本标记，从而在多模态评估和视觉基准测试（如定位、基础、分类）中表现优异。

主要贡献包括：

  1. 设计简单且可扩展的预训练框架：AIMV2通过自回归生成图像块和文本标记的方式进行预训练，这种方法不仅易于实现，而且无需使用过大的批量大小或专用的批间通信方法。
  2. 与大型语言模型（LLM）应用的无缝集成：AIMV2的架构和预训练目标与LLM驱动的多模态应用高度契合，提供了密集的监督信号。
  3. 在多模态和视觉任务中的卓越表现：AIMV2在多模态理解基准上优于当前最先进的视觉-语言预训练方法，并在开放词汇物体检测和指代表达理解等任务中表现突出。
  4. 强大的可扩展性：AIMV2在数据或参数扩展时表现出一致的改进，并与现代工具兼容，支持原生图像分辨率和零样本识别适应性。

这些贡献表明，AIMV2不仅在技术上提供了一个简单而有效的预训练方法，还在实际应用中展示了其强大的通用性和适应性。

PixPin_2026-04-29_20-07-50.png
PixPin_2026-04-29_20-08-01.png

预训练

在AIMV2模型中，标准的单模态自回归框架被扩展到了多模态设置，能够将图像和文本整合成一个统一的序列。具体来说，一幅图像被分割成若干不重叠的patch，每个patch形成一个序列中的令牌。同理，文本序列被分解为子词，这些序列随后被连接起来，使得文本令牌能够关注到图像令牌。尽管图像和文本的连接顺序可以是双向的，但为了训练一个强大的视觉编码器，AIMV2始终将图像置于首位，从而增强对视觉特征的条件依赖。

这种方法导致了一个统一的多模态自回归建模过程，序列的因子化如下：

\[P(S) = \prod_{j=1}^{I+T} P(S_j | S_{ < j}) \\\]

其中，\(S_j\) 表示图像块和文本token的连接序列中的第 \(j\) 个token，\(S_{<j}\) 为前面所有的token

预训练设置包括一个专用的视觉编码器来处理原始图像块，这些图像块随后与嵌入的文本令牌一起传递给多模态解码器。解码器随后在组合序列上执行下一个令牌预测。

目标函数定义了图像和文本域的单独损失函数：

  • 对于图像域，使用的是像素级的回归损失：
\[L_{img} = \frac{1}{I} \sum_{i=1}^{I} \| \hat{x}_i(x_{ < i}; \theta) - x_i \|_2^2 \\\]
  • 对于文本域，使用的是标准的交叉熵损失：
\[L_{text} = - \frac{1}{T} \sum_{t=I+1}^{I+T} \log P(x_t | x_{ < t}; \theta) \\\]

整体目标是最小化 \(L_{text} + \alpha \cdot L_{img}\)，其中 \(\alpha\) 是一个权重系数。

架构

AIMV2的视觉编码器采用了Vision Transformer (ViT)架构，训练了一系列参数在300M到3B之间的视觉编码器。详细的模型规格在下面表格中给出。

  • 前缀注意力：在视觉编码器中引入前缀注意力掩码，允许在推理过程中使用双向注意力。
  • SwiGLU和RMSNorm：视觉编码器和多模态解码器采用SwiGLU作为前馈网络，并用RMSNorm替代所有归一化层。

多模态解码器同时执行图像和文本模态的自回归生成。图像特征和原始文本令牌分别线性投影并嵌入到统一的空间中。解码器接收图像和文本特征的连接序列作为输入，并在自注意力操作中使用因果注意力。解码器的输出通过两个独立的线性头处理，一个用于图像token，另一个用于文本token，以分别预测每种模态中的下一个令牌。

PixPin_2026-04-20_16-38-24.png

数据

AIMV2模型的预训练使用了公共和私有数据集，这些数据集包含了配对的图像和文本。公共数据集包括DFN-2B和COYO，同时还使用了一个高质量图像-文本对的私有数据集（HQITP）。除了alt-text，还使用了合成的标题。关于数据集的详细信息，包括大小和每个数据集的采样概率，在表格中有所提供。除非另有说明，所有AIMV2模型均使用了120亿（12billion）对图像-文本样本进行预训练。

后训练

虽然AIMV2的初始预训练阶段已经产生了性能优异的模型，但研究人员探索了各种后训练策略以进一步提升模型能力。

高分辨率适配

在初始预训练阶段，图像数据的分辨率固定为224像素。然而，许多下游任务（如检测、分割、多模态大模型）需要适应更高分辨率的图像。因此，研究人员对AIMV2模型进行了336和448像素分辨率的微调。高分辨率适配阶段使用了与预训练阶段相同池中的20亿对图像-文本样本，不同的是在这个阶段没有使用合成标题。与其他研究一致，研究人员发现零权重衰减对于保持优化的稳定性非常重要。

原生分辨率微调

为特定分辨率和长宽比训练模型可能对需要处理原始形状图像的应用不够灵活。之前的工作如FlexiViT和NaViT已解决了这些限制。研究人员采用了一种不同的方法来处理可变的长宽比和分辨率。

具体来说，为了在可变分辨率下保持训练的稳定性，作者巧妙地利用了“总计算量守恒”的原理。

首先定义三个核心变量：

  • \(B_i\)：当前 mini-batch \(i\) 中的图像数量（即动态 Batch Size）。
  • \(A_i\)：当前 mini-batch 中每张图像被划分的 Patch 数量。
  • \(C\)：当前 mini-batch 中包含的总 Patch 数量（这是一个预设的固定常数）。

具体的执行步骤如下：

  1. 采样目标面积：对于每一个 mini-batch \(i\)，系统首先随机采样一个目标面积 \(A\)。
  2. 等比缩放：将该批次内的图像在严格保持原始宽高比的前提下，缩放至刚好能放入面积 \(A\) 的尺寸。
  3. 动态平衡批次： 根据公式 \(C = A_i \times B_i\)，动态调整该批次的图像数量 \(B_i\)。

(注：这意味着如果本次采样的图像分辨率大，Batch Size 就会自动减小；反之则增大，从而确保显存占用和计算量始终稳定。)

目标面积 \(A\) 的采样数学模型

为了保证模型能学习到丰富多样的分辨率，目标面积 \(A\) 并非完全随机，而是遵循特定的概率分布：

计算公式为：\(A = 2^n\)，指数 \(n\) 的生成规则：首先从标准正态分布 \(\mathcal{N}(0, 1)\) 中采样，然后将其截断在 \([-1, 1]\) 的范围内，最后线性映射到 \([7, 12]\) 的区间。

Reference

自回归视觉生成模型解读 (四)：大视觉模型的自回归预训练