與自然語(yǔ)言處理類似，對(duì)預(yù)訓(xùn)練視覺(jué)主干的遷移提高了模型在各種視覺(jué)任務(wù)上的性能。更大的數(shù)據(jù)集、可擴(kuò)展的架構(gòu)和新的訓(xùn)練方法都推動(dòng)了模型性能的提升。

然而，視覺(jué)模型仍然遠(yuǎn)遠(yuǎn)落后于語(yǔ)言模型。具體來(lái)說(shuō)，迄今為止最大的視覺(jué)模型 ViT 只有 4B 參數(shù)，而入門級(jí)語(yǔ)言模型通常超過(guò) 10B 參數(shù)，更別說(shuō)具有 540B 參數(shù)的大型語(yǔ)言模型。

為了探索 AI 模型的性能極限，Google Research 最近在 CV 領(lǐng)域的一項(xiàng)研究，率先將 Vision Transformer 參數(shù)量擴(kuò)展到了 22B，提出 ViT-22B，與之前類似的模型參數(shù)量 4B 相比，可以說(shuō)這是迄今為止最大的稠密型 ViT 模型。

論文地址：https://arxiv.org/pdf/2302.05442.pdf

對(duì)比之前最大的 ViT- G 和 ViT-e，表 1 給出了比較結(jié)果，由下表可得，ViT-22B 主要是擴(kuò)展了模型的寬度，使得參數(shù)量更大，深度和 ViT-G 一樣。

當(dāng)前的 ViT 大模型

正如這位知乎網(wǎng)友所說(shuō)，難道是谷歌在 ChatGPT 上輸了一局，勢(shì)必要在 CV 領(lǐng)域爭(zhēng)口氣？

如何做到的？原來(lái)研究早期，他們發(fā)現(xiàn)在擴(kuò)展 ViT 的過(guò)程中，出現(xiàn)了訓(xùn)練不穩(wěn)定性，并且可能會(huì)帶來(lái)架構(gòu)變化。然后研究人員仔細(xì)設(shè)計(jì)模型，并且以前所未有的效率來(lái)實(shí)現(xiàn)模型并行訓(xùn)練。ViT-22B 的質(zhì)量是通過(guò)一套全面的任務(wù)來(lái)評(píng)估的，從（少樣本）分類到密集輸出任務(wù)，在這些任務(wù)中，它達(dá)到或超過(guò)了當(dāng)前 SOTA 水平。例如，即使用作凍結(jié)的視覺(jué)特征提取器，ViT-22B 在 ImageNet 上的準(zhǔn)確率也達(dá)到了 89.5%。通過(guò)訓(xùn)練 text tower 來(lái)匹配這些視覺(jué)特征，它在 ImageNet 上實(shí)現(xiàn)了 85.9% 的零樣本設(shè)置準(zhǔn)確率。此外，該模型可以看作是一個(gè)教師，用作蒸餾目標(biāo)，研究人員訓(xùn)練了一個(gè) ViT-B 學(xué)生模型，在 ImageNet 上的準(zhǔn)確率為 88.6%，達(dá)到了此類規(guī)模模型上 SOTA 水平。

模型架構(gòu)

ViT-22B 是一種基于 Transformer 的編碼器模型，類似于原始 Vision Transformer 架構(gòu)，但包含以下三個(gè)主要修改，以提高效率和大規(guī)模訓(xùn)練的穩(wěn)定性：并行層、查詢 / 鍵（QK）歸一化和 omitted biases。

并行層。正如 Wang 和 Komatsuzaki 研究所述，該研究設(shè)計(jì)了一個(gè) Attention 和 MLP 并行結(jié)構(gòu)：

這可以通過(guò)組合 MLP 和注意力塊的線性投影來(lái)實(shí)現(xiàn)額外的并行化。值得注意的是，用于查詢 / 鍵 / 值投影的矩陣乘法和 MLP 的第一線性層被融合到一個(gè)單獨(dú)的操作中，對(duì)于 MLP 的注意力外投影和第二層線性層也是如此。

QK 歸一化。訓(xùn)練大模型的一個(gè)困難是模型的穩(wěn)定性，在將 ViT 擴(kuò)展的過(guò)程中，研究人員發(fā)現(xiàn)在幾千輪的 step 后訓(xùn)練損失呈發(fā)散性。特別是在 8B 參數(shù)的模型中這種現(xiàn)象尤為突出。為了穩(wěn)定模型訓(xùn)練，研究人員采用 Gilmer 等人的方法，在點(diǎn)積注意力計(jì)算之前對(duì)查詢和鍵應(yīng)用 LayerNorm 歸一化操作，以提升訓(xùn)練的穩(wěn)定性。具體來(lái)說(shuō)，注意力權(quán)重計(jì)算為：

omitted biases。在 PaLM 之后，偏置項(xiàng)從 QKV 投影中移除，并且所有的 Layernorm 都在沒(méi)有偏置的情況下應(yīng)用，從而提高了加速器的利用率 (3%)，且質(zhì)量沒(méi)有下降。然而，與 PaLM 不同的是，研究人員對(duì) MLP 密集層使用了偏置項(xiàng)，即便如此，這種方式在兼顧質(zhì)量的同時(shí)，速度沒(méi)有下降。 

圖 2 展示了一個(gè) ViT-22B 編碼器塊。嵌入層在原有 ViT 的基礎(chǔ)上進(jìn)行了 patch 提取、線性投影和添加位置嵌入等操作。研究人員使用多頭注意力池化來(lái)聚合頭中的每個(gè) token 表示。

ViT-22B 使用 14 × 14 的 patch，圖像分辨率為 224 × 224。ViT-22B 采用了一種學(xué)習(xí)到的一維位置嵌入。在對(duì)高分辨率圖像進(jìn)行微調(diào)期間，研究人員根據(jù)預(yù)訓(xùn)練的位置嵌入在原始圖像中的位置執(zhí)行二維插值。

訓(xùn)練基礎(chǔ)設(shè)施與效率

ViT-22B 使用 FLAX 庫(kù)，實(shí)現(xiàn)方式是 JAX，并在 Scenic 中構(gòu)建。它同時(shí)利用了模型和數(shù)據(jù)并行性。值得一提的是，研究人員使用了 jax. xmap API，它提供了對(duì)所有中間體的分片（例如權(quán)重和激活）以及芯片間通信的顯式控制。研究人員將芯片組織成大小為 t × k 的 2D 邏輯網(wǎng)格，其中 t 是數(shù)據(jù)平行軸的大小，k 是模型軸的大小。然后，對(duì)于 t 組中的每個(gè)組，k 個(gè)設(shè)備獲得相同批次的圖像，每個(gè)設(shè)備只保留 1/k 的激活，并負(fù)責(zé)計(jì)算所有線性層輸出的 1/k（詳細(xì)內(nèi)容如下）。

圖 3：異步并行線性操作（y = Ax）：跨設(shè)備的重疊通信和計(jì)算的模型并行矩陣乘法。

異步并行線性操作。為了最大限度地提高吞吐量，必須考慮計(jì)算和通信。也就是說(shuō)，如果希望這些操作在分析上等效于未分片的情況，就必須盡可能少地進(jìn)行通信，理想情況下讓它們重疊，這樣就可以保持矩陣乘法單元（FLOP 的大部分容量所在）始終處于繁忙狀態(tài)。

參數(shù)分片。該模型在第一個(gè)軸上是數(shù)據(jù)并行的。每個(gè)參數(shù)可以在這個(gè)軸上完全復(fù)制，也可以讓每個(gè)設(shè)備保存它的一個(gè)塊。研究人員選擇從模型參數(shù)中分割一些大張量，以便能夠擬合更大的模型和批量大小。

使用這些技術(shù)，ViT-22B 在 TPUv4 上訓(xùn)練期間，每個(gè)核每秒處理 1.15k token。ViT-22B 的模型 flops 利用率（MFU）為 54.9%，表明硬件的使用非常有效。請(qǐng)注意，PaLM 報(bào)告的 MFU 為 46.2%，而研究人員在相同硬件上為 ViT-e（僅數(shù)據(jù)并行）測(cè)量的 MFU 為 44.0%。

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

實(shí)驗(yàn)探究了 ViT-22B 用于圖像分類的評(píng)估結(jié)果。

表 2 結(jié)果顯示，ViT-22B 在各種指標(biāo)上仍有顯著的改善。此外，研究表明，像 ViT-22B 這樣的大型模型的 Linear probing 可以接近或超過(guò)具有高分辨率的小型模型的 full fine-tuning 性能，通常成本更小、更容易做到。

研究進(jìn)一步在細(xì)粒度分類數(shù)據(jù)集 iNaturalist 2017 上測(cè)試線性可分離性，將 ViT-22B 與其他 ViT 變體進(jìn)行比較。研究測(cè)試了 224px 和 384px 的輸入分辨率。結(jié)果如圖 4。研究觀察到 ViT-22B 明顯優(yōu)于其他 ViT 變體，特別是在標(biāo)準(zhǔn)的 224px 輸入分辨率下。這表明 ViT-22B 中大量的參數(shù)對(duì)于從圖像中提取詳細(xì)信息是有用的。

表 3 顯示了 ViT-22B 對(duì) CLIP、ALIGN、BASIC、CoCa、LiT 模型的零樣本遷移結(jié)果。表 3 底部比較了三個(gè) ViT 模型性能。

在所有的 ImageNet 測(cè)試集中，ViT-22B 取得了相當(dāng)或更好的結(jié)果。值得注意的是，ObjectNet 測(cè)試集上的零樣本結(jié)果與 ViT 模型大小高度相關(guān)。最大的 ViT-22B 將新的 SOTA 設(shè)置在具有挑戰(zhàn)性的 ObjectNet 測(cè)試集中。

Out-of-distribution （OOD）。研究構(gòu)建了一個(gè)從 JFT 到 ImageNet 的標(biāo)簽映射，以及從 ImageNet 到不同分布外數(shù)據(jù)集的標(biāo)簽映射，即 ObjectNet、ImageNet-v2、ImageNet- R 和 ImageNet- A。

目前可以確認(rèn)的結(jié)果是，與 ImageNet 上的改進(jìn)一致，擴(kuò)展模型增加了分布外性能。這適用于只看過(guò) JFT 圖像的模型，以及在 ImageNet 上進(jìn)行微調(diào)的模型。在這兩種情況下，ViT-22B 在更大的模型上都延續(xù)了 OOD 性能更好的趨勢(shì)（圖 5，表 11）。

此外，研究人員還研究了 ViT-22B 模型在語(yǔ)義分割和單目深度估計(jì)任務(wù)中捕獲的幾何和空間信息質(zhì)量。

語(yǔ)義分割。研究人員在三個(gè)基準(zhǔn)上評(píng)估 ViT-22B 作為語(yǔ)義分割主干：ADE20K、Pascal Context 和 Pascal VOC。從表 4 可以看出，當(dāng)只看到少量分割掩碼時(shí)，ViT-22B 主干遷移效果更好。

單目深度估計(jì)。表 5 總結(jié)了研究的主要發(fā)現(xiàn)。從最上面的行（DPT 解碼器）中可以觀察到，與不同的主干相比，使用 ViT-22B 特性產(chǎn)生了最好的性能（在所有指標(biāo)上）。通過(guò)將 ViT-22B 主干與 ViT-e（一個(gè)較小的模型，但在與 ViT-22B 相同的數(shù)據(jù)上進(jìn)行訓(xùn)練）進(jìn)行比較，研究發(fā)現(xiàn)擴(kuò)展架構(gòu)可以提高性能。

此外，將 ViT-e 主干與 ViT-L（與 ViT-e 類似的架構(gòu)，但訓(xùn)練的數(shù)據(jù)更少）進(jìn)行比較，研究發(fā)現(xiàn)這些改進(jìn)也來(lái)自于擴(kuò)展訓(xùn)練前的數(shù)據(jù)。這些發(fā)現(xiàn)表明，更大的模型和更大的數(shù)據(jù)集都有助于提高性能。

該研究還在視頻數(shù)據(jù)集上進(jìn)行了探索。表 6 展示了在 Kinetics 400 和 Moments in Time 數(shù)據(jù)集上的視頻分類結(jié)果，表明可以使用凍結(jié)的主干實(shí)現(xiàn)具有競(jìng)爭(zhēng)力的結(jié)果。研究首先與 ViT-e 進(jìn)行比較，ViT-e 擁有最大的先驗(yàn)視覺(jué)主干模型，由 40 億個(gè)參數(shù)組成，并且也在 JFT 數(shù)據(jù)集上進(jìn)行訓(xùn)練。我們觀察到更大的 ViT-22B 模型在 Kinetics 400 上提高了 1.5 分，在 Moments in Time 上提高了 1.3 分。

最后研究注意到，通過(guò)完整的端到端微調(diào)，還有進(jìn)一步改進(jìn)的空間。

更多技術(shù)細(xì)節(jié)請(qǐng)參閱原論文。