はじめに
この記事はAIで作成されています。multi-modal LLM関連のメモを記載しています。
目次
Vector Quantization
画像や音声を生成するモデルでは、データを連続値の潜在ベクトルとして表現するだけでなく、離散的なトークンに変換して扱う方法がある。
代表的な手法には、次の4つがある。
| 方式 | 離散化の方法 | 学習可能なコードブック | 特徴 |
|---|---|---|---|
| ① VQ-VAE | 潜在ベクトルを最も近いコードブックのベクトルに置き換える | あり | 学習可能な離散表現を獲得する |
| ② RVQ | 残差を段階的に量子化する | あり | 複数段の量子化によって表現能力を高める |
| ③ FSQ | 潜在ベクトルの各次元を固定された離散値に量子化する | なし | コードブックの学習を不要にする |
| ④ LFQ | 潜在ベクトルの各次元を2値化する | なし | ビット列によって離散表現を構成する |
現在の主流
画像・音声生成などでは、目的に応じてさまざまな量子化手法が使われている。
- VQ-VAE:学習可能なコードブックを利用して、入力データに適した離散表現を学習する。
- RVQ:複数段の量子化によって、表現の精度や容量を高める。音声コーデックなどで利用される。
- FSQ:コードブックを学習せず、各次元を固定された有限個の値に量子化する。
- LFQ:各次元を2値化し、ビット列として離散表現を構成する。
これらは単純な新旧関係にあるわけではなく、用途やモデルの設計によって使い分けられる。
① VQ-VAE(Vector Quantized Variational Autoencoder)
論文名: Neural Discrete Representation Learning
URL: https://arxiv.org/abs/1711.00937
発表日: 2017-11-02
✅ ポイント
- 通常の VAE では、潜在変数を確率分布からサンプリングすることで、連続的な潜在表現を学習する。
- VQ-VAE では、Encoder が出力した潜在ベクトルを、学習可能なコードブックに登録されたベクトルのうち、最も近いものに置き換える。
- これにより、潜在表現を連続値ではなく、離散的なインデックスとして扱える。
- 潜在変数の確率分布を直接学習する従来型 VAE とは異なり、離散的なコードを使って表現を学習する。ただし、VQ-VAE 自体がデータの多様性を必ず高めるわけではない。
- コードブックの学習には工夫が必要であり、コードの利用が偏る、特定のコードが使われなくなるといった問題が発生することがある。
🧪 学習の疑似コード
# Encoder、Decoder、コードブックを初期化
initialize(encoder, decoder, codebook)
while not converged:
x = sample_from_dataset()
# Encoder で連続値の潜在ベクトルを生成
z_e = encoder(x)
# 各潜在ベクトルに最も近いコードブックのベクトルを選択
indices = nearest_codebook_indices(z_e, codebook)
z_q = codebook[indices]
# 量子化された潜在表現から入力を再構成
x_reconstructed = decoder(z_q)
# 再構成誤差
reconstruction_loss = compute_loss(x_reconstructed, x)
# Encoder の出力と量子化ベクトルの整合性を学習
commitment_loss = compute_commitment_loss(z_e, z_q)
# コードブックの学習誤差
codebook_loss = compute_codebook_loss(z_e, z_q)
# 損失を統合してモデルを更新
loss = (
reconstruction_loss
+ commitment_loss
+ codebook_loss
)
update(loss)
実際には、コードブックの更新に勾配降下法を使う方式のほか、指数移動平均(EMA)によって更新する方式もある。
また、量子化処理は通常の微分ができないため、Encoder への勾配伝播には Straight-Through Estimator(STE)などの工夫を使う。
🔎 推論時の疑似コード
x = input_image()
# Encoder で連続値の潜在ベクトルを生成
z_e = encoder(x)
# 最も近いコードブックのベクトルを選択
indices = nearest_codebook_indices(z_e, codebook)
z_q = codebook[indices]
# 量子化された潜在表現から画像を再構成
x_reconstructed = decoder(z_q)
推論時には、学習済みの Encoder、Decoder、コードブックを使用する。
画像を再構成するだけでなく、コードブックのインデックスを離散トークンとして保存し、後段の生成モデルに入力することもできる。
🧮 量子化の計算式
Encoder が出力する潜在ベクトルを (z_e(x))、コードブックのベクトルを (e_k) とする。
最も近いコードブックのインデックスは、次の式で求める。
$$
k^* = \arg\min_k |z_e(x)-e_k|_2^2
$$
量子化後の潜在ベクトルは、次のようになる。
$$
z_q(x)=e_{k^*}
$$
つまり、Encoder が生成した連続値のベクトルを、コードブック内のいずれかのベクトルに置き換える。
学習時の代表的な損失関数は、次のように表せる。
$$
\mathcal{L}
=\mathcal{L}_{\mathrm{recon}}
+
|\operatorname{sg}[z_e]-e|_2^2
+
\beta|z_e-\operatorname{sg}[e]|_2^2
$$
ここで、
- $(\mathcal{L}_{\mathrm{recon}})$:入力データの再構成誤差
- $(e)$:選択されたコードブックのベクトル
- $(\operatorname{sg})$:Stop Gradient。勾配を伝播させない操作
- $(\beta)$:Commitment Loss の重み
第2項はコードブックを Encoder の出力に近づけ、第3項は Encoder の出力が選択されたコードブックのベクトルに近づくよう促す。
❓ よくある質問
Q. なぜ連続値ではなく離散値にするの?
- 潜在表現をトークンとして扱えるようにするため。
- 離散トークンに変換することで、Transformer などを用いて画像や音声の生成を学習できる。
- 潜在表現を圧縮して扱いやすくなる場合もある。
Q. VQ-VAE は通常の VAE と何が違うの?
- 通常の VAE は、潜在変数の確率分布を学習し、そこから潜在変数をサンプリングする。
- VQ-VAE は、Encoder の出力をコードブック内のベクトルに置き換え、離散的な潜在表現を得る。
- ただし、VQ-VAE でも、後段に事前分布を学習する生成モデルを組み合わせることができる。
Q. コードブックはどのように学習するの?
- Encoder が出力する潜在ベクトルに近づくように、コードブックのベクトルを更新する。
- コードブックの更新には勾配降下法や EMA などの方法がある。
- コードの利用が偏ると、ほとんど使われないコードが生じることがある。
Q. なぜ Straight-Through Estimator が必要なの?
- 最近傍のコードを選択する処理は、通常の微分では勾配を伝播できないため。
- 順伝播では離散化を行い、逆伝播では量子化前の潜在表現に勾配を近似的に伝える。
② RVQ(Residual Vector Quantization)
RVQ は、Residual Vector Quantization(残差ベクトル量子化)の略である。
代表的な関連論文: SoundStream: An End-to-End Neural Audio Codec
URL: https://arxiv.org/abs/2107.03312
発表日: 2021-07-07
✅ ポイント
- VQ-VAE のように1つのコードブックで量子化するのではなく、複数段のコードブックを使って段階的に量子化する。
- 最初のコードブックで潜在ベクトルを近似し、その近似で表現しきれなかった残差を次のコードブックで近似する。
- 各段の出力を加算することで、元の潜在ベクトルに近い表現を得る。
- 量子化段数を増やすほど表現能力を高められるが、その分、必要なトークン数や計算量などが増える場合がある。
- 音声コーデックでは、量子化段数を調整することで、ビットレートと音質のトレードオフを制御できる。
🧪 学習の疑似コード
# Encoder、Decoder、複数のコードブックを初期化
initialize(encoder, decoder, codebooks)
while not converged:
x = sample_from_dataset()
# Encoder で潜在ベクトルを生成
z = encoder(x)
# 初期残差を設定
residual = z
# 各段の量子化結果を加算するための変数
z_q = zeros_like(z)
for codebook in codebooks:
# 現在の残差に最も近いコードを選択
q = nearest_codebook_vector(residual, codebook)
# 量子化結果を加算
z_q += q
# 量子化で表現できなかった残差を計算
residual = residual - q
# 量子化された潜在表現から入力を再構成
x_reconstructed = decoder(z_q)
# 再構成誤差などを計算
loss = compute_rvq_loss(
x,
x_reconstructed,
z,
z_q,
codebooks
)
# モデルとコードブックを更新
update(loss)
RVQ では、各段のコードブックが残差を順番に近似する。
実際の損失関数やコードブックの更新方法は、採用するモデルによって異なる。
🔎 推論時の疑似コード
x = input_audio()
# Encoder で潜在ベクトルを生成
z = encoder(x)
residual = z
z_q = zeros_like(z)
indices_per_stage = []
for codebook in codebooks:
# 現在の残差を量子化
index = nearest_codebook_index(residual, codebook)
q = codebook[index]
# コードのインデックスを保存
indices_per_stage.append(index)
# 量子化結果を加算
z_q += q
# 次の段で扱う残差を計算
residual = residual - q
# Decoder で再構成
x_reconstructed = decoder(z_q)
推論時には、各段のコードブックからインデックスを取得する。
各段のインデックスを保存すれば、元の潜在ベクトルをコードブックのベクトルの和として復元できる。
🧮 残差量子化の計算式
入力となる潜在ベクトルを (z)、第 (m) 段のコードブックを (C_m) とする。
最初の残差は次のように設定する。
$$
r_0=z
$$
第 (m) 段では、現在の残差に最も近いコードブックのベクトルを選択する。
$$
q_m=
\arg\min_{e\in C_m}
|r_{m-1}-e|_2^2
$$
そして、残差を更新する。
$$
r_m=r_{m-1}-q_m
$$
最終的な量子化表現は、各段の出力を加算したものになる。
$$
z_q=\sum_{m=1}^{M}q_m
$$
ここで、(M) は量子化段数である。
各段が前段までで表現しきれなかった情報を補うため、段数を増やすことで近似精度を高められる。
❓ よくある質問
Q. なぜ残差を段階的に量子化するの?
- 1つのコードブックだけでは表現しきれない情報を、複数段に分けて近似するため。
- 各段が残差を補正するため、量子化段数を増やすことで表現精度を高められる。
Q. VQ-VAE との違いは?
- VQ-VAE の基本形では、1つのコードブックからコードを選択する。
- RVQ では、複数のコードブックを使い、残差を順番に量子化する。
- なお、RVQ は VQ-VAE と対立する概念ではなく、VQ-VAE 系のモデルにも組み込める量子化方式である。
Q. 段数を増やすとどうなるの?
- 一般に、表現できる情報量が増え、量子化誤差を小さくできる。
- 一方、保存・伝送するコードの数や計算量も増える。
- そのため、必要な品質と圧縮率に応じて段数を調整する。
Q. 各段のコードブックは同じものを使うの?
- 一般的な RVQ では、各段にそれぞれのコードブックを用意する。
- 段ごとに異なる残差を扱うため、それぞれ異なる情報を表現するように学習される。
③ FSQ(Finite Scalar Quantization)
論文名: Finite Scalar Quantization: VQ-VAE Made Simple
URL: https://arxiv.org/abs/2309.15505
発表日: 2023-09-26
✅ ポイント
- VQ-VAE では、コードブックの学習や管理が必要になる。
- FSQ では、潜在ベクトルの各次元を、あらかじめ決められた有限個の値に量子化する。
- 各次元の量子化レベルは固定されているため、学習可能なコードブックを持つ必要がない。
- 各次元の離散値の組み合わせによって、多数の離散コードを表現できる。
- コードブックの利用が偏る問題を回避しやすく、コードブックの更新に伴う学習上の複雑さを減らせる。
🧪 学習の疑似コード
# Encoder、Decoder、各次元の量子化レベルを初期化
initialize(encoder, decoder)
levels = [5, 5, 5, 5] # 各次元の量子化レベル数の例
while not converged:
x = sample_from_dataset()
# Encoder で潜在ベクトルを生成
z_e = encoder(x)
# 潜在ベクトルの各次元を固定された値に量子化
z_q = finite_scalar_quantize(z_e, levels)
# 量子化された潜在表現から入力を再構成
x_reconstructed = decoder(z_q)
# 再構成誤差を計算
loss = compute_loss(x_reconstructed, x)
# 量子化処理を通して勾配を伝播させ、モデルを更新
update(loss)
FSQ では、コードブック自体を学習する必要がない。
ただし、量子化処理は離散的であるため、学習時には Straight-Through Estimator などを使って勾配を伝播させる。
また、量子化レベルの数や潜在ベクトルの次元数は、モデルの設計時に決める。
🔎 推論時の疑似コード
x = input_image()
# Encoder で潜在ベクトルを生成
z_e = encoder(x)
# 各次元を固定された量子化レベルに変換
z_q = finite_scalar_quantize(z_e, levels)
# 必要に応じて量子化値を離散インデックスに変換
indices = encode_scalar_levels(z_q, levels)
# Decoder で再構成
x_reconstructed = decoder(z_q)
推論時には、学習済みの Encoder と Decoder、および学習時に設定した量子化レベルを使用する。
各次元の離散値を組み合わせれば、離散トークンとして扱うこともできる。
🧮 量子化の計算式
潜在ベクトルを次のように表す。
$$
z=(z_1,z_2,\ldots,z_d)
$$
各次元 (i) に対して、有限個の量子化レベルを設定する。
例えば、量子化レベルが5段階なら、各次元は5種類の値のいずれかに変換される。
$$
\hat{z}_i=Q_i(z_i)
$$
ここで、(Q_i) は次元 (i) の量子化関数である。
各次元の量子化レベル数を (L_i) とすると、組み合わせによって得られる離散コードの総数は、理論上、次のようになる。
$$
N=\prod_{i=1}^{d}L_i
$$
例えば、4次元の潜在ベクトルの各次元を5段階に量子化する場合、
$$
N=5^4=625
$$
となる。
つまり、学習可能なコードブックを用意しなくても、各次元の有限個の値の組み合わせによって625種類の離散コードを表現できる。
なお、実際の FSQ では、量子化しやすい範囲へのスケーリングや、値を所定の範囲に収める変換などを用いる。上記は仕組みを理解するための簡略化した説明である。
❓ よくある質問
Q. なぜコードブックを使わなくてよいの?
- VQ-VAE は、学習可能なベクトルの集合から最も近いベクトルを選ぶ。
- FSQ は、各次元を固定された量子化レベルに変換する。
- そのため、コードブックのベクトルを学習・管理する必要がない。
Q. コードブックがないのに、どうやって多くのコードを表現するの?
- 各次元に複数の離散値を設定し、それらの組み合わせによってコードを構成する。
- 例えば、各次元が5種類の値を取れる4次元ベクトルなら、理論上625種類の組み合わせがある。
Q. VQ-VAE より必ず性能が高いの?
- 必ずしもそうではない。
- FSQ はコードブックの学習を不要にすることで、学習の安定性や実装の簡潔さを改善しやすい。
- 一方、適切な量子化レベルや潜在次元数は、タスクやモデルに応じて設計する必要がある。
Q. 量子化レベルは学習するの?
- 基本的には、量子化レベルの数や各レベルの値は固定する。
- 学習するのは主に Encoder と Decoder などのモデルパラメータである。
④ LFQ(Lookup-Free Quantization)
LFQ は Lookup-Free Quantization の略である。
代表的な関連論文: Language Model Beats Diffusion — Tokenizer Is Key to Visual Generation
URL: https://arxiv.org/abs/2310.05737
発表日: 2023-10-08
✅ ポイント
- VQ-VAE では、潜在ベクトルに近いコードブックのベクトルを検索する必要がある。
- LFQ では、潜在ベクトルの各次元を2値化し、ビット列として離散表現を構成する。
- 量子化処理に学習可能なコードブックを必要とせず、最近傍ベクトルの検索も行わない。
- 潜在ベクトルの各次元が2種類の値を取る場合、(d) 次元の組み合わせによって理論上 (2^d) 種類のコードを表現できる。
- コードブックの学習を不要にする一方、表現能力は潜在次元数や量子化の設計に依存する。
🧪 学習の疑似コード
# Encoder、Decoder を初期化
initialize(encoder, decoder)
while not converged:
x = sample_from_dataset()
# Encoder で潜在ベクトルを生成
z_e = encoder(x)
# 潜在ベクトルの各次元を2値化
z_q = binary_quantize(z_e)
# 2値化された潜在表現から入力を再構成
x_reconstructed = decoder(z_q)
# 再構成誤差を計算
reconstruction_loss = compute_loss(
x_reconstructed,
x
)
# 必要に応じてエントロピー正則化などを計算
entropy_loss = compute_entropy_loss(z_e)
# 損失を統合
loss = reconstruction_loss + entropy_loss
# モデルを更新
update(loss)
これは基本的な処理の流れを示した疑似コードである。
実際の LFQ 系モデルでは、量子化に加えて、コードの利用状況を改善するためのエントロピー正則化などを導入する場合がある。損失関数の構成は実装によって異なる。
🔎 推論時の疑似コード
x = input_image()
# Encoder で潜在ベクトルを生成
z_e = encoder(x)
# 各次元を2値化
z_q = binary_quantize(z_e)
# 2値の組み合わせを整数などのトークンIDに変換
token_ids = bits_to_integer(z_q)
# Decoder で再構成
x_reconstructed = decoder(z_q)
推論時には、Encoder が出力した潜在ベクトルを2値化する。
各次元の2値をまとめて整数に変換すれば、離散トークンの ID として保存・利用できる。
ただし、実際の LFQ では、量子化された潜在表現の形状やトークンへの変換方法はモデルの設計によって異なる。
🧮 量子化の計算式
潜在ベクトルを次のように表す。
$$
z=(z_1,z_2,\ldots,z_d)
$$
各次元を、符号に応じて2値化する。
$$
\hat{z}_i=
\begin{cases}
+1 & z_i \geq 0\
-1 & z_i < 0
\end{cases}
$$
この処理によって、各次元は (+1) または (-1) のいずれかになる。
各次元が2種類の値を取るため、(d) 次元の潜在ベクトルが表現できる組み合わせの総数は、
$$
N=2^d
$$
となる。
例えば、16次元の潜在ベクトルを2値化した場合、
$$
N=2^{16}=65,536
$$
種類のコードを理論上表現できる。
このように、コードブックを明示的に保持しなくても、多数の離散コードを構成できる。
なお、これは理論上のコード数であり、学習済みモデルがすべてのコードを実際に利用するとは限らない。
❓ よくある質問
Q. なぜ Lookup-Free と呼ばれるの?
- VQ-VAE のように、コードブックの中から最近傍のベクトルを検索する必要がないため。
- 各次元を直接2値化することで、離散表現を構成する。
Q. コードブックがないのに、どうやってコードを識別するの?
- 各次元の2値を組み合わせて、コードを一意に表現する。
- 例えば4次元なら、2値の組み合わせによって16種類のコードを表現できる。
Q. FSQ との違いは?
- FSQ は各次元を有限個の値に量子化する。1次元あたり3種類、5種類など、2種類以外の値も設定できる。
- LFQ は基本的に各次元を2値化する。
- どちらも学習可能なコードブックを必要としないが、離散表現の構成方法が異なる。
Q. 潜在次元数を増やすとどうなるの?
- 理論上のコード数は (2^d) で増加する。
- 一方、潜在次元数を増やすと、モデルの計算量や表現の管理に必要なコストも増える場合がある。
- また、コード数が多くても、学習データやモデルの制約によって利用されないコードが生じることがある。
Q. LFQ では学習が不要なの?
- 量子化処理やコードブックの学習は不要だが、モデル全体の学習は必要である。
- Encoder や Decoder などは、目的のデータを適切に表現・再構成できるように学習する。
4つの手法の比較
| 比較項目 | VQ-VAE | RVQ | FSQ | LFQ |
|---|---|---|---|---|
| 基本的な考え方 | 最近傍のコードを選択 | 残差を段階的に量子化 | 各次元を有限個の値に量子化 | 各次元を2値化 |
| コードブック | 学習可能 | 各段に学習可能なコードブック | 不要 | 不要 |
| 量子化単位 | ベクトル | ベクトルの残差 | スカラー(各次元) | 2値(各次元) |
| コードの検索 | 必要 | 各段で必要 | 不要 | 不要 |
| 表現可能なコード数 | コードブックのサイズに依存 | 各段のコード数の積に依存 | 各次元の量子化レベル数の積 | 潜在次元数 (d) に対して (2^d) |
| 主な利点 | データに適した離散表現を学習できる | 段階的に表現精度を高められる | コードブックの学習が不要 | 単純な2値化で離散表現を構成できる |
| 主な課題 | コードブックの学習・利用の偏り | 段数に応じた計算量・符号量 | 量子化レベルや潜在次元の設計 | 2値化による表現上の制約 |
まとめ
これらの手法の違いは、連続的な潜在ベクトルを、どのような方法で離散的なコードに変換するかにある。
- VQ-VAE:学習可能なコードブックから、最も近いベクトルを選択する。
- RVQ:残差を複数段に分けて量子化し、各段の出力を加算する。
- FSQ:各次元を固定された有限個の値に量子化し、コードブックの学習を不要にする。
- LFQ:各次元を2値化し、その組み合わせによって離散コードを構成する。
VQ-VAE から RVQ への違いは主に量子化の段階化、FSQ と LFQ の違いは主にコードブックを使わずに離散表現を構成する方法にある。
なお、RVQ は VQ-VAE と同時に利用できる量子化方式であり、FSQ と LFQ はコードブックの学習を避けるための別のアプローチである。そのため、4つの手法を単純な進化の順序として捉えるのではなく、それぞれの設計上の特徴として理解することが重要である。