編輯：杜偉、陳萍
基於殘差量化的自迴歸圖像生成，官方已將程式碼公開。
向量量化（Vector quantization，VQ）已經成為自迴歸（AR）模型生成高解析度圖像的一種基本技術，具體來說，該技術將圖像特徵圖通過 VQ 量化後，再進行光柵掃描等排序，之後將圖像表示為離散編碼序列。量化後，訓練的 AR 模型對序列中的編碼進行序列預測。也就是說，AR 模型可以生成高解析度的圖像，而無需預測圖像中的全部像素。
我們假設減少編碼的序列長度對於圖像的 AR 建模很重要。短編碼序列可以顯著降低 AR 模型的計算成本，因為 AR 通常使用先前位置的編碼來預測下一個編碼。然而，以前的研究由於圖像的序列長度問題在速率 – 失真（rate-distortion）不能很好的權衡。也就是說，VQ-VAE 需要呈指數增長的編碼簿 (Codebook)大小來降低量化特徵圖的解析度，同時保持重建圖像的質量。然而，巨大的編碼簿會導致模型參數增加和編碼簿崩潰，使得 VQ-VAE 的訓練不穩定。
來自 Kakao Brain 、韓國浦項科技大學的研究者提出了一種殘差量化 VAE (RQ-VAE)方法，它使用殘差量化 (RQ) 來精確逼近特徵圖並降低其空間解析度。RQ 沒有增加編碼簿大小，而是使用固定大小的編碼簿以從粗到細的方式遞迴量化特徵圖。在 RQ 的 D 次迭代之後，特徵圖表示為 D 個離散編碼的堆疊圖。由於 RQ 可以組成與編碼簿大小一樣多的向量，因此 RQ-VAE 可以精確地逼近特徵圖，同時保留編碼圖像的資訊，而無需龐大的編碼簿。由於精確的近似，RQ-VAE 可以比以前的研究 [14,37,45] 進一步降低量化特徵圖的空間解析度。例如， RQ-VAE 可以使用 8×8 解析度的特徵圖對 256×256 圖像進行 AR 建模。該論文已被 CVPR’22 接收。
論文地址：https://arxiv.org/pdf/2203.01941.pdf
此外，該研究還提出了 RQ-Transformer 來預測 RQ-VAE 提取的編碼。對於 RQ-Transformer 的輸入，該研究首先將 RQ-VAE 中的量化特徵對映轉換為特徵向量序列；然後，RQ-Transformer 預測下一個 D 編碼，以估計下一個位置的特徵向量。由於 RQ-VAE 降低了特徵圖的解析度，RQ-Transformer 可以顯著降低計算成本並輕鬆學習輸入的遠端互動。該研究還為 RQ-Transformer 提出了兩種訓練技術，軟標籤（soft labeling）和用於 RQ-VAE 編碼的隨機取樣。通過解決 AR 模型訓練中的曝光偏差（exposure bias）進一步提高了 RQ-Transformer 的性能。
值得一提的是，該研究近日發佈了在 30M 文字圖像對上訓練的 3.9B 參數的 RQ-Transformer 。據了解，這是公共可用模型中最大的文字到圖像 (T2I) 模型。不過這一結果沒有出現在該論文中。具體內容可參考 GitHub。
程式碼地址：https://github.com/kakaobrain/rq-vae-transformer
如下圖所示，該模型可以生成高質量的圖像。
論文中生成圖像的例子
3.9B 參數的 RQ-Transformer 生成結果，畫框裡帶著眼鏡的貓：
生成沙漠上的埃菲爾鐵塔：
方法概述
研究者提出了用於圖像 AR 建模的 RQ-VAE 和 RQ-Transformer 兩階段框架，如下圖 2 所示。RQ-VAE 使用編碼簿將圖像表示為 D 個離散碼的堆疊圖。然後，使用 RQ-Transformer 自迴歸預測下一個空間位置的下一個 D 碼。他們還解釋了使用 RQ-Transformer 解決 AR 模型訓練中的曝光偏差問題。
階段 1：殘差量化 VAE
研究者首先介紹 VQ 和 VQVAE 的表達方式，然後提出了 RQ-VAE，它可以在不增加編碼簿大小的情況下精確地逼近特徵圖。他們解釋瞭如何將圖像表示為離散碼的堆疊圖。
VQ 和 VQ-VAE 的表達。令編碼簿C為一個有限集
，它包含了成對的程式碼k和程式碼嵌入
，其中K是編碼簿大小，n_z是程式碼嵌入的維數。
考慮到一個向量 z ϵ R^nz，
表示 z 的 VQ，這個程式碼的嵌入離 z 最近，如下公式（1）所示。
在將圖像編碼為離散碼圖後，VQ-VAE 從編碼碼圖重建原始圖像。
對於殘差量化，研究者沒有增加編碼簿大小，而是採用殘差量化（RQ）來離散化向量 z。給定一個量化深度 D，RQ 將 z 表示為一個有序的 D 碼，如下公式（3）所示。
RQ-VAE。在上圖 2 中，研究者提出 RQ-VAE 來精確量化圖像的特徵圖。RQ-VAE 也是由 VQ-VAE 的編解碼器架構組成，不過 VQ 模組被上面的 RQ 模組所取代。
具體而言，深度為 D 的 RQ-VAE 將特徵圖 Z 表示為程式碼 M ϵ [K]^H×W×D 的堆疊圖，並提取了
，成為 d ϵ [D]的深度為 D 的量化特徵圖，得到如下公式（5）。
RQ-VAE 使 AR 模型以低計算成本有效地生成高解析度圖像。對於固定的下采樣因子 f，RQ-VAE 可以產生比 VQ-VAE 更逼真的重建結果，這是因為 RQ-VAE 使用給定的編碼簿大小精確地逼近特徵圖。
階段 2：RQ-Transformer
研究者展示了 RQ-Transformer 如何自迴歸地預測 RQ-VAE 的程式碼堆疊。在對 RQVAE 提取的程式碼的進行 AR 建模之後，他們介紹 RQ-Transformer 如何有效地學習離散程式碼的堆疊圖。此外，研究者還解釋了 RQ-Transformer 的訓練技術，以防止 AR 模型訓練中出現曝光偏差。
深度為 D 的程式碼的 AR 建模。在 RQ-VAE 提取程式碼對映 M ϵ ^[K] H×W×D 後，光柵掃描順序（raster scan order）將 M 的空間索引重新排列為程式碼 S ϵ [ K]^T ×D 的二維陣列，其中 T = HW。也就是說，S_t 是 S 的第 t 行，幷包含了 D 個程式碼，如下公式（8）所示。
如上圖 2 所示，RQ-Transformer 由空間 transformer 和深度 transformer 組成。空間 transformer 是帶掩碼自注意力塊的堆疊，用於提取一個總結先前位置資訊的上下文向量。給定上下文向量 h_t，深度 transformer 自迴歸預測位置 t 處的 D 個程式碼 (S_t1,····,S_tD)。
軟標籤和隨機抽樣。研究者提議對來自 RQ-VAE 的程式碼進行軟標籤和隨機取樣，以解決暴露偏差。
實驗結果
無條件圖像生成
該研究在 LSUN-{cat, bedroom, church}和 FFHQ 資料集上評估了無條件圖像生成的質量。表 1 顯示，RQ-Transformer 模型在無條件生成圖像方面優於其他 AR 模型。對於小規模資料集，如 LSUN-church 和 FFHQ，RQ-Transformer 小幅度優於 DCT 和 VQ-GAN。對於更大規模的資料集，如 LSUN-{cat, bedroom}，RQ-Transformer 明顯優於其他 AR 模型和基於擴散的模型。
該研究推測性能的提高來自於 RQ-VAE 較短的序列長度，因為 SQ-Transformer 可以很容易地在較短的序列長度內學習編碼間的長程互動。圖 3 前兩行展示了 RQ-Transformer 可以無條件生成高質量圖像。
有條件圖像生成
該研究分別使用 ImageNet 和 CC-3M 進行類和文字條件的圖像生成。如表 2 所示，RQ-Transformer 在 ImageNet 上的性能明顯優於以前的模型。RQ-Transformer（480M 參數）與以往 AR 模型具有競爭力，包括 VQ-VAE2， DCT，和 VQ-GAN，雖然 RQ-Transformer 比 VQ-GAN 少 3 倍的參數。此外，具有 821M 參數的 RQ-Transformer 在沒有拒絕取樣的情況下優於以往的 AR 模型。
具有 1.4B 參數的 RQ-Transformer 在沒有拒絕取樣的情況下達到 11.56 的 FID 分數。當將 RQ-VAE 的訓練 epoch 從 10 增加到 50 時，1.4B 參數的 RQ-Transformer 進一步提高了性能，達到了 8.71 FID 分數 。此外，當研究者進一步將參數數量增加到 3.8B 時，RQ-Transformer 在沒有拒絕取樣的情況下達到了 7.55 的 FID 分數，並且可以與 BigGAN 競爭。
RQ-Transformer 還可以根據 CC-3M 的各種文字條件生成高質量圖像。RQ-Transformer 在參數數量相似的情況下明顯高於 VQ-GAN 的性能。圖 3 顯示在 CC-3M 上訓練的 RQ-Transformer 可以使用各種文字條件生成高質量的圖像。
RQ-Transformer 的計算效率
圖 4 評估了 RQ-Transformer 的取樣速度，並與 VQ-GAN 進行了比較，兩種模型參數均為 1.4B。VQ-GAN 和 RQ-Transformer 的輸入分別設置為 16×16×1 和 8×8×4。每個模型使用單塊 NVIDIA A100 GPU 生成 5000 個樣品，批量大小分別為 100、200 和 500。
對於 100 和 200 的批量大小，與 VQ-GAN 相比，RQ-Transformer 顯示出 4.1 倍和 5.6 倍的加速。此外，由於 RQ-VAE 的短序列長度節省了記憶體，RQ-Transformer 可以將批量大小增加到 500，而 VQ-GAN 是不允許的。因此，RQ-Transformer 可以進一步加快取樣速度，每張圖像為 0.02 秒，比批量大小為 200 的 VQ-GAN 快 7.3 倍。因此，RQ-Transformer 比以前的 AR 模型在計算上更高效，同時實現高解析度圖像生成基準的最新結果。
