codebook從崩潰到高效利用！南大&清華&騰訊聯合打造IBQ：自回歸生成最強視覺分詞器
論文鏈接：https://arxiv.org/pdf/2412.02692github鏈接：https://github.com/TencentARC/SEED-Voken
亮點直擊
- 提出了一種簡單而有效的向量量化方法，稱為索引反向傳播量化（Index Backpropagation Quantization，IBQ），用于訓練可擴展的視覺分詞器。
- 通過增加碼本大小、編碼維度和模型規模來研究IBQ的擴展特性。IBQ首次訓練了一個超大碼本（），具有大維度（256）和高使用率，實現了最先進的重建性能。
- 展示了一系列從300M到2.1B的基礎自回歸圖像生成模型，顯著超越了競爭方法，例如LlamaGen和 Open-MAGVIT2。
總結速覽
解決的問題
現有的向量量化（VQ）方法在可擴展性方面存在困難，主要由于訓練過程中僅部分更新的碼本的不穩定性。隨著利用率的降低，碼本容易崩潰，因為未激活代碼與視覺特征之間的分布差距逐漸擴大。
提出的方案
提出了一種新的向量量化方法，稱為索引反向傳播量化（Index Backpropagation Quantization，IBQ），用于碼本embedding和視覺編碼器的聯合優化。通過在編碼特征與碼本之間的單熱編碼分類分布上應用直通估計器，確保所有代碼都是可微的，并與視覺編碼器保持一致的潛空間。
應用的技術
- 使用直通估計器在單熱編碼分類分布上進行優化，使得所有代碼可微。
- 通過IBQ實現碼本embedding和視覺編碼器的聯合優化。
- 研究了IBQ在增加碼本大小、編碼維度和模型規模方面的擴展特性。
達到的效果
- 在標準ImageNet基準上的實驗表明，IBQ在重建（1.00 rFID）和自回歸視覺生成方面取得了具有競爭力的結果。
- 展示了一系列從300M到2.1B的基礎自回歸圖像生成模型，顯著超越了競爭方法，如LlamaGen和Open-MAGVIT2。
效果展示
下圖的上半部分展示了在1024×1024分辨率下，IBQ分詞器在Unsplash數據集上的測試結果。下半部分則展示了IBQ分詞器在256×256分辨率下，針對Imagenet數據集的測試結果。(a)表示原始圖像，(b)表示重建圖像。
Imagenet上256×256類條件生成樣本效果：
方法
我們的框架由兩個階段組成。第一階段是通過索引反向傳播量化學習一個具有高碼本利用率的可擴展視覺分詞器。在第二階段，我們使用自回歸變換器通過下一個標記預測進行視覺生成。
Preliminary
這些方法采用的部分更新策略（即僅優化選定的代碼）逐漸擴大了視覺特征與未激活代碼之間的分布差距。這會導致訓練期間的不穩定性，因為碼本崩潰會阻礙視覺分詞器的可擴展性。
索引反向傳播量化
量化。 為了確保在訓練過程中碼本與編碼特征之間的一致分布，我們引入了一種全代碼更新方法，即索引反向傳播量化（Index Backpropagation Quantization, IBQ）。IBQ的核心在于將梯度反向傳播到碼本的所有代碼，而不僅僅是選定的代碼。算法1提供了IBQ的偽代碼。
具體來說，我們首先對給定的視覺特征與所有代碼embedding進行點積運算作為logits，并通過softmax函數獲得概率（soft one-hot）。
然后我們將soft one-hot 類別分布的梯度復制到hard one-hot索引上：
給定索引，量化后的特征通過以下方式獲得：
通過這種方式，我們可以通過索引將梯度傳遞到碼本的所有代碼上。通過索引反向傳播量化，整個碼本和編碼特征的分布在整個訓練過程中保持一致，從而獲得較高的碼本利用率。
訓練損失。
與 VQGAN類似，分詞器的優化由多種損失的組合完成：
如下圖 3 所示，現有的 VQ 方法在每次反向過程中僅優化有限數量的代碼以接近編碼特征。這逐漸擴大了未激活代碼和編碼特征之間的分布差距，最終導致碼本崩潰。隨著代碼維度和碼本大小的增加，這種情況變得更加嚴重。我們不是直接將直通估計器 [1] 應用于選定的代碼，而是將這種參數化方法應用于視覺特征和所有碼本embedding之間的分類分布，以使梯度能夠反向傳播到所有代碼。通過這種方式，整個碼本和編碼特征之間的分布在整個訓練過程中保持一致。因此，IBQ 實現了具有高代碼維度和利用率的極大碼本。
自回歸Transformer
其中，c 是條件，例如類別標簽。請注意，由于我們專注于視覺分詞器，因此我們采用類似于 Llama 的自回歸transformer的基礎架構，并使用 AdaLN 進行視覺生成。
實驗
數據集和指標
視覺分詞器和自回歸transformer的訓練均在256×256的ImageNet上進行。對于視覺重建，采用重建-FID（記為rFID）、碼本利用率和LPIPS在ImageNet 50k驗證集上來衡量重建圖像的質量。對于視覺生成，我們通過常用的指標FID、IS和Precision/Recall來衡量圖像生成的質量。
實驗細節
主要結果
視覺重建。 下表1展示了IBQ與常見視覺分詞器的定量重建比較?？梢钥吹?，當碼本規模擴大時，現有VQ方法的碼本使用率顯著下降（例如，VQGAN 在1024碼本規模下的使用率為44%，而在16,384碼本規模下的使用率為5.9%），以及代碼維度（例如，LlamaGen 在8維代碼下的使用率為97%，而在256維代碼下的使用率為0.29%）。因此，實際的表示能力受到碼本崩潰的限制。
相比之下，對所有碼本embedding和視覺編碼器的聯合優化確保了它們之間的一致分布，有助于穩定訓練具有高利用率的大規模碼本和embedding視覺分詞器。具體來說，IBQ在16,384碼本規模和256代碼維度下實現了1.37的rFID，超過了在相同下采樣率和碼本規模下的其他VQ方法。通過將碼本規模增加到262,144，IBQ超越了Open-MAGVIT2，實現了最先進的重建性能（1.00 rFID）。我們還在下圖4中與幾種具有代表性的VQ方法進行了定性比較。IBQ在復雜場景如面部和字符上表現出更好的視覺合理性。
視覺生成。 在下表7中，我們將IBQ與其他生成模型進行比較，包括擴散模型、AR模型以及AR模型的變體（VAR和MAR）在類別條件圖像生成任務上的表現。借助IBQ強大的視覺分詞器，我們的模型在擴大模型規模時（從300M到2.1B）表現出持續的改進，并在不同規模的模型下超越了所有之前的基礎自回歸模型。此外，IBQ優于基于擴散的模型DiT，并在AR模型變體中取得了可比的結果。這些AR模型變體專注于第二階段transformer的架構設計，而我們的工作則致力于第一階段更好的視覺分詞器。因此，我們相信，借助我們更強大的分詞器，AR模型及其變體可以進一步提升。
擴大 IBQ
現有的 VQ 方法在擴展時因碼本崩潰而遇到困難。例如，當將 LlamaGen的代碼維度從 8 增加到 256 時，其使用率和 rFID 顯著下降（97% → 0.29%，2.19 rFID → 9.21 rFID），如上表 1 所示。這是由于訓練期間的部分更新逐漸擴大了未激活代碼與編碼特征之間的分布差距。IBQ 在三個方面顯示出有希望的擴展能力：
- 碼本大小：如下表 4 所示，隨著碼本大小從 1024 擴大到 16,384，重建質量顯著提高。此外，IBQ 即使在使用 262,144 個代碼進行訓練時，也能實現高碼本利用率和視覺效果的一致提升。
- 模型大小：下表 6 顯示，通過在編碼器和解碼器中擴展 ResBlock 的數量，可以保證重建性能的提升。
- 代碼維度：有趣的是，觀察到在擴展代碼維度時，碼本使用率顯著增加。我們假設低維代碼辨別力較弱，類似的代碼往往會聚集在一起。這表明在我們的全局更新策略下，具有代表性的代碼更有可能被選擇。相比之下，高維embedding的代碼在表示空間中是高度信息化的，因為它們在表示空間中是相互稀疏的。因此，這些代碼在訓練過程中可以被均勻選擇，從而確保高利用率和更好的性能。通過以上因素，我們實現了一個擁有 262,144 個碼本大小和 256 維度的超大碼本，并且具有高碼本使用率（84%），實現了最先進的重建性能（1.00 rFID）。為了更好地說明擴展特性，我們還在下圖 5 中提供了可視化。
消融實驗
關鍵設計。 為了驗證我們方法的有效性，對幾個關鍵設計進行了消融研究，如下表2所示。重新實現的VQGAN性能為3.98 rFID，碼本利用率為5.3%。與之前的方法不同，將VQ替換為IBQ后，通過使所有代碼可微分，實現了編碼特征與整個碼本之間的一致分布，從而顯著提高了碼本的使用率（從5.3%提高到98%）和重建質量（從3.98 rFID提高到1.67 rFID）。通過引入雙重量化損失來迫使選擇的代碼embedding和編碼視覺特征相互靠近，IBQ保證了更精確的量化。按照MAGVIT-v2 的做法，我們擴大了模型規模以提高緊湊性，重建性能也相應得到了改善。
與LFQ的比較。 為了進行公平的比較，采用了具有16,384個代碼的LFQ，并用我們的基礎Transformer架構替換了其不對稱的token分解。我們在下表5中比較了LFQ在重建和生成方面的表現，我們提出的IBQ表現更好，這表明增加代碼維度可以提高視覺tokenizer的重建能力，并進一步提升視覺生成。
結論
在本文中，我們識別出了當前向量量化（VQ）方法中部分更新策略導致的tokenizer擴展瓶頸，這種策略逐漸加大了編碼特征與未激活代碼之間的分布差距，最終導致碼本崩潰。為了解決這一挑戰，提出了一種簡單而有效的向量量化方法，稱為索引反向傳播量化（IBQ），用于可擴展的tokenizer訓練。該方法通過在視覺特征與所有碼本embedding之間的分類分布上應用直通估計器來更新所有代碼，從而保持整個碼本與編碼特征之間的一致分布。ImageNet上的實驗表明，IBQ實現了高利用率的大規模視覺tokenizer，在重建（1.00 rFID）和生成（2.05 gFID）方面的性能有所提高，驗證了我們方法的可擴展性和有效性。
本文轉自AI生成未來 ，作者：AI生成未來
