首次超越擴散模型和非自回歸Transformer模型！字節(jié)開源RAR：自回歸生成最新SOTA！
文章鏈接：https://arxiv.org/pdf/2411.00776
項目鏈接：https://yucornetto.github.io/projects/rar.html
代碼&模型鏈接：https://github.com/bytedance/1d-tokenizer
亮點直擊
- RAR（隨機排列自回歸訓練策略），這是一種改進的訓練策略，使得標準的自回歸圖像生成器能夠?qū)崿F(xiàn)SOTA性能。
- 引入雙向上下文學習：RAR通過最大化所有可能的分解順序的期望似然值，打破了自回歸模型在視覺任務中的單向上下文限制，使模型能夠在圖像生成中更有效地利用雙向上下文信息。
- 保持與語言建模框架的兼容性：RAR在提升圖像生成性能的同時，保留了自回歸建模的核心結構，它與大語言模型（LLM）的優(yōu)化技術（如KV-cache）完全兼容，相比于MAR-H或MaskBit，采樣速度顯著更快，同時保持了更好的性能，便于在多模態(tài)統(tǒng)一模型中應用。
- 創(chuàng)新的退火訓練策略：訓練初期將輸入序列隨機排列，隨著訓練逐步回歸至標準光柵順序。這一過程使模型在各類上下文排列中均能獲得優(yōu)化，提升生成質(zhì)量。
- 顯著的性能提升：在 ImageNet-256 基準測試中，RAR實現(xiàn)了1.48的 FID 分數(shù)，顯著超越了先前的自回歸圖像生成器，顯示了其在圖像生成任務中的突破性改進。
總結速覽
解決的問題：
RAR（隨機自回歸建模）旨在提升圖像生成任務的表現(xiàn)，同時保持與語言建?？蚣艿耐耆嫒菪?。
提出的方案：
RAR采用了一種簡單的方法，通過標準的自回歸訓練過程并結合下一個 token 預測目標，將輸入序列隨機打亂到不同的分解順序。在訓練過程中，使用概率 r 將輸入序列的排列順序隨機化，其中 r 從 1 開始并逐漸線性衰減至 0，以讓模型學習所有分解順序的期望似然值。
應用的技術：
RAR在訓練中應用了一種退火策略，使模型能夠更好地利用雙向上下文信息來建模，同時仍然保持自回歸建?？蚣艿耐暾裕瑥亩鴮崿F(xiàn)了語言建模的完全兼容性。
達到的效果：
在 ImageNet-256 基準測試中，RAR 獲得了 1.48 的 FID 分數(shù)，超越了之前最先進的自回歸圖像生成器，并優(yōu)于領先的基于擴散和掩碼Transformer的方法。
方法
背景
RAR: 隨機自回歸建模
視覺信號天然存在雙向關聯(lián)，因此全局上下文建模在視覺任務中至關重要。然而，傳統(tǒng)的自回歸模型依賴因果注意力掩碼，僅允許 token 序列呈單向依賴，這與視覺數(shù)據(jù)的雙向關聯(lián)性不符。已有研究表明，對于視覺模態(tài)，雙向注意力顯著優(yōu)于因果注意力。
此外，圖像 token 排列成因果序列時沒有統(tǒng)一的“正確”方式，常用的光柵掃描順序雖有效果，但在訓練過程中引入了偏差，因為每個 token 僅依賴于掃描順序中的前序 token，限制了模型學習來自其他方向的依賴關系。
為了解決這些問題，本文提出了一種隨機自回歸建模方法，將雙向上下文的優(yōu)化目標納入自回歸建模。
與 BERT 風格 或 MaskGIT 風格方法不同，本文的方法采用了置換目標的方法，在所有可能的分解順序上以自回歸方式訓練模型。這使得模型在保持自回歸框架的同時，能夠在期望上收集雙向上下文信息。公式表示如下：
雖然方法簡單，但這種修改顯著提升了圖像生成性能，突顯了雙向上下文在提升自回歸圖像生成器能力方面的重要性。與自然語言處理 (NLP) 中的自回歸訓練觀察結果一致。
討論：盡管置換目標允許在自回歸框架中實現(xiàn)雙向上下文學習，但在生成過程中完全捕獲“全局上下文”仍具挑戰(zhàn)性。這是因為在生成過程中，總會有一些 tokens 在其他 tokens 之前生成，無法完全利用全局上下文。這種限制不僅存在于自回歸方法中，也存在于非自回歸模型中。重新采樣或精煉等技術可能有助于確保每個 token 都能在充分的上下文下生成。然而，這些設計可能會增加系統(tǒng)的復雜性，因此探索這些解決方案超出了本文的范圍，留待未來研究。
隨機退火：雖然通過置換的隨機自回歸訓練使模型能夠在單向框架中捕獲雙向上下文，但它可能會因兩個主要因素而引入次優(yōu)的視覺生成行為：(1) 可能的置換數(shù)量極多，可能導致模型過于關注如何處理不同的置換順序，而不是提高生成質(zhì)量。例如，對于長度為 256 的 token 序列，可能的置換數(shù)為 (256! > 10^{506})，這會讓模型不堪重負，降低訓練效率。(2) 盡管圖像可以按任意順序處理，某些掃描順序往往優(yōu)于其他順序。例如，[22] 評估了六種不同的掃描順序（行優(yōu)先、向內(nèi)螺旋、向外螺旋、Z 曲線、子采樣和交替順序），并發(fā)現(xiàn)行優(yōu)先（即柵格順序）始終表現(xiàn)最佳，這使其成為視覺生成中最廣泛使用的順序。
為了解決這些問題，本文提出了隨機退火策略，以平衡置換的隨機性與行優(yōu)先順序的已知效果。此方法引入了一個控制隨機置換和行優(yōu)先順序使用概率的單一參數(shù) r。在訓練開始時，r = 1，意味著模型完全使用隨機置換。隨著訓練的進行，r 線性衰減至 0，逐步將模型轉換為行優(yōu)先順序。具體而言，定義了r 的訓練調(diào)度，由兩個超參數(shù) ??start??? 和 ??end?? 控制，分別表示 r 開始退火和退火結束的訓練輪次。公式如下：
其中，??epoch??? 表示當前的訓練輪次。在實驗中，我們將對超參數(shù) ??start??? 和 ??end?? 進行消融實驗。該調(diào)度策略允許模型在初期探索多種隨機置換，以更好地學習雙向表示，最終收斂到更有效的行優(yōu)先掃描順序，從而提升視覺生成質(zhì)量，類似于其他典型的自回歸方法。值得注意的是，該策略不僅提升了生成性能，還保持了與先前工作中使用的標準掃描順序的兼容性。
實驗結果
本節(jié)首先介紹方法的實現(xiàn)細節(jié)。接這展示了關于關鍵設計選擇的消融研究。再討論了主要結果，最后還包括了縮放研究和可視化內(nèi)容。
實現(xiàn)細節(jié)
本文在語言建模自回歸框架的基礎上實現(xiàn)了RAR方法，做了最小的改動。
VQ分詞器：采用了與先前工作 [10, 22] 相似的VQ分詞器，將輸入圖像轉換為離散的token 序列。我們使用的是基于CNN的MaskGIT-VQGAN [10] 分詞器，使用在ImageNet上訓練的官方權重。該分詞器將256 × 256的圖像轉化為256個離散的token （下采樣因子為16），并且字典大?。丛~匯表大?。?024。
自回歸Transformer：本文使用不同配置的視覺Transformer（ViT），包括RAR-S（133M）、RAR-B（261M）、RAR-L（461M）、RAR-XL（955M）和RAR-XXL（1499M）。對于所有這些模型變體，我們在自注意力模塊中應用了因果注意力掩碼，并使用QK層歸一化來穩(wěn)定大規(guī)模模型的訓練。為了加速實驗，在消融研究中使用了普通的ViT，而在最終模型中增強了AdaLN 。架構配置和模型大小可以在下表1中找到。
位置嵌入：本文為ViT中的原始位置嵌入和目標感知位置嵌入使用了可學習的嵌入。值得注意的是，由于我們的模型在訓練結束后會退火到基于柵格順序的自回歸圖像生成，最終這兩種位置嵌入可以合并為一個，使得最終模型與傳統(tǒng)的自回歸圖像生成器相同。
數(shù)據(jù)集：研究者們在ImageNet-1K訓練集上訓練我們的模型，該數(shù)據(jù)集包含128,1167張圖像，涵蓋1000個物體類別。我們使用MaskGIT-VQGAN分詞器對整個訓練集進行預分詞，以加速訓練。對于消融研究，僅使用中心裁剪和水平翻轉數(shù)據(jù)增強進行預分詞，而對于最終模型，使用了十裁剪變換來增強數(shù)據(jù)集的多樣性。
訓練協(xié)議：研究者們所有模型變體使用相同的訓練超參數(shù)。模型使用批量大小2048訓練400個epoch（250k步）。在前100個epoch（熱身階段）內(nèi)，學習率從0線性增加到4 × 10??，然后按照余弦衰減計劃逐漸衰減至1 × 10??。使用AdamW優(yōu)化器，其中beta1為0.9，beta2為0.96，權重衰減為0.03。我們對梯度進行了裁剪，最大梯度范數(shù)為1.0。在訓練過程中，類別條件會以0.1的概率被丟棄。消融研究和主要結果中所有RAR模型變體的訓練設置保持一致。
采樣協(xié)議：本文使用[18]的評估代碼對50,000張圖像進行FID計算。我們不使用任何基于top-k或top-p的過濾技術。還遵循先前的工作使用無分類器指導。在消融研究中，使用更簡單的線性指導調(diào)度，而在最終模型中使用改進的冪余弦指導調(diào)度。
消融研究
本文研究了RAR的不同配置，包括隨機退火策略和RAR最終收斂的掃描順序。
隨機退火策略：在下表2中，比較了不同的隨機退火策略。采用了線性衰減的調(diào)度，并通過改變超參數(shù) ??start??? 和 ??end??? 來研究何時應該開始和結束隨機化退火，具體定義見公式(5)。對于持續(xù)400個epoch的訓練，我們枚舉了每100個epoch的所有可能組合。例如，當 ??start = 200??? 和 ??end = 300??? 時，模型在前200個epoch采用隨機排列，在后100個epoch采用柵格順序。在第200到300個epoch之間，模型通過以概率 ??r??? 進行隨機排列，或者以概率 ??1?r??? 進行柵格順序訓練，其中 ??r??? 按照公式(5)計算。值得注意的是，當 ??start = end = 0??? 時，模型僅使用柵格順序進行訓練，即標準的自回歸訓練；當 ??start = end = 400??? 時，模型始終使用隨機排列的輸入序列進行訓練。兩種情況都是提出的隨機退火方法的重要基準，分別達到了FID得分3.08和3.01。令人感興趣的是，我們觀察到所有其他變體都比這兩個基準取得了顯著的改進。例如，簡單地將前100個epoch的柵格順序替換為隨機排列（即，??start = 100??? 和 ??end = 100??），就將FID得分提高到了2.48，提升了0.6。此外，模型傾向于保留一些初期的epoch進行純隨機排列訓練，并且在最后一些epoch更好地適應柵格順序，這通常比其他變體表現(xiàn)更好。所有結果表明，通過引入帶有排列目標的隨機化自回歸訓練，有助于自回歸視覺生成器的性能，并提升FID得分，這得益于改進的雙向表示學習過程。
此外，在所有變體中，發(fā)現(xiàn)當 ??start = 200??? 和 ??end = 300?? 時表現(xiàn)最佳，將基準（純柵格順序）的FID從3.08提高到2.18。該策略將稍多的計算分配給隨機排列順序的訓練，并將最后100個epoch專注于純柵格順序。因此，我們默認采用這種退火策略用于所有RAR模型。
不同的掃描順序（除了柵格順序）：盡管行主序（即柵格掃描）一直是視覺生成中事實上的掃描順序，但缺乏系統(tǒng)的研究來比較它與其他掃描順序的優(yōu)劣。我們注意到，四年前的工作 [22] 進行了類似的研究。然而，考慮到近年來生成模型取得的顯著進展，值得重新審視這一結論。具體來說，我們考慮了6種不同的掃描順序（行主序、螺旋內(nèi)、螺旋外、Z曲線、子采樣和替代掃描順序），這些掃描順序是RAR可能最終收斂的目標。本文沒有像那樣報告訓練損失和驗證損失作為對比指標，而是直接評估它們的生成性能。結果總結在下表3中。有趣的是，我們觀察到所有變體的得分都相當不錯，這表明RAR能夠處理不同的掃描順序?？紤]到行主序（柵格掃描）仍然在其他掃描順序中表現(xiàn)出優(yōu)勢，我們因此為所有最終RAR模型使用柵格掃描順序。
主要結果
本文報告了RAR與最先進的圖像生成器在ImageNet-1K 256×256基準測試上的結果。
如下表4所示，RAR相較于之前的AR圖像生成器表現(xiàn)出了顯著更好的性能。具體來說，最緊湊的RAR-B（僅有261M參數(shù)）就達到了FID得分1.95，已經(jīng)顯著超越了當前最先進的AR圖像生成器LlamaGen-3B-384（3.1B，F(xiàn)ID 2.18，裁剪尺寸384）和 Open-MAGVIT2-XL（1.5B，F(xiàn)ID 2.33），并且分別減少了91%和81%的模型參數(shù)。它還超越了廣泛使用的擴散模型，例如DiT-XL/2（FID 1.95 vs. 2.27）和SiT-XL（FID 1.95 vs. 2.06），并且僅使用了相對于這些模型的39%的參數(shù)。
在表4中，進一步探討了不同模型尺寸下的RAR（從261M到1.5B），我們觀察到RAR在不同尺寸下具有強大的可擴展性，并且隨著模型尺寸的增大，性能不斷提升。特別地，最大的變體RAR-XXL在ImageNet基準測試上創(chuàng)下了新的最先進結果，F(xiàn)ID得分為1.48。與其他兩種近期方法VAR和MAR相比，這兩種方法都嘗試通過改進AR公式來提升視覺生成質(zhì)量，RAR不僅在性能上表現(xiàn)更優(yōu)（RAR的FID為1.48，而VAR為1.73，MAR為1.55），而且保持了整個框架與語言建模的兼容性，因此更適合將成熟的優(yōu)化和加速技術應用于大型語言模型，從而推動視覺生成的發(fā)展。
此外，RAR在不同框架中的表現(xiàn)超越了最先進的視覺生成器。它在對比領先的自回歸模型、擴散模型和掩蔽Transformer模型時，表現(xiàn)得更好，分別超越了LlamaGen-3B-384、MDTv2-XL/2和 MaskBit（RAR的FID為1.48，相比之下LlamaGen為2.18，MDTv2為1.58，MaskBit為1.52）。據(jù)所知，這是首次語言建模風格的自回歸視覺生成器超越最先進的擴散模型和掩蔽Transformer模型。
采樣速度：自回歸方法的一個關鍵優(yōu)勢是能夠利用LLM中已建立的優(yōu)化技術，如KV緩存。在表5中，我們將RAR與其他類型的生成模型進行采樣速度比較（以每秒生成圖像數(shù)為衡量標準），包括擴散模型、masked transformer器、VAR和 MAR。其中，自回歸模型（RAR）和VAR模型（VAR-d30）與KV緩存優(yōu)化兼容，因此在生成速度上顯著優(yōu)于其他方法。如表5所示，RAR不僅在FID得分上達到了最先進的水平，同時在生成速度上也大大超越了其他方法。例如，在FID得分約為1.5時，MaskBit和 MAR-H的生成速度分別為每秒0.7和0.3張圖像。相比之下，RAR-XL不僅獲得了更好的FID得分，還能每秒生成8.3張高質(zhì)量視覺樣本——比MaskBit快11.9倍，比MAR-H快27.7倍。最大的RAR變體RAR-XXL進一步提高了FID得分，同時保持了顯著的速度優(yōu)勢，速度是MaskBit的9.1倍，是MAR-H的21.3倍。此外，RAR可能進一步受益于LLM優(yōu)化技術，例如vLLM，這一點與其他AR方法一致。
擴展性行為
本文研究了RAR的擴展性行為。具體來說，我們繪制了訓練損失曲線和FID得分曲線（有無無分類器引導的情況）如下圖4所示。如圖所示，我們觀察到RAR在不同模型尺寸下均表現(xiàn)出良好的擴展性，較大的模型尺寸在訓練損失和FID得分上持續(xù)表現(xiàn)出較好的性能，無論是否使用無分類器引導增強。我們注意到，由于RAR保持了AR公式和框架的完整性，它也繼承了AR方法的可擴展性。
可視化
在下圖5中可視化了不同RAR變體生成的樣本，結果表明RAR能夠生成高質(zhì)量、具有高度保真度和多樣性的樣本。更多可視化結果見附錄。
結論
本文提出了一種簡單而有效的策略，以增強與語言建模兼容的自回歸圖像生成器的視覺生成質(zhì)量。通過采用隨機化排列目標，本文的方法在保持自回歸結構的同時，改善了雙向上下文學習。因此，所提出的RAR模型不僅超越了以前的最先進自回歸圖像生成模型，還超過了領先的非自回歸Transformer和擴散模型。希望本研究有助于推動自回歸Transformer朝著視覺理解與生成統(tǒng)一框架的方向發(fā)展。
本文轉自AI生成未來 ，作者：AI生成未來
