今回の論文
今回取り上げるのは、Aaron van den Oord、Oriol Vinyals、Koray Kavukcuogluによる「Neural Discrete Representation Learning」です。研究分野は教師なし表現学習・生成モデルで、NIPS 2017の論文です。会議の公開ページとarXivの論文ページから参照できます。arXiv初版は2017年11月2日で、本記事では2018年5月30日改訂のv2本文を確認しています。
選んだ理由は、「画像や音声を、そのまま予測するには細かすぎる」という問題に対して、予測しやすい単位自体を学習する考え方を示しているためです。生成器を大きくする前に、入力表現と復元処理をどう分担するかを考える材料になります。
どんな技術か
VQ-VAE（Vector Quantised-Variational AutoEncoder）は、画像や音声を少数の離散コードへ変換し、そのコードから入力を復元するニューラルネットワークです。
例えば画像なら、エンコーダが小さな特徴マップを作り、各位置の特徴ベクトルを「代表ベクトルの一覧」であるコードブックの中から最も近いものへ置き換えます。保存するのは各代表ベクトルの番号です。デコーダは番号からベクトルを取り出して画像へ戻します。
重要なのは、この番号が人手で決めた物体名や音素記号ではないことです。復元に役立つよう、エンコーダ、コードブック、デコーダを一緒に学習します。さらに番号の並び方を別のモデルで学べば、新しい番号列を生成して画像や音声へ変換できます。原論文、3節
課題
画像の画素や音声の波形には、近くの値同士の相関や細かな揺らぎが大量にあります。それらを一つずつ予測するモデルは、局所的な細部にも計算と表現能力を使います。長い時間の発話内容や画像全体の構成を扱うには、より短い表現で予測する仕組みが欲しくなります。
連続値の潜在表現を使うVAEにも圧縮の役割がありますが、強力な自己回帰デコーダを組み合わせると、デコーダだけで入力を予測できてしまい、潜在変数を無視することがあります。これが事後崩壊です。見た目の生成品質がよくても、別の処理へ渡せる有用な表現が得られない場合があります。
一方、表現を離散化すると、最も近い番号を選ぶ操作を通常の微分で学習できないという問題が生まれます。VQ-VAEは、有限個のコードを選ぶ処理と、ニューラルネットワークの勾配による学習を両立させようとします。
実際の開発では、音声生成の系列を短くしたい、画像の復元と構成の予測を分けたい、といった場面に対応します。ただし短くするほど細部の情報は落ちるため、圧縮量だけでなく、後段が必要とする情報を残せるかが課題です。原論文、1〜3節
用語解説
- 潜在表現とオートエンコーダ
- 入力をエンコーダで特徴へ変換し、デコーダで入力へ戻す構成です。その中間の特徴が潜在表現です。VQ-VAEでは、この中間部分を離散コードに制限しても復元できるよう学びます。
- ベクトル量子化とコードブック
- 連続値ベクトルを、有限個の代表ベクトルのうち最も近いものへ置き換える操作と、その一覧です。VQ-VAEでは代表ベクトルの値も学び、番号が生成モデルの予測対象になります。
- 事後崩壊（posterior collapse）
- 入力から得た潜在変数をデコーダが無視し、潜在表現に情報が残りにくくなる現象です。本論文は、強力なデコーダと組み合わせても有用な離散表現を使える学習方式を狙います。
- Straight-through estimatorとstop-gradient
- 前者は、離散的な選択を逆伝播時だけ通り抜けるように扱う勾配の近似です。後者は、値は使いつつ勾配の伝達を止める操作です。両者を組み合わせ、エンコーダとコードブックの更新先を分けます。
- 自己回帰の事前分布
- 既に決まったコードから次のコードの確率を予測するモデルです。VQ-VAEの復元器とは別に学び、入力がない状態から新しいコード列を作るために使います。
技術の仕組み
特徴を最も近い代表ベクトルへ置き換えます
コードブックを、K個のD次元ベクトル e_1, ..., e_K とします。各位置のエンコーダ出力 z_e(x) に対して、次の番号とベクトルを選びます。
k = argmin_j ||z_e(x) - e_j||²
z_q(x) = e_k
これは「特徴と代表ベクトルの距離をすべて比べ、最も近い代表を採用する」という意味です。番号kは離散値ですが、デコーダへ入るのは対応するD次元ベクトルです。Kは候補数、Dは一つの候補が持つ特徴の次元で、異なる設定です。
画像では2次元、音声では時間方向、動画では時間と空間の潜在配置を使います。一つの画像全体を必ず一番号にするわけではなく、位置ごとに番号を持ちます。画像実験の一例は、128×128画素を32×32個のコードへ変換する構成です。原論文、3.1・4.2節
順伝播では離散化し、逆伝播では勾配を近似します
番号の選択は、入力の微小な変化では変わらず、境界を越えると急に変わります。そのため、通常の滑らかな関数と同じようには微分できません。
そこで順伝播では実際に最近傍のベクトルを使い、逆伝播ではデコーダ入力に届いた勾配をエンコーダ出力へそのまま渡します。stop-gradientを sg と書くと、概念的には次のように表せます。
z_st = z_e + sg(z_q - z_e)
値としては z_q になり、勾配は z_e 側へ流れます。これは離散選択の厳密な微分ではなく、学習を進めるための近似です。この経路では、再構成損失からコードブックへ直接勾配を流しません。原論文、3.2節
三つの損失で更新先を分担します
学習の目的を、最小化する損失として書くと次の形です。式は一つの潜在位置について示しています。
L = L_rec(x, decoder(z_st)) + ||sg(z_e) - e_k||² + β ||z_e - sg(e_k)||²
第一項は再構成損失で、入力を復元するようエンコーダとデコーダを更新します。確率モデルなら負の対数尤度を使い、画像の実験では画素のMSEも扱います。
第二項は、選ばれたコードブックの代表をエンコーダ出力へ近づけます。エンコーダ側は固定値として扱います。第三項のcommitment lossは逆に、エンコーダ出力を選んだ代表へ近づけ、両者が離れ続けるのを抑えます。こちらでは代表側を固定します。
原論文はβ＝0.25を使っています。ただし再構成損失の尺度との関係があるため、どのデータ・実装でも同じ値が最適とは限りません。複数の潜在位置がある場合は位置ごとの損失を集約します。付録には指数移動平均によるコードブック更新もありますが、本文はこの論文の実験では使っていないと述べています。原論文、3.2節・付録A.1
復元を学んだ後で、コードの並び方を学びます
復元器の学習中は、各潜在変数の事前分布を一様分布として固定します。エンコーダが一つの番号を確定的に選ぶため、このときのKL項は潜在変数一つ当たりlog Kの定数になり、エンコーダの学習から除けます。一般的なVAEのように、入力由来の分布を事前分布へ近づける圧力が、潜在情報を減らす構造とは異なります。
その後、学習済みエンコーダでデータをコード列に変換し、その並び方を別途学習します。論文では画像にPixelCNN、音声にWaveNetの事前モデルを使います。復元時は「入力→コード→出力」、新規生成時は「事前モデル→コード→出力」と入口が変わります。復元器と事前モデルの同時学習は、本論文では将来課題です。原論文、3.3節
この二段階構成を発展させた例として、VARの残差型マルチスケール表現と次スケール予測があります。VQ-VAE原論文が一つずつのコードを扱う基礎を作り、VARではコードの配置と予測単位を粗い地図から細かい地図へ変えています。
実験と結果
原論文は画像、音声、動画を使い、離散表現でも学習ができるか、短いコードに有用な情報が残るかを検証しています。複数の種類のデータで試していますが、画像と音声を共通のコードブックへ統合する実験ではありません。
CIFAR-10では連続VAEに近い性能を示します
同じ基本ネットワーク構成を用い、潜在変数の容量を変えながら比較しています。学習はAdam、学習率0.0002、バッチサイズ128で25万ステップです。論文が報告する尤度下界をbits/dimへ換算した値は次のとおりです。負の対数尤度に対応するため、小さい方が良好です。
| 方式 | bits/dim |
|---|---|
| 連続潜在変数のVAE | 4.51 |
| VQ-VAE | 4.67 |
| 離散潜在変数の学習法VIMCO | 5.14 |
VQ-VAEはこの比較のVIMCOより良好ですが、連続VAEを上回ってはいません。離散表現を得ながら連続方式に近い値へ到達したことがポイントです。報告値は厳密な尤度そのものではなく下界に基づく評価で、現在の画像生成モデルとの画質ランキングを示すものでもありません。原論文、4.1節
画像を短いコードの地図へ変換できます
ImageNetの128×128×3の画像を、K＝512のコードで32×32の潜在マップへ変換しています。512候補の番号は固定長なら9ビットなので、元画像を各色8ビットとすると符号量は次の比率です。
(128 × 128 × 3 × 8) / (32 × 32 × 9) ≈ 42.6
これは元の画素値とコード番号の名目的なビット数の比較です。モデル・コードブックの配布容量、ファイル形式、整数配列の保存方法は含まれません。JPEGなどと比べた実ファイルの圧縮率や、42.6倍の推論高速化を意味しません。
図2では、元画像よりややぼけるものの内容を残した再構成を示し、図3ではコードのPixelCNN事前モデルから画像を生成しています。またDeepMind Labでは二段階のモデルを用い、上段に3個のコード、計27ビットという強い制約をかけても、場面の構造を利用する例を示しています。こちらは細部まで完全再現する結果ではありません。原論文、4.2節・図2〜5
音声ではコードと音素の対応を調べています
VCTKの109話者を使う実験では、エンコーダで時間方向を64分の1へ縮め、K＝512のコードと話者IDをデコーダへ与えています。コードを固定し話者IDを変える話者変換を、サンプルで示しています。発話内容を残しても波形や韻律は変わり、原音を厳密に保存する方式ではありません。
別の分析では、128種類・毎秒25コードの表現について、各コードを最も対応しやすい音素へ割り当てると、41音素の分類精度が49.3%でした。比較基準として、最頻音素を予測する場合の7.2%を示しています。音素ラベルはVQ-VAEの学習には使わず、学習後の対応分析に使っています。
この結果はコードが音素に関連する情報を持つ証拠ですが、文字起こしの精度49.3%という意味ではありません。音声生成や話者変換の例にも、一般的な音声認識ベンチマークの優位性まで読み込むことはできません。原論文、4.3節
何に使える？
ここからは、技術の構造に基づく応用案です。原論文が各業務用途での効果を検証したわけではありません。
専用の画像生成器の前処理を作ります
ゲームの小さなマップや限られた種類の素材など、対象を絞った画像をコードへ変換し、その配置を別の生成器で予測する構成が考えられます。復元器が画素への展開を担うため、生成器は短いコードの構成に集中できます。
ただし、細い線や小さな文字が復元段階で失われれば、コードの予測器を強化するだけでは忠実な再現が難しくなります。最初に対象データの再構成品質を確かめる必要があります。
音声を予測モデルが扱う単位へ変換します
波形サンプルより短いコード列を作り、その列の補完や生成を学ぶ構成に使えます。原論文の話者IDによる条件付けも、内容と出力条件を別々に与える設計例です。
音の忠実な圧縮を目的にするなら、単段の量子化から複数段で誤差を補う方式へ進む選択肢があります。EnCodecの残差ベクトル量子化と聴感品質の学習は、その違いを考える具体例です。VQ-VAEの音素に関連する表現と、高忠実度コーデックの表現は、同じ目的とは限りません。
環境の未来を潜在空間で予測します
原論文のDeepMind Lab動画実験では、最初の6フレームと行動条件から10フレーム先までの系列を生成しています。まず潜在コード列を予測し、必要な段階で画像へ戻す構成です。原論文、4.4節
これを踏まえると、簡単なゲーム環境で複数の行動候補を比較する研究に応用する余地があります。ただし、長期間の物理整合性や方策の改善を保証する結果ではなく、実システムの制御に使うには別の評価が必要です。
開発や事業へのヒント
生成器の前に、表現の合格条件を決めます
小規模な検証では、いきなり大きな事前モデルを学習するより、対象データを符号化・復元する工程から始められます。画像なら輪郭、文字、色、音声なら発話内容や聞きやすさなど、用途上失えない情報を決めます。
そのうえでコードの個数と候補数Kを変え、再構成品質、符号量、実行時間を比較します。元データを入れても復元できない情報と、コードの生成を間違えて失った情報を分けると、改善対象を見つけやすくなります。
圧縮・復元と、予測の評価を分けます
復元がよくても、新しいコード列をうまく生成できるとは限りません。逆に、次コードの予測損失が小さくても、用途に重要な細部が最初からコードに入っていなければ、出力は役に立ちません。
開発では「元データからの復元」と「生成したコードからの出力」を別に採点します。コードが実際に何種類使われているかも記録すれば、候補を増やしたのに少数の番号へ偏っている状況を把握できます。これは実装上の検証案で、本論文が確立した業務評価手順ではありません。
コード列を再利用するなら版を固定します
コードを事前計算しておけば、予測器の実験ごとに生データをエンコードし直す処理を減らせます。ただし番号の意味は学習したコードブックに依存します。番号37がどのモデルでも同じ特徴を表すわけではありません。
保存時にはエンコーダ、コードブック、デコーダの版をひも付けます。トークナイザを再学習する場合には、過去のコード列や予測器を使い続けられるか、再変換・再学習が必要かを確認します。表現を共有部品にするなら、この互換性も設計対象になります。
限界
VQ-VAEは情報を落とす表現です。画像のぼけや音声の韻律変化のように、見聞きして自然な出力と入力への忠実性は一致しません。業務用途で必要な微細情報が残るかは、対象データで評価する必要があります。
離散化の学習は勾配の近似に依存します。論文が事後崩壊の回避を示していても、どんなデコーダ・学習設定でも有用なコードが得られる保証ではありません。コードが十分に使われない問題と、デコーダが潜在変数を無視する問題も区別して調べるべきです。
計算コストも残ります。潜在位置数をNとすると、全候補を比較する素朴な最近傍探索はおおむねO(NKD)です。これは探索手順から分かる計算量で、論文の実測値ではありません。距離行列をまとめて作る実装ではメモリにも注意が必要です。エンコーダ・デコーダの実行、別途学習する事前モデルの費用もかかります。
さらに、原論文の事前モデルは自己回帰です。コード列が短くなっても、逐次生成の待ち時間がなくなるわけではありません。学習データと異なる画像・言語・録音条件への品質、現行モデルに対する優位性、日本語業務データへの適合は、この論文の結果だけから判断できません。
コード番号そのものも意味検索用の距離ではありません。復元を目的に学んだ局所コードをそのままRAGの文書埋め込みの代わりにする根拠はなく、検索へ応用するなら専用の表現設計と評価が必要です。
よくある質問
Q. VQ-VAEだけで新しい画像を生成できますか？
A. エンコーダ・量子化器・デコーダは、入力をコードへ変換して復元する部分です。入力なしで自然な新規画像を作るには、原論文のPixelCNNのように、妥当なコードの並び方を学んだ事前モデルを組み合わせます。
Q. コードの番号は物体名や単語を意味しますか？
A. 人が決めたラベルではありません。復元を通じて学習した代表ベクトルへの参照番号です。音素との関連が観察されても、一つの番号が必ず一つの単語や物体に対応するわけではありません。
Q. ベクトル量子化は、LLMの4ビット量子化と同じですか？
A. 本記事の量子化は、入力から得た特徴ベクトル全体をコードブックの代表へ置き換えるものです。モデルの重みを低ビットで保存して省メモリ化する方式とは対象が違います。VQ-VAEを導入するだけでモデルの全重みが小さくなるわけではありません。
Q. Kを大きくすれば品質は上がりますか？
A. 代表候補は増えますが、品質改善は保証されません。固定長で保存する番号のビット数や探索コストが増え、候補を十分に使い分けられるかも学習に依存します。潜在位置数とKを分けて変え、復元品質とコード使用率を確認します。
Q. 学習済みVQ-VAEを別データへそのまま使えますか？
A. 入出力の形が合えば試せますが、品質は保証されません。例えば自然画像で覚えたコードが、細かい図面の線を残せるとは限りません。まず生成器を使わずに対象データの符号化・復元を行い、重要情報が保持されるかを測るのが先です。
今日の学び
この論文は、画像や音声を短く有用な表現へ変えつつ、離散選択を含むモデルを学習する課題を扱いました。最近傍のコード選択、勾配の近似、更新先を分けた三つの損失を組み合わせ、コードの予測と入力への復元を別のモデルへ分担します。
得られるヒントは、生成モデルの性能を考える際に「何を予測するか」も設計できることです。予測器の拡大に進む前に、用途に必要な情報が残る表現を作り、復元品質と生成品質を分けて測ることが開発の出発点になります。