Perceiver IO: 一種適用於任何模態的可擴展、全注意力模型

Perceiver IO 是一種可擴展的神經網路架構，它能夠在任何模態（包括文本、圖像、音訊、影片和點雲）中使用 Transformer 的自注意力機制，而不會產生通常與高維數據相關的二次方計算與記憶體成本。透過在一個小型且固定大小的潛在空間（latent space）中進行大部分計算，Perceiver IO 消除了輸入大小與自注意力層複雜度之間的依賴關係。

克服 Transformer 的擴展限制

標準 Transformer 架構在計算與記憶體方面的擴展性較差，因為其自注意力機制在每一層都需要對所有輸入進行成對點積。為了處理高維數據，現有模型通常依賴於特定模態的預處理，將數據離散化為 token。例如，ViT 使用圖像補丁（image patches），Wav2Vec2 使用特徵編碼器處理音訊。

Perceiver IO 透過對一組潛在變數（latent variables）而非輸入本身進行自注意力運算來解決此問題。輸入僅用於與這些潛在變數進行交叉注意力（cross-attention）運算。這確保了 Transformer 編碼器的複雜度與輸入大小呈線性關係，而潛在注意力則保持與輸入大小無關。Perceiver IO 進一步透過在輸出端使用類似的交叉注意力機制，讓模型能夠處理任意形狀的輸出。

核心架構與工作流程

所有 Perceiver 變體都建立在 PerceiverModel 類別之上，該類別利用三個選用的組件來處理不同的數據類型：

Preprocessor：將原始輸入（文本、圖像等）嵌入為向量。
Decoder：將潛在狀態的最終隱藏狀態解碼為實用的格式，例如分類邏輯值（logits）或光流（optical flow）。
Postprocessor：將解碼器輸出轉換為特定特徵，主要用於自動編碼任務。

處理流程

Input Stage：輸入可選擇性地由 preprocessor 進行處理。
Cross-Attention：潛在變數產生查詢（Q），而預處理後的輸入產生鍵（K）與值（V）。這將高維輸入映射到潛在空間。
Latent Processing：一個可重複的自注意力層區塊會更新潛在嵌入。輸出是一個包含潛在變數「最後隱藏狀態」的張量。
Output Stage：一個選用的 decoder 執行另一個交叉注意力運算，其中可訓練的嵌入產生查詢（Q），而潛在變數產生鍵（K）與值（V）。接著，postprocessor 可能會進一步精煉這些輸出。

模態特定實作

文本分類

Perceiver IO 可以直接處理原始的 UTF-8 位元組，從而消除了對顯式 token 化（如 WordPiece 或 BPE）的需求。對於文本分類，PerceiverTextPreprocessor 會嵌入位元組 ID 並加入絕對位置嵌入。接著，PerceiverClassificationDecoder 會使用交叉注意力將潛在變數映射到分類邏輯值。在遮罩語言模型（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 補丁，沿著通道維度連接圖像幀，並展平空間維度。接著，PerceiverOpticalFlowDecoder 會使用預處理後的輸入作為查詢（Q），將潛在變數解碼回預測的光流圖。這種方法在 Kinetics-700 數據集上訓練後，於 Sintel 和 KITTI 基準測試中取得了尖端（state-of-the-art）的結果。

多模態自動編碼

使用 PerceiverForMultimodalAutoencoding 模型，該架構可以學習多個模態之間的聯合分佈。例如，在 Kinetics-700 數據集上，它可以同時處理影片幀、音訊樣本和類別標籤：

Images：透過「space to depth」轉換並與 2D Fourier 嵌入連接。
Audio：與固定的 Fourier 位置嵌入連接。
Labels：透過 one-hot 預處理器處理。

這些模態會被填充（padded）到統一的通道維度並進行連接。接著，解碼器會以分塊（chunks）形式自動編碼數據，以重建原始模態。該模型在影片分類中達到了 45% 的 top-1 準確率，同時在影片重建方面保持了 20.7 PSNR。

廣泛的應用與多功能性

除了主要的模態之外，Perceiver IO 協助在其他複雜領域也展現了成功：

3D Point Clouds：在 ModelNet40 上訓練後，模型達到了 85.7% 的 top-1 準確率，足以與 PointNet++ 等專門化模型競爭。
Reinforcement Learning：在 StarCraft II 中，Perceiver 取代了 AlphaStar 中的原始 Transformer，在無需額外參數調整的情況下，對抗 Elite bot 時維持了 87% 的勝率。

由於該架構是任務無關的（task-agnostic），因此可以透過設計新的預處理器與解碼器，將其適應於任何問題，並潛在於將各種模態統一到一個共享的潛在空間中。