過(guò)去幾年，大規(guī)模預(yù)訓(xùn)練模型在 NLP 領(lǐng)域取得了成功，這類(lèi)模型可以通過(guò)幾個(gè)示例解決復(fù)雜的推理任務(wù)，也可以按照指令進(jìn)行推理。

眾所周知的是，預(yù)訓(xùn)練模型能取得成功的一個(gè)理論基礎(chǔ)是，隨著容量（即參數(shù)量）或預(yù)訓(xùn)練數(shù)據(jù)量的增加，模型能夠不斷提升性能。

這很自然地引發(fā)了研究者們的聯(lián)想：使用自回歸目標(biāo)對(duì) Transformers 進(jìn)行縮放的成功是否僅限于文本？

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

• 論文鏈接：https://arxiv.org/pdf/2401.08541.pdf
• 項(xiàng)目地址：https://github.com/apple/ml-aim

先說(shuō)結(jié)論：研究者發(fā)現(xiàn)，模型容量可以輕松擴(kuò)展到數(shù)十億個(gè)參數(shù)，并且 AIM 能夠有效利用大量未經(jīng)整理的圖像數(shù)據(jù)。

他們利用包括 ViT、大規(guī)模網(wǎng)絡(luò)數(shù)據(jù)集和 LLM 預(yù)訓(xùn)練最新進(jìn)展在內(nèi)的工具集，重新審視了 iGPT 等自回歸表征學(xué)習(xí)方面的前期工作，此外還引入了兩處架構(gòu)修改，以適應(yīng)視覺(jué)特征的自回歸預(yù)訓(xùn)練。

首先，研究者并沒(méi)有像 LLM 通常那樣將自注意力限制為完全因果關(guān)系，而是采用了 T5 中的前綴注意力。這一選擇使得能夠在下游任務(wù)中轉(zhuǎn)向完全雙向的注意力。其次，研究者使用了參數(shù)化程度較高的 token-level 預(yù)測(cè)頭，其靈感來(lái)自對(duì)比學(xué)習(xí)中使用的預(yù)測(cè)頭。他們觀察到，這種修改大大提高了后續(xù)特征的質(zhì)量，而在訓(xùn)練過(guò)程中的開(kāi)銷(xiāo)卻很小?？傮w來(lái)說(shuō)，AIM 的訓(xùn)練與最近的 LLM 訓(xùn)練類(lèi)似，而且不依賴(lài)于監(jiān)督式或自監(jiān)督式方法所需的任何 stability-inducing 技術(shù)。

隨后，研究者對(duì)一系列模型展開(kāi)了研究，這些模型的參數(shù)從 600M 到 7B 不等，都是使用 20 億帶許可的未編輯圖像進(jìn)行預(yù)訓(xùn)練的。如圖 1 所示，以 15 個(gè)圖像識(shí)別基準(zhǔn)的平均準(zhǔn)確率來(lái)衡量，AIM 模型在與模型規(guī)模的關(guān)系上表現(xiàn)出很強(qiáng)的擴(kuò)展性，容量越大的模型下游性能越好。更重要的是，驗(yàn)證集上的目標(biāo)函數(shù)值與后續(xù)凍結(jié)特征的質(zhì)量之間存在相關(guān)性。這一觀察結(jié)果證明，自回歸目標(biāo)足以滿(mǎn)足視覺(jué)特征的訓(xùn)練要求。此外，隨著對(duì)更多圖像進(jìn)行訓(xùn)練，研究者還觀察到了下游性能的持續(xù)改善，且沒(méi)有飽和的跡象?？傮w而言，這些觀察結(jié)果與之前關(guān)于擴(kuò)展大型語(yǔ)言模型的研究結(jié)果是一致的。

但同時(shí)引起注意的是，本文實(shí)驗(yàn)所使用的模型規(guī)模有限，是否能在更大參數(shù)量級(jí)的模型上驗(yàn)證此規(guī)律，有待進(jìn)一步探索。

方法概覽

本文的訓(xùn)練目標(biāo)遵循應(yīng)用于圖像 patch 序列的標(biāo)準(zhǔn)自回歸模型。更準(zhǔn)確地說(shuō)，圖像 x 被分割為 K 個(gè)不重疊 patch x_k 組成的網(wǎng)格 k ∈ [1, K]，這些 patch 共同形成 token 序列。

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

預(yù)測(cè)損失。研究者的訓(xùn)練目標(biāo)自然會(huì)產(chǎn)生某些損失變體，每個(gè)變體對(duì)應(yīng)分布 P (x_k | x_<k) 的選擇。他們還考慮通過(guò)使用離線 tokenizer，將帶有 patch 的交叉熵?fù)p失轉(zhuǎn)換為離散 token。消融實(shí)驗(yàn)表明，這些設(shè)計(jì)是有效的，盡管不會(huì)產(chǎn)生像像素級(jí)損失那樣顯著的特征。

架構(gòu)

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

AIM 整體模型架構(gòu)如下圖 2 所示。

在預(yù)訓(xùn)練期間，研究者將因果掩碼用于自注意力層，以對(duì)給定先前 patch 的 patch 進(jìn)行概率建模。更準(zhǔn)確地說(shuō)，給定一個(gè)自注意力層，patch i 的嵌入被計(jì)算如下：

前綴 Transformer。預(yù)訓(xùn)練中的自回歸目標(biāo)需要在自注意力操作中使用因果掩碼，不過(guò)這與 ViT 模型在下游任務(wù)中的標(biāo)準(zhǔn)使用方法不同，后者要部署雙向自注意力。下圖 3 說(shuō)明了因果和前綴注意力之間的差別。

MLP 預(yù)測(cè)頭。使用這些頭的目的是防止主干（trunk）特征在預(yù)訓(xùn)練目標(biāo)中變得過(guò)于專(zhuān)門(mén)化，從而增強(qiáng)對(duì)下游任務(wù)的遷移能力。研究者選擇了一種簡(jiǎn)單設(shè)計(jì)，在最終 transformer 層頂部使用 N 個(gè)多層感知機(jī)（MLP）塊，從而獨(dú)立地對(duì)每個(gè) patch 進(jìn)行處理。

直接實(shí)現(xiàn)。研究者觀察到，AIM 使用相同的優(yōu)化超參數(shù)集來(lái)擴(kuò)展模型大小，無(wú)需進(jìn)一步調(diào)整。

下游適應(yīng)。研究者專(zhuān)注的場(chǎng)景中，所有模型權(quán)重對(duì)下游任務(wù)都是固定的。在這種情況下，他們只訓(xùn)練一個(gè)分類(lèi)頭，這可以減輕小型下游數(shù)據(jù)集過(guò)擬合的風(fēng)險(xiǎn)，并顯著降低適應(yīng)成本。

實(shí)驗(yàn)結(jié)果

首先，研究者從參數(shù)和訓(xùn)練數(shù)據(jù)的角度衡量了擴(kuò)展本文方法所產(chǎn)生的影響。特別是，他們研究了預(yù)訓(xùn)練目標(biāo)和下游性能在不同基準(zhǔn)之間是否存在相關(guān)性；還研究了縮放對(duì)損失函數(shù)值的影響。所有這些實(shí)驗(yàn)都報(bào)告了在 IN-1k 驗(yàn)證集上的損失函數(shù)值。

從圖 4 可以看出，在整個(gè)訓(xùn)練過(guò)程中，兩個(gè)探針都得到了相應(yīng)的改善，這表明優(yōu)化目標(biāo)會(huì)直接帶來(lái)更好的下游性能。

他們同時(shí)發(fā)現(xiàn)，隨著模型容量的增加，損失值和下游任務(wù)的準(zhǔn)確性都有所提高。這一觀察結(jié)果與在 LLMs 中觀察到的趨勢(shì)一致，可直接歸因于目標(biāo)函數(shù)的優(yōu)化，這反過(guò)來(lái)又會(huì)導(dǎo)致學(xué)習(xí)到更強(qiáng)的表征。

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

在 IN-1k 上進(jìn)行的訓(xùn)練很快就會(huì)帶來(lái)較低的驗(yàn)證損失，然而這種損失在訓(xùn)練結(jié)束時(shí)會(huì)惡化，這表明了對(duì)訓(xùn)練數(shù)據(jù)的過(guò)擬合。當(dāng)在未經(jīng)整理的 DFN-2B 數(shù)據(jù)集上進(jìn)行訓(xùn)練時(shí)，模型開(kāi)始時(shí)的驗(yàn)證損失較高，但損失持續(xù)減少，沒(méi)有過(guò)擬合的跡象。

當(dāng)在同一數(shù)據(jù)集上添加少量 IN-1k 數(shù)據(jù)，可以觀察到性能進(jìn)一步提高，最終超過(guò)了在 IN-1k 數(shù)據(jù)集上的預(yù)訓(xùn)練。表 2 證實(shí)了這一點(diǎn)。

由于在使用 DFN-2B+ 數(shù)據(jù)集進(jìn)行訓(xùn)練時(shí)沒(méi)有觀察到過(guò)擬合的跡象，因此研究者繼續(xù)研究延長(zhǎng)預(yù)訓(xùn)練計(jì)劃長(zhǎng)度的影響。圖 6 展示了將預(yù)訓(xùn)練時(shí)間表的長(zhǎng)度從 50 萬(wàn)次迭代增加到 120 萬(wàn)次迭代的影響?？梢杂^察到，使用更長(zhǎng)計(jì)劃進(jìn)行預(yù)訓(xùn)練的模型的驗(yàn)證損失明顯降低。這表明可以通過(guò)增加模型容量或使用更長(zhǎng)的時(shí)間表進(jìn)行預(yù)訓(xùn)練來(lái)提高 AIM 的性能。

同時(shí)，研究者討論了模型和訓(xùn)練目標(biāo)的一些變化所產(chǎn)生的影響。這些消融實(shí)驗(yàn)使用 AIM-0.6B 模型進(jìn)行，該模型已在 IN-1k 數(shù)據(jù)集上進(jìn)行了預(yù)訓(xùn)練和評(píng)估。表 3 展示了消融實(shí)驗(yàn)的結(jié)果。

研究者還使用自回歸目標(biāo)訓(xùn)練的架構(gòu)與 BERT 在語(yǔ)言領(lǐng)域以及 BEiT 和 MAE 在視覺(jué)領(lǐng)域流行的掩蔽目標(biāo)進(jìn)行了比較。他們?cè)谂c AIM 相同的設(shè)置中應(yīng)用了掩蔽目標(biāo)，從而將預(yù)訓(xùn)練目標(biāo)對(duì)性能的影響與 AIM 和其他方法之間的其他設(shè)計(jì)選擇的不同之處隔離開(kāi)來(lái)。表 5 顯示，AIM 在使用自回歸目標(biāo)時(shí)比使用掩蔽目標(biāo)時(shí)表現(xiàn)更好。

表 6 展示了 AIM 與其他 SOTA 方法在 15 種不同基準(zhǔn)中的 Attentive Probing 性能。

除此之外，研究者還探索了 LoRA 這種高效的微調(diào)方法，表 8 展示了對(duì) AIM 進(jìn)行 LoRA 微調(diào)的結(jié)果。LoRA 與 AIM 兼容，與凍結(jié)主干評(píng)估相比，性能有了很大提升。例如，AIM-7B 提高了 3.9%（與上一層性能相比），而微調(diào)的主干參數(shù)只提升 0.1%。

更多技術(shù)細(xì)節(jié)和實(shí)驗(yàn)結(jié)果請(qǐng)參閱原論文。