Transformer無疑是促進自然語言處理領域繁榮的最大功臣，也是GPT-4等大規(guī)模語言模型的基礎架構。

不過相比語言模型動輒成千上萬億的參數(shù)量，計算機視覺領域吃到Transformer的紅利就沒那么多了，目前最大的視覺Transformer模型ViT-e的參數(shù)量還只有40億參數(shù)。

最近谷歌發(fā)布了一篇論文，研究人員提出了一種能夠高效且穩(wěn)定訓練大規(guī)模Vision Transformers(ViT)模型的方法，成功將ViT的參數(shù)量提升到220億。

論文鏈接：https://arxiv.org/abs/2302.05442

為了實現(xiàn)模型的擴展，ViT-22B結合了其他語言模型（如PaLM模型）的思路，使用 QK 歸一化改進了訓練穩(wěn)定性，提出了一種異步并行線性操作（asynchronous parallel linear operations）的新方法提升訓練效率，并且能夠在硬件效率更高的Cloud TPU上進行訓練。

在對ViT-22B模型進行實驗以評估下游任務性能時，ViT-22B也表現(xiàn)出類似大規(guī)模語言模型的能力，即隨著模型規(guī)模的擴大，性能也在不斷提升。

ViT-22B 還可以應用于PaLM-e中，與語言模型結合后的大模型可以顯著提升機器人任務的技術水平。

研究人員還進一步觀察到規(guī)模帶來的其他優(yōu)勢，包括更好地平衡公平性和性能，在形狀/紋理偏見方面與人類視覺感知的一致性，以及更好的穩(wěn)健性。

模型架構

ViT-22B 是一個基于Transformer架構的模型，和原版ViT架構相比，研究人員主要做了三處修改以提升訓練效率和訓練穩(wěn)定性。

并行層（parallel layers）

ViT-22B并行執(zhí)行注意力塊和MLP塊，而在原版Transformer中為順序執(zhí)行。

PaLM模型的訓練也采用了這種方法，可以將大模型的訓練速度提高15%，并且性能沒有下降。

query/key (QK) normalization

在擴展ViT的過程中，研究人員在80億參數(shù)量的模型中觀察到，在訓練幾千步之后訓練損失開始發(fā)散(divergence)，主要是由于注意力logits的數(shù)值過大引起的不穩(wěn)定性，導致零熵的注意力權重（幾乎one-hot）。

為了解決這個問題，研究人員在點乘注意力計算之前對Query和Key使用LayerNorm

在80億參數(shù)模型上的實驗結果如下圖所示，歸一化可以緩解發(fā)散問題。

刪除QKV投影和LayerNorms上的偏置項

和PaLM模型一樣，ViT-22B從QKV投影中刪除了偏置項，并且在所有LayerNorms中都沒有偏置項（bias）和centering，使得硬件利用率提高了3%，并且質(zhì)量沒有下降。

不過與PaLM不同的是，ViT-22B對（內(nèi)部和外部）MLP稠密連接層使用了偏置項，可以觀察到質(zhì)量得到了改善，并且速度也沒有下降。

ViT-22B的編碼器模塊中，嵌入層，包括抽取patches、線性投影和額外的位置嵌入都與原始ViT中使用的相同，并且使用多頭注意力pooling來聚合每個頭中的per-token表征。

ViT-22B的patch尺寸為14×14，圖像的分辨率為224×224（通過inception crop和隨機水平翻轉(zhuǎn)進行預處理）。

異步并聯(lián)線性運算（asynchronous parallel linear operations）

大規(guī)模的模型還需要分片（sharding），即將模型參數(shù)分布在不同的計算設備中，除此之外，研究人員還把激活（acctivations，輸入的中間表征）也進行分片。

因為輸入和矩陣本身都是分布在各種設備上的，即使是像矩陣乘法這樣簡單的操作也需要特別小心。

研究人員開發(fā)了一種稱為異步并行線性運算的方法，可以在矩陣乘法單元（在TPU 中占據(jù)絕大多數(shù)計算能力的單元）中計算時，同時對設備之間的激活和權值進行通信。

異步方法最小化了等待傳入通信的時間，從而提高了設備效率。

異步并行線性運算的目標是計算矩陣乘法 y = Ax，但矩陣 A 和激活 x 都分布在不同的設備上，需要通過跨設備的重疊通信和計算來實現(xiàn)這一點。矩陣 A 在設備之間進行列分片（column-shard），每個矩陣包含一個連續(xù)的切片，每個塊表示為 Aij，更多細節(jié)請看原始論文。

實驗結果

為了說明ViT-22B學習到的表征非常豐富，研究人員使用LiT-tuning訓練一個文本模型來生成一些表征用來對齊文本和圖像。

下面是用Parti 和 Imagen 生成的分布外（out-of-distribution）圖像得到的實驗結果，可以看到ViT-22B的zero-shot圖像分類泛化能力非常強，僅從web上爬取的自然圖像就能識別出沒見過的物體和場景。

論文中還討論了ViT-22B在視頻分類、深度估計和語義分割任務上的效果。

與人類目標識別對齊

為了驗證 ViT-22B 分類決策與人類分類決策的一致性，研究人員對 ViT-22B 進行了微調(diào)，對分布外(OOD)數(shù)據(jù)集的不同分辨率進行了微調(diào)，其中人類比較數(shù)據(jù)可通過model-vs-human toolbox獲得。

該工具箱主要衡量三個關鍵指標: 模型如何處理失真(準確性) ？人和模型的精度(精度差)有什么不同？人和模型的錯誤模式(錯誤一致性)有多相似？

形狀偏差評估(值越大代表更多的形狀偏差)。許多視覺模型具有低形狀/高紋理偏差，而在 ImageNet 上進行微調(diào)的 ViT-22B具有迄今為止在 ML 模型中記錄的最高形狀偏差，更接近于人類形狀偏見

實驗結果顯示，雖然并非所有的微調(diào)解決方案都表現(xiàn)得很好，但 ViT-22B 變體在所有三個指標上都達到了新高。

此外，ViT-22B 模型在視覺模型中也有最高的形狀偏差記錄。這意味著他們主要使用目標的形狀，而不是目標的紋理來進行分類決策，策略結果類似于人類的感知(其形狀偏差為96%)。

標準模型(例如，ResNet-50有20-30% 的形狀偏差)通常根據(jù)紋理來分類，而高形狀偏差的模型則傾向于關注形狀(下圖識別為貓)，盡管人類和模型的感知之間仍然存在許多差異，但是 ViT-22B 顯示出與人類視覺對象識別更多的相似性。

貓還是大象？車還是鐘？鳥還是自行車？具有某個物體的形狀和另一個不同物體紋理的圖像，可用于測量形狀/紋理偏差

分布外（out-of-distribution）性能

測量 OOD 數(shù)據(jù)集的性能有助于評估模型泛化性。

在這個實驗中，研究人員構建了從 JFT 到 ImageNet 的標簽映射，以及從 ImageNet 到不同的分布外數(shù)據(jù)集(如 ObjectNet)的標簽映射。

對這些數(shù)據(jù)進行預訓練后的結果如下圖所示，然后在 ImageNet 上對模型進行完全微調(diào)。

可以觀察到縮放 Vision Transformers 可以提高 OOD 性能: 即使 ImageNet 的精度達到飽和，也可以看到 ObjectNet 上從 ViT-e 換成 ViT-22B 模型可以顯著提升性能。

線性探測Linear Probe

線性探測是一種將單個線性層置于凍結模型之上的技術，與完全微調(diào)相比，這種方法的訓練成本更低，設置起來也更容易。

在 ImageNet 上訓練的線性探測結果，在 ImageNet-Real，ImageNet-v2，ObjectNet，ImageNet-R 和 ImageNet-A 數(shù)據(jù)集上評估，提供高分辨率微調(diào) ViT-e/14作為參考

從結果中可以觀察到，ViT-22B 的線性探測性能接近于使用高分辨率圖像對較小模型進行全面微調(diào)的最先進水平，其中具有較高分辨率的訓練通常要昂貴得多，但可以在許多任務上取得更好的結果。

蒸餾

利用蒸餾法，可以將較大模型的知識轉(zhuǎn)化為較小模型的知識，可以提升成本更高、運行速度更慢的大模型的運行效率。

從實驗結果中可以發(fā)現(xiàn)，ViT-22B 的知識可以遷移到更小的模型，如 ViT-B/16和 ViT-L/16，并在同等模型尺寸下在ImageNet上刷新了性能記錄。

公平性與偏見

機器學習模型容易受到意想不到的不公平偏見的影響，例如找到錯誤的相關性或者在各個子群體之間存在性能差距，研究人員發(fā)現(xiàn)，擴大模型規(guī)模有助于緩解這些問題。

首先，規(guī)模是一個有前景的權衡方式，即使模型經(jīng)過訓練后再進行后處理，將其人口平等（demographic parity）水平控制在規(guī)定的、可容忍的水平之下，性能也會隨著規(guī)模的增加而提高。

上圖: 去偏前 CelebA 中每個子組的精度。下圖: y 軸顯示了在這個例子中突出顯示的兩個特定亞組(女性和男性)的表現(xiàn)的絕對差異。與較小的 ViT 模型相比，ViT-22B 在性能的差距很小。

更重要的是，這不僅適用于以準確性衡量性能的情況，而且適用于其他度量，例如校準，即對模型估計概率的真實性的統(tǒng)計測量，所有子群的分類隨著規(guī)模的增大而趨于改善，并且ViT-22B 降低了各子群之間的性能差距。

結論

研究人員提出了一個目前最大的視覺Transformer模型 ViT-22B，包含220億參數(shù)。

通過對原始模型架構進行微小但關鍵的修改后，實現(xiàn)了更高的硬件利用率和訓練穩(wěn)定性，從而得到了一個在幾個基準測試上提高了模型的上限性能。

使用凍結模型生成嵌入，只需要在頂部訓練幾層，即可獲得很好的性能，并且評估結果進一步表明，與現(xiàn)有模型相比，ViT-22B 在形狀和紋理偏差方面顯示出與人類視知覺更多的相似性，并且在公平性和穩(wěn)健性方面提供了優(yōu)勢。