半年多來，Meta 開源的 LLaMA 架構在 LLM 中經受了考驗並大獲成功（訓練穩定、容易做 scaling）。
沿襲 ViT 的研究思路，我們能否藉助創新性的 LLaMA 架構，真正實現語言和影像的架構統一？
在這一命題上，最近的一項研究 VisionLLaMA 取得了進展。VisionLLaMA 在影像生成（包含 Sora 依賴的底層的 DIT）和理解（分類、分割、檢測、自監督）等多個主流任務上相較於原 ViT 類方法提升顯著。
論文標題：VisionLLaMA: A Unified LLaMA Interface for Vision Tasks
論文地址：https://arxiv.org/abs/2403.00522
程式碼地址：https://github.com/Meituan-AutoML/VisionLLaMA
該研究在統一影像和語言架構方面的嘗試，可以複用 LLM 社群在 LLaMA 上的訓練（穩定且有效的 scaling）、部署等一系列成果。
研究背景
大語言模型是當前學術界研究的熱點，其中，LLaMA 是最具影響力和代表性的工作之一，許多最新的研究工作都基於該架構開展，各種應用的解決方案大都建立在該系列的開源模型之上。在多模態模型的進展中，其中許多方法都依賴 LLaMA 進行文字處理、並依賴類似 CLIP 的視覺 transformer 進行視覺感知。同時，許多工作致力於加快 LLaMA 的推理速度、降低 LLaMA 的儲存成本。總而言之，LLaMA 現在是事實上最通用、最重要的大語言模型架構。
LLaMA 架構的成功使得本文作者提出了一個簡單而有趣的設想：該架構是否可以在視覺模態上同樣成功？如果答案是肯定的，那麼視覺模型和語言模型都可以使用相同的統一架構，並從為 LLaMA 設計的各種動態部署技術中受益。然而，這是一個複雜的問題，因為這兩種模態之間存在一些明顯的差異。
首先存在維度差異：文字序列是一維的，而視覺任務需要處理兩個或更多維度的資料；其次存在結構差異：許多視覺任務依賴於金字塔結構的骨幹網路以獲得更好的效能，而 LLaMA 是一個結構上樸素的編碼器；第三，需要有效處理不同解析度的影像和影片輸入。
本文旨在解決這些挑戰，並彌合不同模態之間的架構差距，具體為提出適應視覺任務的 LLaMA 架構，解決與模態差異相關的難題，並實現透過一種統一的方法對視覺和語言資料進行處理。
本文主要貢獻如下：
1. 本文提出 VisionLLaMA，一種類似於 LLaMA 的視覺 transformer 架構，以減少語言和視覺之間的架構差異。
2. 本文研究了使 VisionLLaMA 適應常見視覺任務的方法，包括影像理解和建立（圖 1）。本文研究了兩種廣為人知的視覺架構方案（常規結構和金字塔結構），並評估它們在監督和自監督學習場景下的效能。此外，本文還提出了 AS2DRoPE（即自動縮放 2D RoPE），它將旋轉位置編碼從 1D 擴充套件到 2D，並利用插值縮放來適應任意解析度。
3. 在精確的評估下，VisionLLaMA 在影像生成、分類、語義分割和目標檢測等許多代表性任務中明顯優於當前主流且被精確微調過的視覺 transformer。大量實驗表明，VisionLLaMA 與現有視覺 transformer 相比具有更快的收斂速度和更好的效能。
VisionLLaMA 總體架構設計
常規 Transformer
本文提出的常規 VisionLLaMA 遵循 ViT 的流程，並且儘可能保留 LLaMA 的架構設計。對於一張影像，首先將其變換並展平為一個序列，然後在序列的開頭新增一個類別 token，整個序列透過 L 個 VisionLLaMA block 進行處理。與 ViT 不同，VisionLLaMA 不向輸入序列新增位置編碼，因為 VisionLLaMA 的 block 包含位置編碼。具體來說，該 block 與標準的 ViT block 有兩點不同：具有位置編碼（RoPE）的自注意力和 SwiGLU 啟用。本文仍然使用 LayerNorm 而不是 RMSNorm，因為本文透過實驗發現前者表現更好（參見表 11g）。block 的結構如圖 2 (a) 所示。本文發現在視覺任務中直接應用 1D RoPE 不能很好地推廣到不同的解析度上，因此將其擴充套件到二維形式：
金字塔結構 Transformer
將 VisionLLaMA 應用於類似 Swin 的基於視窗的 transformer 非常簡單，因此本文選擇在更強的基線 Twins 上探索如何構建強大的金字塔結構 transformer。Twins 的原始架構利用了條件位置編碼、以區域性 - 全域性注意力的形式進行交錯的區域性 - 全域性資訊交換。這些元件在各種 transformer 中十分常見，這意味著在各類 transformer 變體中應用 VisionLLaMA 並不困難。
本文的目標不是發明一種全新金字塔結構的視覺 transformer ，而是如何在現有設計的基礎上調整 VisionLLaMA 的基本設計，因此本文遵循對架構和超引數進行最少修改的原則。遵循 ViT 的命名方式，兩個連續的 block 可以寫為：
其中 LSA 是組內的區域性自注意力操作，GSA 是透過與每個子視窗中的代表性鍵值互動而進行的全域性子取樣的注意力。本文移除了金字塔結構 VisionLLaMA 中的條件位置編碼，因為 AS2DRoPE 中已經包含了位置資訊。此外，還移除了類別 token，並在分類頭之前使用 GAP（全域性平均池化），該設定下的 block 結構如圖 2 (b) 所示。
超越序列長度限制的訓練或推理
將一維 RoPE 擴充到二維：對不同的輸入解析度進行處理是視覺任務中的常見需求。卷積神經網路使用滑動視窗機制來處理可變長度。與之相比，大多數視覺 transformer 應用區域性視窗操作或插值，例如 DeiT 在不同解析度上訓練時採用雙三次插值；CPVT 使用基於卷積的位置編碼。本文中評估了 1D RoPE 的效能，發現其在 224×224 解析度上擁有最高的精度，然而當解析度上升到 448×448 時，精度急劇下降甚至為 0。因此，本文將一維 RoPE 擴充套件到二維。對於多頭自注意力機制，二維 RoPE 在不同頭之間共享。
位置插值有助於二維 RoPE 更好地泛化：受一些使用插值來擴充套件 LLaMA 的上下文視窗的工作啟發，在更高解析度的參與下，VisionLLaMA 採用類似方式擴充套件二維上下文視窗。與具有擴大的固定上下文長度的語言任務不同，目標檢測等視覺任務通常在不同的迭代中處理不同的取樣解析度。本文使用 224×224 的輸入解析度對小模型進行訓練，並在不重新訓練的情況下評估更大解析度的效能，指引本文能夠更好的應用內插值或外差值策略。經過實驗，本文選擇應用基於 “錨點解析度” 的自動縮放插值（AS2DRoPE）。對 H × H 的方形影像和 B × B 的錨點解析度進行處理的計算方式如下：
這種計算方式效率高並且不會引入額外的成本。如果訓練解析度保持不變，AS2DRoPE 會退化為 2 維 RoPE。
由於需要將位置資訊新增到彙總的鍵值中，本文對於金字塔結構設定下的 GSA 進行了特殊處理。這些子取樣的鍵值是透過特徵圖上的抽象生成的。本文使用核心大小為 k×k 且步長為 k 的卷積。如圖 3 所示，生成的鍵值的座標可以表示為取樣特徵的平均值。
實驗結果
本文全面評估了 VisionLLaMA 在影像生成、分類、分割和檢測等任務上的有效性。預設情況下，本文所有模型均在 8 個 NVIDIA Tesla A100 GPU 上進行訓練。
影像生成
基於 DiT 框架的影像生成：本文選擇在 DiT 框架下應用 VisionLLaMA，因為 DiT 是使用視覺 Transformer 和 DDPM 進行影像生成的代表性工作。本文用 VisionLLaMA 替換了 DiT 原來的視覺 transformer，同時保持其他元件與超引數不變。該實驗證明了 VisionLLaMA 在影像生成任務上的通用性。與 DiT 相同，本文設定 DDPM 的 sample steps 為 250，實驗結果如表 1 所示。與大多數方法保持一致，FID 被視為主要指標，並在其他次要指標上例如 sFID、Precision/Recall、Inception Score 進行評估。結果表明，VisionLLaMA 在各種模型尺寸上都顯著優於 DiT。本文還將 XL 模型的訓練步數擴充套件到 2352k，以評估本文的模型是否具有更快的收斂優勢，或者在更長的訓練週期設定下仍然表現更好。DiT-LLaMA-XL/2 的 FID 比 DiT-XL/2 低 0.83，表明 VisionLLaMA 不僅具有更好的計算效率，而且比 DiT 具有更高的效能。圖 1 中展示了使用 XL 模型生成的一些示例。
基於 SiT 框架的影像生成：SiT 框架顯著提高了使用視覺 transformer 生成影像的效能。本文用 VisionLLaMA 替換 SiT 中的視覺 transformer，以評估更好的模型架構帶來的收益，本文將其稱為 SiT-LLaMA。實驗保留了 SiT 中其餘所有設定與超引數，所有模型都使用相同數量的步驟進行訓練，在所有實驗中都使用線性插值（linear interpolant）和快速模型（velocity model）。為了進行公平比較，本文還重新執行已釋出的程式碼，並使用 250 steps 的 SDE 取樣器（Euler）對 50k 256×256 影像進行取樣，結果如表 2 中所示。SiT-LLaMA 在各種容量級別的模型中均優於 SiT。與 SiT-L/2 相比，SiT-LLaMA-L/2 降低了 5.0 FID，其幅度大於新框架帶來的提升（4.0 FID）。本文還在表 13 中展示了更高效的 ODE 取樣器 (dopri5)，與本文方法的效能差距仍然存在。可以得出與與 SiT 論文中的類似的結論：SDE 比其對應的 ODE 具有更好的效能。
ImageNet 上的影像分類
全監督訓練
本節重點關注模型在 ImageNet-1K 資料集上的全監督訓練，排除其他資料集或蒸餾技巧的影響，所有模型均使用 ImageNet-1K 訓練集進行訓練，並在表 3 中展示了在驗證集上的準確性結果。
常規視覺 Transformer 的比較：DeiT3 是當前最先進的常規視覺 transformer，它提出了一種特殊的資料增強並執行廣泛的超引數搜尋以提高效能。DeiT3 對超引數敏感並且容易出現過擬合，用 GAP（全域性平均池化）替換類別 token 會導致 DeiT3-Large 模型在經過 800 個 epoch 訓練後準確率下降 0.7%。因此，本文在常規 transformer 中使用類別 token 而不是 GAP。結果如表 3 中所示，其中 VisionLLaMA 取得了與 DeiT3 相當的 top-1 精度。單一解析度上的準確性並不能提供全面的比較，本文還評估了不同影像解析度的效能，結果如表 4 所示。對於 DeiT3，本文使用雙三次插值來進行可學習的位置編碼。儘管這兩個模型在 224×224 解析度下具有相當的效能，但當解析度增加時，差距會擴大，這意味著本文的方法在不同解析度下具有更好的泛化能力，這對於目標檢測等許多下游任務來說至關重要。
金字塔結構的視覺 transformer 比較：本文使用與 Twins-SVT 相同的架構，詳細配置列於表 17。本文移除了條件位置編碼，因為 VisionLLaMA 已經包含一種旋轉位置編碼。因此，VisionLLaMA 是一種無卷積架構。本文沿用 Twins-SVT 中的包含超引數在內的所有設定，與 Twins-SVT 保持一致，本文不使用類別 token，而是應用 GAP。結果如表 3 所示，本文的方法在各個模型級別上都實現了與 Twins 相當的效能，並且始終優於 Swin。
自監督訓練
本文使用 ImageNet 資料集評估自監督視覺 transformer 的兩種常見方法，同時將訓練資料限制為 ImageNet-1K，移除了任何使用 CLIP、DALLE 或蒸餾等可以提高效能的元件，本文的實現基於 MMPretrain 框架，利用 MAE 框架並使用 VisionLLaMA 替換編碼器，同時保持其他元件不變。該對照實驗能夠評估本文方法的有效性。此外，本文使用與所比較方法相同的超引數設定，在這種設定下，與強大的基線相比依然實現了顯著的效能提升。
Full fine-tuning 設定：在當前設定下，模型首先使用預訓練的權重進行初始化，然後使用完全可訓練的引數進行額外的訓練。VisionLLaMA-Base 在 ImageNet 上經過 800 個 epoch 的訓練，達到了 84.0% 的 top-1 準確率，比 ViT-Base 提高了 0.8%。本文的方法訓練速度比 SimMIM 快約 3 倍。本文還將訓練週期增加到 1600，以驗證 VisionLLaMA 能否在足夠的訓練資源下保持優勢。VisionLLaMA-Base 在 MAE 變體中取得了新的 SOTA 結果，top-1 準確率達到 84.3%，比 ViT-Base 提高了 0.9%。考慮到 full fine-tuning 具有效能飽和風險，本文方法的提升十分顯著。
Linear probing：最近的一項工作認為線性探測度量（linear probing metric）是對錶示性學習更加可靠的評估。在當前設定下，模型由 SSL 階段的預訓練權重初始化。然後，在訓練過程中，除了分類器頭之外，整個骨幹網路都被凍結。結果如表 5 所示：在訓練成本為 800 個 epoch 的情況下，VisionLLaMA-Base 的效能優於 ViTBase-MAE 4.6%。它還超過了訓練了 1600 個 epoch 的 ViT-Base-MAE。當 VisionLLaMA 訓練 1600 個 epoch 時，VisionLLaMA-Base 達到了 71.7% 的 top1 準確率。本文方法還擴充套件到 VisionLLaMA-Large，相比 ViT-Large 提高了 3.6%。
ADE20K 資料集上的語義分割
全監督訓練
按照 Swin 的設定，本文在 ADE20K 資料集上使用語義分割來評估本文方法的有效性。為了進行公平比較，本文限制基線模型僅使用 ImageNet-1K 進行預訓練。本文使用 UperNet 框架，並用金字塔結構 VisionLLaMA 替換主幹網路。本文的實現基於 MMSegmentation 框架。模型訓練步數設定為 160k，全域性 batch size 為 16。結果如表 6 中所示，在相近的 FLOP 下，本文的方法比 Swin 和 Twins 的效能高出 1.2% mIoU 以上。
自監督訓練
本文使用 UperNet 框架在 ADE20K 資料集上進行語義分割，用 VisionLLaMA 替換 ViT 主幹，同時保持其他元件和超引數不變。本文的實現基於 MMSegmentation，結果如表 7 所示。對於 800 個 epoch 的預訓練組，VisionLLaMA-B 將 ViT-Base 顯著提升了 2.8% mIoU。本文方法還明顯優於其他一些改進，例如引入額外的訓練目標或特徵，這些方法會給訓練過程帶來額外的開銷並降低訓練速度。相比之下，VisionLLaMA 僅涉及基礎模型的替換，並且具有快速的訓練速度。本文進一步評估了 1600 個較長預訓練 epoch 的效能，VisionLLaMA-B 在 ADE20K 驗證集上實現了 50.2% mIoU，這使得 ViT-B 的效能提高了 2.1% mIoU。
COCO 資料集上的目標檢測
全監督訓練
本文評估了金字塔結構 VisionLLaMA 在 COCO 資料集上的目標檢測任務的效能。本文使用 Mask RCNN 框架並用金字塔結構 VisionLLaMA 替換主幹網路，類似於 Swin 的設定，該金字塔結構 VisionLLaMA 在 ImageNet-1K 資料集上預訓練了 300 個 epoch。因此，本文的模型具有與 Twins 相同數量的引數和 FLOP。該實驗能夠用於驗證本文方法在目標檢測任務上的有效性。本文的實現基於 MMDetection 框架，表 8 中展示了標準的 36 個 epoch 訓練週期 (3×) 的結果，本文的模型優於 Swin 和 Twins。具體來說，VisionLLaMA-B 比 Swin-S 高出 1.5% 的 box mAP 和 1.0% mask mAP。與更強的基線 Twins-B 相比，本文的方法具有在 box mAP 上高出 1.1% ，在 mask mAP 上高出 0.8% 的優勢。
自監督訓練
本文應用基於 ViTDet 框架的 VisionLLaMA，該框架利用常規視覺 transformer 來實現與對應金字塔結構視覺 transformer 相當的效能。本文使用 Mask RCNN 檢測器，並用 VisionLLaMA-Base 模型替換 vit-Base 主幹網路，該模型使用 MAE 預訓練 800 輪。原始的 ViTDet 收斂緩慢，需要專門的訓練策略，例如更長的訓練週期才能實現最佳效能。在訓練過程中，本文發現 VisionLLaMA 在 30 個 epoch 後達到了相似的效能，因此，本文直接應用標準的 3x 訓練策略。本文方法的訓練成本僅為基線的 36%。與所比較方法不同，本文方法不進行最佳超引數搜尋。結果如表 9 所示，VisionLLaMA 在 Box mAP 上優於 ViT-B 0.6%，在 mask mAP 上優於 ViT-B 0.8%。
消融實驗與討論
消融實驗
本文預設選擇在 ViT-Large 模型上進行消融實驗，因為本文觀察到該模型在多次執行中產生的方差較小。
FFN 和 SwiGLU 的消融：本文用 SwiGLU 替換 FFN ，結果如表 11a 中所示。由於明顯效能差距，本文選擇使用 SwiGLU 以避免對 LLaMA 架構引入額外的修改。
歸一化策略的消融：本文對 transformer 中兩種廣泛使用的歸一化方法 RMSNorm 和 LayerNorm 進行了比較，結果如表 11g 中所示。後者具有更好的最終效能，這表明重新居中不變性（re-centering invariance）在視覺任務中也很重要。本文還計算了每次迭代花費的平均時間用來衡量訓練速度，其中 LayerNorm 僅比 RMSNorm 慢 2%。因此，本文選擇 LayerNorm 而不是 RMSNorm 以獲得更均衡的效能。
部分位置編碼：本文使用 RoPE 調整全部 channel 的比率，結果如表 11b 中所示，結果表明將比率設定在小閾值上即可獲得良好的效能，不同的設定之間沒有觀察到存在顯著的效能差異。因此，本文保留 LLaMA 中的預設設定。
基礎頻率：本文對基礎頻率進行更改與比較，結果如表 11c 中所示，結果表明，效能對於大範圍的頻率來說是穩健的。因此，本文保留 LLaMA 中的預設值以避免部署時的額外特殊處理。
每個注意力頭之間共享位置編碼：本文發現，在不同頭之間共享相同的 PE（每個頭中的頻率從 1 到 10000 變化）比獨立的 PE（所有通道中的頻率從 1 到 10000 變化）要好，結果如表 11d 所示。
特徵抽象策略：本文在大引數規模的模型（-L）上比較了兩種常見的特徵提取策略：類別 token 和 GAP ，結果如表 11e 中所示，使用類別 token 比 GAP 更好，這與 PEG [13] 中所得到的結論不同。然而，兩種方法的訓練設定截然不同。本文還使用 DeiT3-L 進行了額外的實驗，得到了類似的結論。本文進一步評估 “小型”（-S）和 “基礎”（-B）模型的效能。有趣的是，在小模型中觀察到了相反的結論，有理由懷疑 DeiT3 中使用的較高丟棄路徑率（drop-path rate）使得諸如 GAP 之類的無引數抽象方法（parameter-free abstraction）難以達到應有的效果。
位置編碼策略：本文還在金字塔結構 VisionLLaMA-S 上評估了其他絕對位置編碼策略，例如可學習位置編碼 和 PEG。由於存在強大的基線，本文使用 “小” 模型，結果顯示在表 11f 中：可學習的 PE 不會提高效能，PEG 將基線從 81.6% 略微提高到 81.8%。出於三個原因，本文並沒有將 PEG 作為基本組成部分。首先，本文嘗試對 LLaMA 進行最小程度的修改。其次，本文的目標是為 ViT 等各種任務提出一種通用方法。對於像 MAE 這樣的遮蔽影像框架（masked image frameworks），PEG 增加訓練成本，並可能損害下游任務上的效能。原則上，可以在 MAE 框架下應用稀疏 PEG，但會引入部署不友好的運算元。稀疏卷積是否與其密集版本一樣包含足夠的位置資訊仍然是一個未解決的問題。第三，無模態束縛的設計為進一步研究涵蓋文字和視覺之外的其他模態鋪平了道路。
對輸入尺寸的敏感性：在未訓練的前提下，本文進一步比較了增大解析度和常用解析度的效能，結果如表 12 中所示。這裡使用了金字塔結構 transformer，因為其在下游任務中比對應的非層次結構版本更受歡迎。1D-RoPE 的效能因解析度變化而受到嚴重影響並不奇怪。α = 2 的 NTK-Aware 插值實現了與 2D-RoPE 類似的效能，2D-RoPE 實際上是 NTKAware (α = 1)。AS2DRoPE 展示出了在較大解析度上的最佳效能。
討論
收斂速度：對於影像生成，本文研究了不同訓練步數下的表現，分別在 100k、200k、300k 和 400k 次迭代時儲存權重來計算保真度指標。由於 SDE 明顯慢於 ODE，因此本文選擇使用 ODE 取樣器。表 10 中的結果表明 VisionLLaMA 在所有模型上的收斂速度都比 ViT 快得多。具有 30 萬次訓練迭代的 SiT-LLaMA 效能甚至優於具有 40 萬次訓練次數的的基線模型。
本文還與圖 4 中 ImageNet 上使用 DeiT3-Large 全監督訓練 800 個 epoch 的 top-1 精度進行了比較，表明 VisionLLaMA 比 DeiT3-L 收斂得更快。本文進一步比較了 MAE 框架下 ViT-Base 模型的 800 個 epoch 的訓練損失，並在圖 5 中進行了說明。VisionLLaMA 在開始時具有較低的訓練損失，並將該趨勢保持到最後。