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