自回歸模型迎來全能選手!FlexVAR一模型通吃圖像生成/修補,推理速度與質(zhì)量自由調(diào)節(jié)
解讀:AI生成未來
文章鏈接:https://arxiv.org/pdf/2502.20313
Git鏈接:https://github.com/jiaosiyu1999/FlexVAR
亮點直擊
- 生成各種分辨率和寬高比的圖像,甚至可以超過訓練分辨率;
- 支持圖像到圖像的任務(wù),例如圖像修復(fù)、圖像細化和圖像擴展,而無需微調(diào);
- 享受靈活的推理步驟,允許通過減少步驟加速推理或通過增加步驟提高圖像質(zhì)量。
總結(jié)速覽
原文來源
本文作者來自北京交通大學、悉尼科技大學、美團、佐治亞理工學院。
解決的問題
- 靈活性不足:現(xiàn)有的視覺自回歸模型(如VAR)通常只能生成固定分辨率的圖像,缺乏生成不同分辨率和寬高比圖像的靈活性。
- 殘差預(yù)測的限制:現(xiàn)有模型采用殘差預(yù)測的方式,依賴于固定的步驟設(shè)計,限制了圖像生成的適應(yīng)性和靈活性。同時,不同尺度的殘差缺乏語義連續(xù)性,可能限制模型表示多樣圖像變化的能力。
提出的方案
- FlexVAR模型:提出了一種新的視覺自回歸建模范式FlexVAR,摒棄了殘差預(yù)測,改為在每個步驟中直接預(yù)測真實值(ground-truth),從而確保相鄰尺度之間的語義連貫性。
- 多尺度約束的VQVAE:設(shè)計了一個可擴展的VQVAE tokenizer,增強了VQVAE對不同隱空間尺度的魯棒性,使得圖像可以在任意分辨率下重建。
- 可擴展的2D位置嵌入:提出了可擴展的2D位置嵌入方法,使用2D可學習查詢初始化,使得尺度自回歸建模可以擴展到各種分辨率和步驟,包括訓練時未使用的分辨率和步驟。
應(yīng)用的技術(shù)
- 自回歸建模:利用自回歸模型學習圖像的概率分布,生成圖像。
- VQVAE(向量量化變分自編碼器):用于圖像的分割和重建,增強模型對不同尺度的適應(yīng)性。
- Transformer架構(gòu):用于學習多尺度隱空間特征的概率分布,建模下一個尺度的真實值。
- 2D位置嵌入:引入2D可學習查詢,擴展模型的分辨率和步驟適應(yīng)性。
達到的效果
- 生成多樣性:FlexVAR能夠生成各種分辨率和寬高比的圖像,甚至超過訓練圖像的分辨率。
- 多任務(wù)支持:支持多種圖像到圖像的任務(wù),包括圖像細化、圖像修復(fù)、圖像擴展等。
- 自適應(yīng)步驟:可以根據(jù)需要調(diào)整自回歸步驟,實現(xiàn)更快的推理速度或更高的圖像質(zhì)量。
- 性能提升:在ImageNet 256×256基準測試中,1.0B模型優(yōu)于其VAR對應(yīng)模型。在零樣本轉(zhuǎn)移生成過程中,13步生成的圖像性能進一步提升,F(xiàn)ID達到2.08,優(yōu)于現(xiàn)有的自回歸模型和流行的擴散模型。
- 零樣本轉(zhuǎn)移:在ImageNet 512×512基準測試中,F(xiàn)lexVAR在零樣本轉(zhuǎn)移情況下取得了與全監(jiān)督訓練的VAR 2.3B模型相競爭的結(jié)果。
方法
FlexVAR 概述
FlexVAR 是一種靈活的視覺自回歸圖像生成范式,它通過真實值(ground-truth)預(yù)測而非殘差預(yù)測來實現(xiàn)自回歸學習,從而能夠在任何步驟獨立生成合理的圖像。在本文的方法中:(1) 一個可擴展的 VQVAE tokenizer 將輸入圖像量化為多尺度的 token 并重建圖像。(2) 一個 FlexVAR transformer 通過尺度自回歸建模進行訓練,去除了殘差預(yù)測。
多尺度量化與圖像重建
主流的 VQVAE tokenizer 在單一分辨率下表現(xiàn)良好。然而,當擴展隱空間時,它們通常無法重建圖像(如下圖 3 所示)。這一觀察促使我們探索一種可擴展的 tokenizer,能夠?qū)⑤斎雸D像量化為多尺度的 token 并重建圖像。具體來說,所提出的可擴展 tokenizer 首先將圖像編碼為多尺度的隱空間,然后使用量化器將隱空間特征轉(zhuǎn)換為離散的 token,最后使用解碼器從每個尺度的離散 token 中重建原始圖像。
無殘差的視覺自回歸建模
實驗
實現(xiàn)細節(jié)
最先進的圖像生成
在 ImageNet-1K 基準測試中將 FlexVAR 與現(xiàn)有的生成方法進行比較,包括 GAN、擴散模型、隨機掃描、光柵掃描和尺度掃描自回歸模型。結(jié)果如下表 2 和表 3 所示。
ImageNet 256×256 的總體比較。為了確保公平比較,在表 2 中展示了參數(shù)規(guī)模小于 1B 的模型。我們的 FlexVAR 在所有生成方法中實現(xiàn)了最先進的性能,并且與 VAR 相比表現(xiàn)尤為出色。具體來說,在不同模型規(guī)模下,我們分別實現(xiàn)了 FID 提升 -0.45、-0.56 和 -0.12。
零樣本推理(更多步驟)。在未訓練的情況下使用 13 個步驟進行圖像生成,如表 2 的最后一行所示。FlexVAR 可以靈活地采用更多步驟來提高圖像質(zhì)量。通過使用 13 個推理步驟,F(xiàn)lexVAR 進一步將性能提升至 2.08 FID 和 315 IS,展現(xiàn)了強大的靈活性和泛化能力。具體的步驟設(shè)計詳見附錄。
ImageNet 512×512 基準測試的零樣本推理。使用 FlexVAR-d24 生成 512×512 圖像,并在未訓練的情況下在 ImageNet-512 基準測試上進行評估,如前面表 3 所示。盡管 FlexVAR 僅在分辨率 ≤256×256 上訓練且僅有 1.0B 參數(shù),但我們的 FlexVAR-d24 表現(xiàn)出了與 VAR 相競爭的性能。
消融實驗
本文對 FlexVAR 中的各種設(shè)計選擇進行了消融實驗。由于計算資源有限,在 下表 4、5以及前面表6 中報告了使用短期訓練方案(即 40 輪訓練,約 70K 次迭代)得到的結(jié)果。
組件級消融實驗
為了理解每個組件的影響,從標準 VAR 開始,并逐步添加每個設(shè)計:
- 基線:VAR 采用殘差預(yù)測范式,表現(xiàn)良好(第 1 個結(jié)果),但其在圖像生成方面的靈活性未達到預(yù)期)。
- 預(yù)測類型:直接將預(yù)測類型轉(zhuǎn)換為 GT 是不可行的(第 2 和第 3 個結(jié)果)。我們采用了 VAR 和 Llamagen 的 VQVAE 分詞器,二者均表現(xiàn)較差。這并不意外,因為當前的分詞器缺乏對具有不同隱空間的圖像的魯棒性,而我們在訓練過程中強制這些分詞器獲取多尺度隱空間特征(在前面圖3中提供了詳細分析)。
- 分詞器:可擴展的分詞器在訓練過程中獲取了合理的多尺度隱空間特征,使 FID 下降了 13.87(第 4 個結(jié)果)。然而,靈活的圖像生成仍未實現(xiàn)。
- 位置嵌入:引入我們的可擴展 PE(最后一個結(jié)果)后,圖像生成的靈活性顯著提升,并進一步將 FID 降低至 3.71。
使用不同 VQVAE 重建圖像
在前面圖3中,通過縮放 VQVAE 分詞器的隱空間特征來重建多尺度圖像?,F(xiàn)有的 VQVAE 分詞器通常不支持在小到大范圍內(nèi)縮放隱空間特征。VAR 的 VQVAE 采用基于殘差的訓練策略,直接將其應(yīng)用于非殘差圖像重建未能達到預(yù)期效果(第 1 行)。Llamagen 的 VQVAE 分詞器僅在原始隱空間下表現(xiàn)出色,表明其不適用于尺度自回歸建模(第 2 行)。
將 FlexVAR 遷移至 Mamba
最新研究 AiM 使用 Mamba 架構(gòu)進行基于 token 的自回歸建模。受此啟發(fā),對 FlexVAR 進行了 Mamba 適配,并評估其性能(前面表 5)。在類似的模型參數(shù)下,Mamba 相較于 Transformer 模型展現(xiàn)了競爭力,這表明 GT 預(yù)測范式可以有效適應(yīng) Mamba 這樣的線性注意力機制。然而,考慮到該 Mamba 架構(gòu)并未體現(xiàn)速度優(yōu)勢,我們未將其集成到最終版本中。
Mamba 具有固有的單向注意力機制,使得圖像 token 在同一尺度下無法實現(xiàn)全局注意力。為了解決這一問題,我們在不同的 Mamba 層中采用了 8 條掃描路徑來捕獲全局信息。具體的 Mamba 架構(gòu)細節(jié)見附錄。
位置嵌入
在前面表 6 中,對多種步長 PE 和 x-y 坐標 PE 進行了實驗。為了使模型在推理步驟上更具魯棒性,并能夠生成任意分辨率的圖像,移除了固定長度的步長嵌入(第 2 行的結(jié)果),性能僅出現(xiàn)輕微變化。本文采用了一種類似于 ViT的非參數(shù)變體,與可學習變體相比,F(xiàn)ID 僅相差 0.03。
GT 預(yù)測范式分析
收斂速度
本文對比了 VAR 和 FlexVAR 的訓練損失,如下圖4 所示。FlexVAR 展示了顯著更低的損失值和更快的收斂速度,這表明預(yù)測真實值(ground-truth)而非殘差對訓練更加友好。這可能是由于不同尺度的殘差缺乏語義連續(xù)性,而這種隱式預(yù)測方法可能會限制訓練的收斂速度。
生成任意分辨率的圖像
在下圖1 和 圖5 中,使用 FlexVAR-d24 生成了不同分辨率的圖像。通過控制推理步驟,F(xiàn)lexVAR 能夠生成任意分辨率的圖像,盡管其僅在分辨率 ≤ 256px 的圖像上進行了訓練。生成的圖像在多個尺度上表現(xiàn)出較強的語義一致性,并且更高分辨率的圖像顯示出更清晰的細節(jié)。
生成任意比例的圖像
在 上圖1 和 下圖6 中,展示了不同長寬比的生成樣本。通過在推理過程中控制每一步的長寬比,F(xiàn)lexVAR 能夠生成不同長寬比的圖像,體現(xiàn)了其在圖像生成中的靈活性和可控性。
生成任意推理步長的圖像
在下圖7中,研究了使用 6 到 16 步生成 256×256 圖像的 FID 和 IS,實驗涵蓋了 3 種不同的模型尺寸(depth 16、20、24)。隨著步長的增加,生成圖像的質(zhì)量提升。更大規(guī)模的模型(如 FlexVAR-d24)提升更為顯著,因為更大的 Transformer 結(jié)構(gòu)能夠?qū)W習更復(fù)雜和精細的圖像分布。在訓練過程中,我們最多使用 10 步以避免 OOM(內(nèi)存溢出)問題。然而,在推理階段,使用 13 步可使 FID 降低 0.13,這表明我們的 FlexVAR 對推理步長具有較強的適應(yīng)性,可以通過減少步長加速圖像生成,或通過增加步長獲得更高質(zhì)量的圖像。更多步長設(shè)計的細節(jié)請見附錄。
高分辨率圖像細化
在下圖8 中,輸入低分辨率圖像(如 256×256),并讓 FlexVAR 輸出高分辨率的細化圖像。盡管僅在 ≤ 256px 的圖像上進行了訓練,F(xiàn)lexVAR 仍然能夠通過提高輸入圖像分辨率來細化圖像細節(jié),例如示例中狗的眼睛。這展示了 FlexVAR 在圖像到圖像生成任務(wù)中的高度靈活性。
圖像修補與擴展
對于圖像修補(in-painting)和擴展(out-painting),對掩碼外的 GT token 進行教師強制(teacher-force),并讓模型僅生成掩碼內(nèi)的 token。同時,還注入了類別標簽信息。結(jié)果如 下圖9 所示。在不修改架構(gòu)設(shè)計或訓練方式的情況下,F(xiàn)lexVAR 在這些圖像到圖像的任務(wù)上表現(xiàn)良好。
圖像擴展
對于圖像擴展任務(wù),生成目標類別的 1:2 長寬比圖像,并在中心強制使用 ground-truth token。FlexVAR 在圖像擴展方面表現(xiàn)良好,這表明其具有較強的泛化能力。
失敗案例
當生成的圖像分辨率達到訓練分辨率的 3 倍或更高時,F(xiàn)lexVAR 生成的圖像會出現(xiàn)明顯的波浪狀紋理和模糊區(qū)域(見下圖11)。這種失敗可能是由于當前訓練數(shù)據(jù)集結(jié)構(gòu)過于同質(zhì)化,例如 ImageNet-1K 通常缺乏從粗到細的多尺度目標,因此在生成高分辨率圖像的細節(jié)時會產(chǎn)生誤差。
推測如果使用包含更精細細節(jié)的復(fù)雜數(shù)據(jù)集進行訓練,模型可能會在更高分辨率的生成任務(wù)上表現(xiàn)得更加穩(wěn)健。
結(jié)論
本文提出了 FlexVAR,一種靈活的視覺自回歸圖像生成范式,實現(xiàn)了無需殘差預(yù)測的自回歸學習。本文設(shè)計了可擴展的 VQVAE tokenizer 和 FlexVAR-Transformer。該 ground-truth 預(yù)測范式賦予了自回歸模型極大的靈活性和可控性,使其能夠在訓練時未涉及的分辨率、長寬比和推理步長下進行圖像生成。此外,它還能零樣本遷移至多種圖像到圖像的生成任務(wù)。希望 FlexVAR 能成為一個穩(wěn)固的基線,并促進視覺自回歸建模及相關(guān)領(lǐng)域的研究。
局限性
發(fā)現(xiàn),當生成的圖像分辨率達到訓練圖像的3倍或更高時,會出現(xiàn)明顯的波浪狀紋理。這一問題可能是由于ImageNet-1K訓練集結(jié)構(gòu)過于同質(zhì)化導致的。未來將進一步研究如何在零樣本高分辨率圖像生成任務(wù)中保持穩(wěn)定性。
本文轉(zhuǎn)自AI生成未來 ,作者:AI生成未來
