【論文解説】HQ-VAE: 階層的な離散表現学習における確率的量子化(SQ)の導入
本記事では、以下の論文
で提案された HQ-VAE の概要と主要な成果について解説します。
先行研究との対応が分かりやすいように、自作の図を交えながら一貫した変数表記で説明します。
要約
本論文では、階層的な離散潜在表現を学習する新たな変分オートエンコーダ (Variational Autoencoder; VAE) フレームワークである HQ-VAE が提案されている。
従来広く用いられてきた VQ-VAE 系列(VQ-VAE-2 や RQ-VAE など)は、離散コードブックを使って入力データを量子化し、高次元のデータ(画像や音声など)の圧縮や生成タスクに応用されてきた。しかし、これらの手法には
- コードブック崩壊(コードブック要素が十分に活用されず、上位層がうまく機能しなくなる)、
- 重み付きハイパーパラメータ(
など)のチューニングやヒューリスティックな対処(Stop-Gradient, EMA 更新, コードブックリセットなど)が必要、𝛽
といった問題があった。
これに対し、本論文では「変分ベイズ」の枠組みに基づく階層的離散潜在表現学習を行う HQ-VAE を提案する。HQ-VAE は確率的量子化(Stochastic Quantization; SQ)を用いており、学習中に量子化過程が自己アニーリング的に「連続的→離散的」へ収束していくため、コードブックを効率的に使いながら安定した学習が可能になる。また、これにより従来必要とされてきた多数のハイパーパラメータ調整や EMA 更新といったヒューリスティックが不要になる利点が示されている。
さらに、従来の VQ-VAE-2、RQ-VAE などを HQ-VAE の特別な場合として位置付けられることから、既存の階層的離散表現手法を変分ベイズの観点で統一的に理解・拡張できる、といった貢献も強調されている。
1. 導入
1.1 VAEとVQ-VAE
VAE (Variational AutoEncoder)
- 入力データを低次元の潜在空間へ圧縮し、そこから再構成を行う生成モデル。
- 潜在変数を確率分布として扱いながら、エンコーダとデコーダを同時に学習する。
- 高い柔軟性を持つ一方、再構成がぼやけやすいという欠点がある。
- 学習は、以下の変分下界(ELBO)の最大化により行う:
VQ-VAE (Vector Quantized VAE)
- 潜在空間を離散化し、コードブック上の埋め込みベクトルを使って表現することで、シャープな画像生成や離散表現の学習が可能に。
- 「入力特徴を最も近い埋め込みベクトルへ量子化→再構成」する仕組み。
- 典型的な損失関数は以下のように書かれる(
は勾配停止演算の省略記号):\mathrm{sg}[\cdot]
- 一方、学習が進むと「コードブックの一部しか使われなくなる(コードブック崩壊)」が生じやすいのが難点。
1.2 階層版VQ-VAE
VQ-VAE-2
VQ-VAE-2は、VQ-VAEに階層的な離散潜在空間を導入したモデル。従来のVQ-VAEが単一のコードブックのみを扱うのに対し、複数レベル(上位層・下位層)のコードブックを用いて、異なる解像度・粒度の情報を表現する。
-
トップレベル(上位層)
- 画像全体の大まかな構造を量子化し、グローバルな色や形を表現。
-
ボトムレベル(下位層)
- トップレベルではカバーしきれない細部の情報を量子化し、テクスチャなど局所的特徴を表現。
トップダウンとボトムアップ
-
ボトムアップのエンコーダ
入力画像から下位層の特徴マップを得て、それをさらにダウンサンプリングして上位層の特徴マップを導出。両層で量子化を実施。 -
トップダウンのデコーダ
上位層のコードから大まかな構造を復元し、下位層のコードを追加して細部を再現。
学習方法
- VQ-VAEで用いられる勾配停止やEMA更新を階層的に適用。
- しかし階層が増えることで、特定の層だけがほぼ使われないなどのコードブック崩壊が起こりやすい点が指摘されている。
特徴・メリット
- 高解像度生成: 層を増やすほど細部を捉えやすくなるため、単層VQ-VAEより高精細な画像を生成しやすい。
- グローバルとローカルの分離: 上位層で大きな構造、下位層で細部を学習するため、解釈性や編集応用がしやすい。
RQ-VAE (Residual-Quantized VAE)
RQ-VAEは「残差を再量子化する」手法で階層化を実現。VQ-VAE-2のように解像度で階層を分けるのではなく、「粗→細」と段階的に情報を補完する。
- 1層目の量子化: 入力特徴マップをコードブックで量子化し、大まかな近似を得る。
- 残差計算: 「元の特徴マップ − 量子化後の近似」を取り、取りこぼした情報を抽出。
- 2層目量子化: その残差を別コードブックで量子化し、より細かい補正を与える。
- 必要に応じて繰り返し: 粗い表現から徐々に残差を埋めることで、多層的に情報をカバー。
コードブックの「実質的サイズ」を拡大
- 1層あたりのコードブックサイズが
でも、K 層積めば最大でL 通りの表現組合せが可能。K^L - 段階的に学習できるため、大規模コードブックを直接学習するより崩壊を緩和しやすい面もあるが、やはり崩壊は起こりうる。
特徴・メリット
- 柔軟性: 残差を少しずつ埋めていく構造のため、拡張が容易。
- 層追加による表現力向上: 追加の層でさらなるビットレート・解像度を実現できる。
VQ-VAE-2とRQ-VAEの比較
-
階層の設計
- VQ-VAE-2: グローバル→局所という形で解像度別に階層を分ける。
- RQ-VAE: 前層の残差を次層で量子化し、粗→細へと精緻化。
-
コードブックの解釈
- VQ-VAE-2: 上位層が大まかな形状、下位層がテクスチャなど、役割が比較的明確に分かれる。
- RQ-VAE: 残差を補っていくため、層ごとの意味分担は必ずしも明瞭ではないが、積むほど表現力が増す。
1.3 VQのコードブック崩壊と対策
-
VQ-VAEでのコードブック崩壊
学習が進むほどに一部のコードブック要素だけが使われる問題。表現の偏りや再構成精度の低下をもたらす。 -
ヒューリスティクスでの対策
- EMA (Exponential Moving Average) による埋め込みベクトル更新
- 使用率が低いコードブック要素のリセット
- ハイパーパラメータ調整
-
階層版VQ-VAEでの崩壊
- 複数階層が絡むため、「上位層が全く使われない」「一部の要素だけが集中して使われる」などの形で崩壊が発生しやすい。
1.4 VQからSQ(確率的量子化)へ
従来の決定論的量子化では「最も近い埋め込みベクトル」を一意に選択していた:
一方、SQ (Stochastic Quantization) では確率的に選択することで、表現選択の極端な偏りを回避しやすくする:
このSQを変分ベイズの枠組みを用いてVQ-VAEに自然に導入したのが、次に紹介するSQ-VAEである。
SQ-VAE (Stochastically Quantized VAE)
- VQ-VAEをSQ化することで、コードブック崩壊を緩和。
- 確率的な量子化を徐々に決定論的に収束させる「セルフアニーリング効果」を利用。
- 従来のヒューリスティクスをほぼ不要にし、ハイパーパラメータも1種類に削減。
- 学習は以下に示すELBOの最大化により行われる。
SQ-VAEのELBO一般式
SQ‐VAE の ELBO は以下の形で与えられる。
変数:
-
: 入力データ\bm{x} -
: 離散潜在変数(量子化されたコード)\bm{Z} -
: 連続潜在変数(量子化前の連続表現)\tilde{\bm{Z}} -
: エントロピーH(\cdot)
確率分布:
-
はエンコーダ (q_{(\bm{\phi},\bm{\varphi})}(\tilde{\bm{Z}}|\bm{x}) )\bm{x} \to \tilde{\bm{Z}} -
は脱量子化過程に対応する分布 (p_{\bm{\varphi}}(\tilde{\bm{Z}}|\bm{Z}) )\bm{Z} \to \tilde{\bm{Z}} -
は事前分布P(\bm{Z}) -
は量子化過程に対応する分布(確率質量)(\hat{P}_{\bm{\varphi}}(\bm{Z}|\tilde{\bm{Z}}) )\tilde{\bm{Z}} \to \bm{Z} -
はデコーダ (p_{\bm{\theta}}(\bm{x}|\bm{Z}) )\bm{Z} \to \bm{x}
パラメータ:
-
: デコーダのパラメータ\bm{\theta} -
: エンコーダのメインパラメータ\bm{\phi} -
: 量子化/脱量子化過程のパラメータ\bm{\varphi} -
: エンコーダのパラメータセット全体(\bm{\phi}, \bm{\varphi})
なおこの導出で、一様事前分布(
を利用した。
生成過程に等方性正規分布を仮定する
生成過程の尤度を以下のように仮定する。
ここで、
-
は潜在変数\mu_{\bm{\theta}}(\bm{Z}) から計算される平均ベクトル\bm{Z} -
はスカラーであり、学習可能な分散パラメータである。\sigma^2 -
はデコーダのパラメータ\bm{\theta}
その対数尤度は次のようになる。
ここで
脱量子化過程に(等方性)正規分布を仮定する
SQ-VAEのELBO一般式のうち、
この部分について考える。
従来どおり、脱量子化は正規分布と仮定する。すなわち、
その対数は、
同様に、エンコードも内部で同じ脱量子化過程を含む。つまり、
よって、
次に
残った二次形式の項のうち、2番目に導出した方の期待値は、
となり消える。(
なお、
そのため、この項の期待値はカイ二乗分布の性質から自由度に等しくなり、結果として
以上より、最終的に期待値の中で残るのは、
という項だけである。
Gaussian Ⅰ SQ-VAEの場合
脱量子化過程の正規分布にさらに等方性共分散行列の仮定を導入すると、
となる。ここで
要するに
ここで、
-
はエンコーダによる潜在変数(量子化前の連続的な潜在変数)\tilde{\bm{Z}} -
は量子化後の潜在変数(離散表現)\bm{Z} -
は脱量子化分布の分散(スカラー)、s^2 の一部\bm{\varphi}
以上の簡略化により、項が直観的に潜在変数の量子化前後の距離(二乗誤差)に比例した形で表されることが明らかとなる。
1.5 階層版VQ-VAEもSQ化したい → HQ-VAEの提案
2. 提案手法: HQ-VAE
HQ-VAEは、階層的に離散潜在変数を扱う変分ベイズモデルで、各階層に確率的量子化(SQ)を導入し、コードブック崩壊を抑制する。以下、各構成要素を順を追って見ていく。
2.1 階層的離散表現
-
階層構造
潜在変数を とし、各階層\bm{Z}_{1:L} にコードブックl (埋め込み数\mathbf{B}_l 、次元K_l )を設定。d_l -
事前分布と生成
各階層の潜在変数は一様事前分布 を仮定し、生成モデルは等方性正規分布P(\boldsymbol{z}_{l,i} = \boldsymbol{b}_k) = 1/K_l p_{\theta}(\boldsymbol{x} \mid \bm{Z}_{1:L}) = \mathcal{N}\bigl(\boldsymbol{x};\, \mathbf{f}_{\theta}(\bm{Z}_{1:L}),\, \sigma^{2}\mathbf{I}\bigr). -
変分推論
真の事後分布 は計算困難なので、近似事後を導入し、ELBOの最大化により学習を行う。p_{\theta}(\mathbf{Z}_{1:L}\mid x)
2.2 トップダウン層の構成
HQ-VAEでは、トップダウン方向の情報伝達を行う層として下記の3種類を定義している。
-
Firstトップダウン層
- 全HQ-VAEインスタンスに必ず存在する最上位層。
- これのみの構成は単層のSQ-VAEに相当。
-
Injected型トップダウン層
- ボトムアップパスの高解像度特徴を「注入」し、VQ-VAE-2的な「グローバル情報→局所情報」の構造を実現。
-
Residual型トップダウン層
- 一つ前の階層で量子化した残差をさらに量子化して補完する。RQ-VAEに相当。
これらを自在に組み合わせることで、任意の階層設計を行える。
2.3 DequantizationとQuantization
各階層
と表される。
補助変数
3. HQ-VAEのインスタンスと実験結果
トップダウン層の構成によって、大きく2種類のHQ-VAEインスタンスが示されている。
3.1 SQ-VAE-2 (Injected Top-Down Only)
-
概要
Injected型トップダウン層のみを用いた、VQ-VAE-2のSQ版に相当するモデル。 -
損失関数
以下の を最小化することで学習を行う:-\text{ELBO} \begin{aligned} &\mathcal{J}_{\text{SQ-VAE-2}}(\bm{x};\,\bm{\theta},\bm{\phi},s^{2},\mathcal{B}) \\ &= \frac{D}{2}\log\sigma^{2} + \mathbb{E}_{\mathcal{Q}(\bm{Z}_{1:L}, \hat{\bm{Z}}_{1:L}\mid \bm{x})} \Biggl[ \frac{\|\bm{x}-\bm{f}_{\bm{\theta}}(\bm{Z}_{1:L})\|_{2}^{2}}{2\sigma^{2}} + \sum_{l=1}^{L} \Bigl( \frac{\|\hat{\bm{Z}}_{l}-\bm{Z}_{l}\|_{F}^{2}}{2s_{l}^{2}} - H\bigl(\hat{P}_{s_{l}^{2}}(\bm{Z}_{l}\mid \hat{\bm{Z}}_{l})\bigr) \Bigr) \Biggr]. \end{aligned} -
実験結果
再構成性能はVQ-VAE-2を大きく上回り、特に下位層のコードブック使用率が高まることが報告されている。
Dataset | Model | RMSE ↓ | LPIPS ↓ | SSIM ↑ | exp(H(Q(Z₁))) | exp(H(Q(Z₂))) | ... |
---|---|---|---|---|---|---|---|
ImageNet | VQ-VAE-2 | 6.071 | 0.265 | 0.751 | 106.8 | 288.8 | |
SQ-VAE-2 | 4.603 | 0.096 | 0.855 | 406.2 | 355.5 | ||
FFHQ | VQ-VAE-2 | 4.866 | 0.323 | 0.814 | 24.6 | 41.3 | ... |
SQ-VAE-2 | 2.118 | 0.166 | 0.909 | 125.8 | 398.7 | ... |
3.2 RSQ-VAE (Residual Top-Down Only)
-
概要
Residual型トップダウン層のみで構成した、RQ-VAEのSQ版。 -
特徴
- RQ-VAEより高い再構成性能を達成。
- 上位層のコードブック使用率が特に向上することが確認。
-
損失関数
事前分布の定義を微修正し、各階層の量子化残差に対応する補助変数を導入する形でELBOを定式化(詳細は論文参照)。 -
実験結果
RQ-VAEを上回る再構成精度を示し、上位層が活発に利用される。
3.3 階層的離散表現の可視化
- 再構成画像を階層ごとに比較すると、上位層は大まかな構造、下位層はより微細な情報を担当している様子が見て取れる。
4. 付録: Progressive Coding の概念
論文では階層構造を用いて、漸進的に画像を高精細化するprogressive codingのアイデアも提示されている。具体的には、各階層
のような補助的な生成を行い、低次元情報のみで粗い再構成、上位層を追加するたびに精細化する仕組みを導入。ノイズを加えることで段階的な圧縮と再構成を実証している(詳細は論文参照)。
まとめ
- HQ-VAE は階層型VQ-VAE(VQ-VAE-2やRQ-VAE)を確率的量子化(SQ)で拡張し、コードブック崩壊を効果的に防ぎながら再構成性能を向上させる。
- SQ-VAE-2(Injected型)と RSQ-VAE(Residual型)の2例が示され、どちらも従来手法を上回る性能とコードブック使用率を実現。
- progressive coding により、漸進的な画質向上が可能になり、生成モデルとしての柔軟性も拡がる。
以上、HQ-VAEの概要と主要な成果についての解説でした。階層的離散表現学習と確率的量子化の組み合わせが生み出す新たな可能性を、ぜひ論文本体でも確認してみてください。
Discussion