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

亮點(diǎn)直擊

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

總結(jié)速覽

解決的問(wèn)題

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

提出的方案

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

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

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

達(dá)到的效果

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

效果展示

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

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

方法

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

Preliminary

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

索引反向傳播量化

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

具體來(lái)說(shuō)，我們首先對(duì)給定的視覺(jué)特征與所有代碼embedding進(jìn)行點(diǎn)積運(yùn)算作為logits，并通過(guò)softmax函數(shù)獲得概率（soft one-hot）。

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

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

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

訓(xùn)練損失。 

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

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

自回歸Transformer

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

實(shí)驗(yàn)

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

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

實(shí)驗(yàn)細(xì)節(jié)

主要結(jié)果

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

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

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

擴(kuò)大 IBQ

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

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

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

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

消融實(shí)驗(yàn)

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

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

結(jié)論

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

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

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