Perceiver IO: 一种适用于任何模态的可扩展、全注意力模型

Perceiver IO 是一种可扩展的神经网络架构，它能够在任何模态（包括文本、图像、音频、视频和点云）上使用 Transformer 的自注意力机制，而不会产生通常与高维数据相关的二次方计算和内存成本。通过在小型、固定大小的潜空间（latent space）中执行大部分计算，Perceiver IO 消除了输入大小与自注意力层复杂度之间的依赖关系。

克服 Transformer 的扩展限制

标准的 Transformer 架构在计算和内存方面的扩展性较差，因为其自注意力机制在每一层都需要对所有输入进行成对的点积。为了处理高维数据，现有模型通常依赖于特定模态的预处理，将数据离散化为 token。例如，ViT 使用图像 patch，Wav2Vec2 使用音频特征编码器。

Perceiver IO 通过对一组潜变量（latent variables）而非输入本身应用自注意力来解决这个问题。输入仅用于与这些潜变量进行交叉注意力（cross-attention）计算。这确保了 Transformer 编码器的复杂度与输入大小呈线性关系，而潜注意力则保持独立于输入大小。Perceiver IO 通过在输出端使用类似的交叉注意力机制进一步扩展了这一点，允许模型处理任意形状的输出。

核心架构与工作流程

所有 Perceiver 变体都构建在 PerceiverModel 类之上，该类利用三个可选组件来处理不同的数据类型：

Preprocessor：将原始输入（文本、图像等）嵌入为向量。
Decoder：将潜变量的最终隐藏状态解码为有用的格式，例如分类 logits 或光流。
Postprocessor：将解码器输出转换为特定特征，主要用于自动编码任务。

处理流水线

Input Stage：输入可以选择性地通过 preprocessor 进行处理。
Cross-Attention：潜变量产生查询（Q），而预处理后的输入产生键（K）和值（V）。这会将高维输入映射到潜空间。
Latent Processing：一个可重复的自注意力层块更新潜嵌入。输出是潜变量的“最后隐藏状态”张量。
Output Stage：一个可选的 decoder 执行另一个交叉注意力操作，其中可训练的嵌入产生查询（Q），而潜变量产生键（K）和值（V）。随后，postprocessor 可能会对这些输出进行精细化处理。

特定模态的实现

文本分类

Perceiver IO 可以直接处理原始 UTF-8 字节，从而消除了对显式分词（如 WordPiece 或 BPE）的需求。对于文本分类，PerceiverTextPreprocessor 嵌入字节 ID 并添加绝对位置嵌入。随后，PerceiverClassificationDecoder 使用交叉注意力将潜变量映射到分类 logits。在掩码语言建模（PerceiverForMaskedLM）中，PerceiverBasicDecoder 将潜变量映射回输入序列长度以预测缺失的字节。

图像分类

对于视觉任务，模型使用 PerceiverImagePreprocessor。作者测试了三种预处理方法：

使用 1x1 卷积层和学习到的 1D 位置嵌入来展平像素。
使用固定的 2D Fourier 位置嵌入来展平像素。
使用 2D 卷积 + maxpool 层结合固定的 2D Fourier 位置嵌入。

在 JFT 上进行预训练后，PerceiverForImageClassificationConvProcessing 变体在 ImageNet 上达到了 84.5% 的 top-1 准确率。值得注意的是，仅使用 1D 学习到的位置嵌入的版本达到了 72.7% 的准确率，这证明了该模型在没有显式 2D 结构知识的情况下执行任务的能力。

光流估计

Perceiver IO 可以估计两帧图像之间的 2D 像素位移。PerceiverForOpticalFlow 模型提取每个像素周围的 3x3 的 patch，将帧在通道维度上进行拼接，并展平空间维度。随后，PerceiverOpticalFlowDecoder 使用预处理后的输入作为查询（Q），将潜变量解码回预测的光流图。这种方法在 AutoFlow 数据集上训练后，在 Sintel 和 KITTI 基准测试中取得了最先进的结果。

多模态自动编码

使用 PerceiverForMultimodalAutoencoding 模型，该架构可以学习跨多个模态的联合分布。例如，在 Kinetics-700 数据集上，它同时处理视频帧、音频样本和类别标签：

Images：通过 "space to depth" 变换并与 2D Fourier 嵌入拼接。
Audio：与固定的 Fourier 位置嵌入拼接。
Labels：通过 one-hot 预处理器处理。

这些模态被填充（padding）到统一的通道维度并进行拼接。随后，解码器以分块（chunks）形式对数据进行自动编码，以重建原始模态。该模型在视频分类任务中达到了 45% 的 top-1 准确率，同时保持了 20.7 的 PSNR 用于视频重建。

更广泛的应用与多功能性

除了主要模态外，Perceiver IO 架构在其他复杂领域也取得了成功：

3D 点云：在 ModelNet40 上训练后，模型达到了 85.7% 的 top-1 准确率，可以与 PointNet++ 等专门的模型相媲美。
强化学习：在 StarCraft II 中，Perceiver 替换了 AlphaStar 中的原始 Transformer，在无需额外参数微调的情况下，对阵 Elite bot 保持了 87% 的胜率。

由于该架构是任务无关的，它可以通过设计新的预处理器和解码器来适应任何问题，从而有可能将各种模态统一到同一个共享的潜空间中。