模型越大，性能越好？蘋果自回歸視覺模型AIM：沒錯

作者：機器之心 2024-01-18 10:57:35

視覺模型，同樣遵循「參數越多性能越強」的規律？剛剛，一項來自蘋果公司的研究驗證了這個猜想。

過去幾年，大規模預訓練模型在 NLP 領域取得了成功，這類模型可以通過幾個示例解決復雜的推理任務，也可以按照指令進行推理。

眾所周知的是，預訓練模型能取得成功的一個理論基礎是，隨著容量（即參數量）或預訓練數據量的增加，模型能夠不斷提升性能。

這很自然地引發了研究者們的聯想：使用自回歸目標對 Transformers 進行縮放的成功是否僅限于文本？

在最新的一篇論文《Scalable Pre-training of Large Autoregressive Image Models》中，蘋果的研究者提出了自回歸圖像模型（AIM），探討了用自回歸目標訓練 ViT 模型是否能在學習表征方面獲得與 LLMs 相同的擴展能力。

論文鏈接：https://arxiv.org/pdf/2401.08541.pdf
項目地址：https://github.com/apple/ml-aim

先說結論：研究者發現，模型容量可以輕松擴展到數十億個參數，并且 AIM 能夠有效利用大量未經整理的圖像數據。

他們利用包括 ViT、大規模網絡數據集和 LLM 預訓練最新進展在內的工具集，重新審視了 iGPT 等自回歸表征學習方面的前期工作，此外還引入了兩處架構修改，以適應視覺特征的自回歸預訓練。

首先，研究者并沒有像 LLM 通常那樣將自注意力限制為完全因果關系，而是采用了 T5 中的前綴注意力。這一選擇使得能夠在下游任務中轉向完全雙向的注意力。其次，研究者使用了參數化程度較高的 token-level 預測頭，其靈感來自對比學習中使用的預測頭。他們觀察到，這種修改大大提高了后續特征的質量，而在訓練過程中的開銷卻很小。總體來說，AIM 的訓練與最近的 LLM 訓練類似，而且不依賴于監督式或自監督式方法所需的任何 stability-inducing 技術。

隨后，研究者對一系列模型展開了研究，這些模型的參數從 600M 到 7B 不等，都是使用 20 億帶許可的未編輯圖像進行預訓練的。如圖 1 所示，以 15 個圖像識別基準的平均準確率來衡量，AIM 模型在與模型規模的關系上表現出很強的擴展性，容量越大的模型下游性能越好。更重要的是，驗證集上的目標函數值與后續凍結特征的質量之間存在相關性。這一觀察結果證明，自回歸目標足以滿足視覺特征的訓練要求。此外，隨著對更多圖像進行訓練，研究者還觀察到了下游性能的持續改善，且沒有飽和的跡象。總體而言，這些觀察結果與之前關于擴展大型語言模型的研究結果是一致的。

但同時引起注意的是，本文實驗所使用的模型規模有限，是否能在更大參數量級的模型上驗證此規律，有待進一步探索。

方法概覽

本文的訓練目標遵循應用于圖像 patch 序列的標準自回歸模型。更準確地說，圖像 x 被分割為 K 個不重疊 patch x_k 組成的網格 k ∈ [1, K]，這些 patch 共同形成 token 序列。

研究者假設所有圖像的序列順序是固定的，因此除非另有說明，他們默認使用光柵（行優先）排序。給定上述順序，一張圖像的概率可以被分解為 patch 條件概率的乘積。

預測損失。研究者的訓練目標自然會產生某些損失變體，每個變體對應分布 P (x_k | x_<k) 的選擇。他們還考慮通過使用離線 tokenizer，將帶有 patch 的交叉熵損失轉換為離散 token。消融實驗表明，這些設計是有效的，盡管不會產生像像素級損失那樣顯著的特征。

架構

在骨干選擇上，研究者采用了 Vision Transformer（ViT）架構。為了擴展模型容量，他們遵循語言建模中的常見做法，優先考慮擴展寬度而不是深度。下表 1 展示了 AIM 的設計參數，包括它的寬度和深度以及數據量、每個模型容量的優化方案。

AIM 整體模型架構如下圖 2 所示。

在預訓練期間，研究者將因果掩碼用于自注意力層，以對給定先前 patch 的 patch 進行概率建模。更準確地說，給定一個自注意力層，patch i 的嵌入被計算如下：

前綴 Transformer。預訓練中的自回歸目標需要在自注意力操作中使用因果掩碼，不過這與 ViT 模型在下游任務中的標準使用方法不同，后者要部署雙向自注意力。下圖 3 說明了因果和前綴注意力之間的差別。

MLP 預測頭。使用這些頭的目的是防止主干（trunk）特征在預訓練目標中變得過于專門化，從而增強對下游任務的遷移能力。研究者選擇了一種簡單設計，在最終 transformer 層頂部使用 N 個多層感知機（MLP）塊，從而獨立地對每個 patch 進行處理。

直接實現。研究者觀察到，AIM 使用相同的優化超參數集來擴展模型大小，無需進一步調整。

下游適應。研究者專注的場景中，所有模型權重對下游任務都是固定的。在這種情況下，他們只訓練一個分類頭，這可以減輕小型下游數據集過擬合的風險，并顯著降低適應成本。

實驗結果

首先，研究者從參數和訓練數據的角度衡量了擴展本文方法所產生的影響。特別是，他們研究了預訓練目標和下游性能在不同基準之間是否存在相關性；還研究了縮放對損失函數值的影響。所有這些實驗都報告了在 IN-1k 驗證集上的損失函數值。

從圖 4 可以看出，在整個訓練過程中，兩個探針都得到了相應的改善，這表明優化目標會直接帶來更好的下游性能。

他們同時發現，隨著模型容量的增加，損失值和下游任務的準確性都有所提高。這一觀察結果與在 LLMs 中觀察到的趨勢一致，可直接歸因于目標函數的優化，這反過來又會導致學習到更強的表征。

圖 5 中展示了在由 100 萬張圖像組成的小型數據集（即 IN-1k）或由 20 億張圖像組成的大型數據集（即 DFN-2B+）上進行預訓練時驗證損失的變化情況。

在 IN-1k 上進行的訓練很快就會帶來較低的驗證損失，然而這種損失在訓練結束時會惡化，這表明了對訓練數據的過擬合。當在未經整理的 DFN-2B 數據集上進行訓練時，模型開始時的驗證損失較高，但損失持續減少，沒有過擬合的跡象。

當在同一數據集上添加少量 IN-1k 數據，可以觀察到性能進一步提高，最終超過了在 IN-1k 數據集上的預訓練。表 2 證實了這一點。

由于在使用 DFN-2B+ 數據集進行訓練時沒有觀察到過擬合的跡象，因此研究者繼續研究延長預訓練計劃長度的影響。圖 6 展示了將預訓練時間表的長度從 50 萬次迭代增加到 120 萬次迭代的影響。可以觀察到，使用更長計劃進行預訓練的模型的驗證損失明顯降低。這表明可以通過增加模型容量或使用更長的時間表進行預訓練來提高 AIM 的性能。

同時，研究者討論了模型和訓練目標的一些變化所產生的影響。這些消融實驗使用 AIM-0.6B 模型進行，該模型已在 IN-1k 數據集上進行了預訓練和評估。表 3 展示了消融實驗的結果。

研究者還使用自回歸目標訓練的架構與 BERT 在語言領域以及 BEiT 和 MAE 在視覺領域流行的掩蔽目標進行了比較。他們在與 AIM 相同的設置中應用了掩蔽目標，從而將預訓練目標對性能的影響與 AIM 和其他方法之間的其他設計選擇的不同之處隔離開來。表 5 顯示，AIM 在使用自回歸目標時比使用掩蔽目標時表現更好。