論文鏈接：https://arxiv.org/pdf/2412.02692github鏈接：https://github.com/TencentARC/SEED-Voken

亮點直擊

• 提出了一種簡單而有效的向量量化方法，稱為索引反向傳播量化（Index Backpropagation Quantization，IBQ），用于訓(xùn)練可擴展的視覺分詞器。
• 通過增加碼本大小、編碼維度和模型規(guī)模來研究IBQ的擴展特性。IBQ首次訓(xùn)練了一個超大碼本（），具有大維度（256）和高使用率，實現(xiàn)了最先進的重建性能。
• 展示了一系列從300M到2.1B的基礎(chǔ)自回歸圖像生成模型，顯著超越了競爭方法，例如LlamaGen和 Open-MAGVIT2。

總結(jié)速覽

解決的問題

現(xiàn)有的向量量化（VQ）方法在可擴展性方面存在困難，主要由于訓(xùn)練過程中僅部分更新的碼本的不穩(wěn)定性。隨著利用率的降低，碼本容易崩潰，因為未激活代碼與視覺特征之間的分布差距逐漸擴大。

提出的方案

提出了一種新的向量量化方法，稱為索引反向傳播量化（Index Backpropagation Quantization，IBQ），用于碼本embedding和視覺編碼器的聯(lián)合優(yōu)化。通過在編碼特征與碼本之間的單熱編碼分類分布上應(yīng)用直通估計器，確保所有代碼都是可微的，并與視覺編碼器保持一致的潛空間。

應(yīng)用的技術(shù)

• 使用直通估計器在單熱編碼分類分布上進行優(yōu)化，使得所有代碼可微。
• 通過IBQ實現(xiàn)碼本embedding和視覺編碼器的聯(lián)合優(yōu)化。
• 研究了IBQ在增加碼本大小、編碼維度和模型規(guī)模方面的擴展特性。

達到的效果

• 在標(biāo)準ImageNet基準上的實驗表明，IBQ在重建（1.00 rFID）和自回歸視覺生成方面取得了具有競爭力的結(jié)果。
• 展示了一系列從300M到2.1B的基礎(chǔ)自回歸圖像生成模型，顯著超越了競爭方法，如LlamaGen和Open-MAGVIT2。

效果展示

下圖的上半部分展示了在1024×1024分辨率下，IBQ分詞器在Unsplash數(shù)據(jù)集上的測試結(jié)果。下半部分則展示了IBQ分詞器在256×256分辨率下，針對Imagenet數(shù)據(jù)集的測試結(jié)果。(a)表示原始圖像，(b)表示重建圖像。

Imagenet上256×256類條件生成樣本效果：

方法

我們的框架由兩個階段組成。第一階段是通過索引反向傳播量化學(xué)習(xí)一個具有高碼本利用率的可擴展視覺分詞器。在第二階段，我們使用自回歸變換器通過下一個標(biāo)記預(yù)測進行視覺生成。

Preliminary

這些方法采用的部分更新策略（即僅優(yōu)化選定的代碼）逐漸擴大了視覺特征與未激活代碼之間的分布差距。這會導(dǎo)致訓(xùn)練期間的不穩(wěn)定性，因為碼本崩潰會阻礙視覺分詞器的可擴展性。

索引反向傳播量化

量化。 為了確保在訓(xùn)練過程中碼本與編碼特征之間的一致分布，我們引入了一種全代碼更新方法，即索引反向傳播量化（Index Backpropagation Quantization, IBQ）。IBQ的核心在于將梯度反向傳播到碼本的所有代碼，而不僅僅是選定的代碼。算法1提供了IBQ的偽代碼。

具體來說，我們首先對給定的視覺特征與所有代碼embedding進行點積運算作為logits，并通過softmax函數(shù)獲得概率（soft one-hot）。

然后我們將soft one-hot 類別分布的梯度復(fù)制到hard one-hot索引上：

給定索引，量化后的特征通過以下方式獲得：

通過這種方式，我們可以通過索引將梯度傳遞到碼本的所有代碼上。通過索引反向傳播量化，整個碼本和編碼特征的分布在整個訓(xùn)練過程中保持一致，從而獲得較高的碼本利用率。

訓(xùn)練損失。 

與 VQGAN類似，分詞器的優(yōu)化由多種損失的組合完成：

如下圖 3 所示，現(xiàn)有的 VQ 方法在每次反向過程中僅優(yōu)化有限數(shù)量的代碼以接近編碼特征。這逐漸擴大了未激活代碼和編碼特征之間的分布差距，最終導(dǎo)致碼本崩潰。隨著代碼維度和碼本大小的增加，這種情況變得更加嚴重。我們不是直接將直通估計器 [1] 應(yīng)用于選定的代碼，而是將這種參數(shù)化方法應(yīng)用于視覺特征和所有碼本embedding之間的分類分布，以使梯度能夠反向傳播到所有代碼。通過這種方式，整個碼本和編碼特征之間的分布在整個訓(xùn)練過程中保持一致。因此，IBQ 實現(xiàn)了具有高代碼維度和利用率的極大碼本。

自回歸Transformer

其中，c 是條件，例如類別標(biāo)簽。請注意，由于我們專注于視覺分詞器，因此我們采用類似于 Llama 的自回歸transformer的基礎(chǔ)架構(gòu)，并使用 AdaLN 進行視覺生成。

實驗

數(shù)據(jù)集和指標(biāo)

視覺分詞器和自回歸transformer的訓(xùn)練均在256×256的ImageNet上進行。對于視覺重建，采用重建-FID（記為rFID）、碼本利用率和LPIPS在ImageNet 50k驗證集上來衡量重建圖像的質(zhì)量。對于視覺生成，我們通過常用的指標(biāo)FID、IS和Precision/Recall來衡量圖像生成的質(zhì)量。

實驗細節(jié)

主要結(jié)果

視覺重建。 下表1展示了IBQ與常見視覺分詞器的定量重建比較?？梢钥吹?，當(dāng)碼本規(guī)模擴大時，現(xiàn)有VQ方法的碼本使用率顯著下降（例如，VQGAN 在1024碼本規(guī)模下的使用率為44%，而在16,384碼本規(guī)模下的使用率為5.9%），以及代碼維度（例如，LlamaGen 在8維代碼下的使用率為97%，而在256維代碼下的使用率為0.29%）。因此，實際的表示能力受到碼本崩潰的限制。

相比之下，對所有碼本embedding和視覺編碼器的聯(lián)合優(yōu)化確保了它們之間的一致分布，有助于穩(wěn)定訓(xùn)練具有高利用率的大規(guī)模碼本和embedding視覺分詞器。具體來說，IBQ在16,384碼本規(guī)模和256代碼維度下實現(xiàn)了1.37的rFID，超過了在相同下采樣率和碼本規(guī)模下的其他VQ方法。通過將碼本規(guī)模增加到262,144，IBQ超越了Open-MAGVIT2，實現(xiàn)了最先進的重建性能（1.00 rFID）。我們還在下圖4中與幾種具有代表性的VQ方法進行了定性比較。IBQ在復(fù)雜場景如面部和字符上表現(xiàn)出更好的視覺合理性。

視覺生成。 在下表7中，我們將IBQ與其他生成模型進行比較，包括擴散模型、AR模型以及AR模型的變體（VAR和MAR）在類別條件圖像生成任務(wù)上的表現(xiàn)。借助IBQ強大的視覺分詞器，我們的模型在擴大模型規(guī)模時（從300M到2.1B）表現(xiàn)出持續(xù)的改進，并在不同規(guī)模的模型下超越了所有之前的基礎(chǔ)自回歸模型。此外，IBQ優(yōu)于基于擴散的模型DiT，并在AR模型變體中取得了可比的結(jié)果。這些AR模型變體專注于第二階段transformer的架構(gòu)設(shè)計，而我們的工作則致力于第一階段更好的視覺分詞器。因此，我們相信，借助我們更強大的分詞器，AR模型及其變體可以進一步提升。

擴大 IBQ

現(xiàn)有的 VQ 方法在擴展時因碼本崩潰而遇到困難。例如，當(dāng)將 LlamaGen的代碼維度從 8 增加到 256 時，其使用率和 rFID 顯著下降（97% → 0.29%，2.19 rFID → 9.21 rFID），如上表 1 所示。這是由于訓(xùn)練期間的部分更新逐漸擴大了未激活代碼與編碼特征之間的分布差距。IBQ 在三個方面顯示出有希望的擴展能力：

• 碼本大小：如下表 4 所示，隨著碼本大小從 1024 擴大到 16,384，重建質(zhì)量顯著提高。此外，IBQ 即使在使用 262,144 個代碼進行訓(xùn)練時，也能實現(xiàn)高碼本利用率和視覺效果的一致提升。

• 模型大小：下表 6 顯示，通過在編碼器和解碼器中擴展 ResBlock 的數(shù)量，可以保證重建性能的提升。

• 代碼維度：有趣的是，觀察到在擴展代碼維度時，碼本使用率顯著增加。我們假設(shè)低維代碼辨別力較弱，類似的代碼往往會聚集在一起。這表明在我們的全局更新策略下，具有代表性的代碼更有可能被選擇。相比之下，高維embedding的代碼在表示空間中是高度信息化的，因為它們在表示空間中是相互稀疏的。因此，這些代碼在訓(xùn)練過程中可以被均勻選擇，從而確保高利用率和更好的性能。通過以上因素，我們實現(xiàn)了一個擁有 262,144 個碼本大小和 256 維度的超大碼本，并且具有高碼本使用率（84%），實現(xiàn)了最先進的重建性能（1.00 rFID）。為了更好地說明擴展特性，我們還在下圖 5 中提供了可視化。

消融實驗

關(guān)鍵設(shè)計。 為了驗證我們方法的有效性，對幾個關(guān)鍵設(shè)計進行了消融研究，如下表2所示。重新實現(xiàn)的VQGAN性能為3.98 rFID，碼本利用率為5.3%。與之前的方法不同，將VQ替換為IBQ后，通過使所有代碼可微分，實現(xiàn)了編碼特征與整個碼本之間的一致分布，從而顯著提高了碼本的使用率（從5.3%提高到98%）和重建質(zhì)量（從3.98 rFID提高到1.67 rFID）。通過引入雙重量化損失來迫使選擇的代碼embedding和編碼視覺特征相互靠近，IBQ保證了更精確的量化。按照MAGVIT-v2 的做法，我們擴大了模型規(guī)模以提高緊湊性，重建性能也相應(yīng)得到了改善。

與LFQ的比較。 為了進行公平的比較，采用了具有16,384個代碼的LFQ，并用我們的基礎(chǔ)Transformer架構(gòu)替換了其不對稱的token分解。我們在下表5中比較了LFQ在重建和生成方面的表現(xiàn)，我們提出的IBQ表現(xiàn)更好，這表明增加代碼維度可以提高視覺tokenizer的重建能力，并進一步提升視覺生成。

結(jié)論

在本文中，我們識別出了當(dāng)前向量量化（VQ）方法中部分更新策略導(dǎo)致的tokenizer擴展瓶頸，這種策略逐漸加大了編碼特征與未激活代碼之間的分布差距，最終導(dǎo)致碼本崩潰。為了解決這一挑戰(zhàn)，提出了一種簡單而有效的向量量化方法，稱為索引反向傳播量化（IBQ），用于可擴展的tokenizer訓(xùn)練。該方法通過在視覺特征與所有碼本embedding之間的分類分布上應(yīng)用直通估計器來更新所有代碼，從而保持整個碼本與編碼特征之間的一致分布。ImageNet上的實驗表明，IBQ實現(xiàn)了高利用率的大規(guī)模視覺tokenizer，在重建（1.00 rFID）和生成（2.05 gFID）方面的性能有所提高，驗證了我們方法的可擴展性和有效性。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/ABR-RuLHV4G_2zjfWn-1hA??