画像生成におけるVAEの周波数分離とその検証

背景と課題

画像や動画の拡散モデルでは VRAM 不足が課題となっており、扱う Tensor のサイズを縮小するために、画像の圧縮率を可能な限り高める必要があります。
そのためには、圧縮率を高めつつ再構成誤差を最小限に抑えられるよう、VAE の品質を向上させることが重要です。

しかし、画像の抽象的・構造的な部分（低周波）を捉えるための重要なチャンネルが、多数の詳細部分（高周波）の中に分散してしまうという問題が知られています（ Channel-wise dispersion に関する研究 (arXiv:2508.00413) ）。

潜在チャンネルを増やすことで VAE 自体の再構成誤差は減少しますが、拡散モデルに用いた場合の画像生成の品質は逆に低下してしまいます。さらに、潜在チャンネルの増加は本質的に圧縮率の低下を招くという問題も抱えています。

---

離散ウェーブレット変換による分離

低周波を捉える重要なチャンネルが高周波の中に分散してしまう問題をケアするために、離散ウェーブレット変換（Discrete Wavelet Transform, DWT）を用いて画像を抽象的・構造的な情報（低周波）と詳細な情報（高周波）に分離します。これにより低周波と高周波の交雑を防いだうえで、潜在変数を出力するように設計します。

手法1: Organized AE

離散ウェーブレット変換を使用して、画像を数学的に異なる周波数成分へ分解します。

• 低周波成分 : 画像の全体的な形状や輪郭などの構造に対応します。
• 高周波成分 : 画像のテクスチャや細かいエッジなどのディテールに対応します。

情報の流れを制御するために、エンコーダとデコーダに特別な Transformer を導入します。attention mask を使用し、ある周波数帯の情報を処理する際には、それより低い周波数帯（より抽象的な構造情報）のみを参照するように制限します。あわせて自己回帰的なアーキテクチャを採用します。この制限により、低周波の重要な構造情報を破壊することなく、高周波の詳細情報を適切に学習させることを狙います。

この手法は Organized AE の論文 (arXiv:2412.09607) を参考にしました。

アーキテクチャ

入力は 2 の累乗の大きさの画像を受け付けることを想定しています。
低周波トークン（approx_tokens）のサイズ $(B,N,D)$ に合わせて、高周波トークン（detail_tokens）のサイズを一致させていくプロセスを採ります。

パッチの数 $N$ を揃えるために、高周波側の広い解像度を折りたたみ（Pixel Unshuffle）、溢れた情報をチャンネル方向に詰め込んだうえで、最後に線形層（Linear）で特徴量次元 $D$ をちょうどのサイズに削ります。

位置埋め込みは 2 種類を用います。

• intra_scale_pos_embedding : パッチが画像内の「どこ（上・真ん中・下など）」にあるかを示します。
• scale_embeddings : その情報が「どの周波数帯」にあるかを示します。

上記の処理によって整理されたデータを、周波数スケール間で自己回帰的（粗から細へ）に処理します。すなわち「大まかなスケールのコンテキストに基づいて、詳細なスケールの表現を段階的に予測・エンコードする」という自己回帰的な情報流（粗 $\to$ 細）を強制します。その過程で周波数帯域を交錯させないように次元削減を行い、エンコーダ出力の Tensor の shape に合わせます。

このようなアーキテクチャを採用しました。

Organized AE の結果

ImageNet で 20 epoch、GAN なしで学習しました。圧縮は $(B,3,256,256)\to (B,4,32,32)$ です。
損失関数は低周波・高周波で共通とし、次式を用います。

$$L_{total}=L_{MAE}+L_{LPIPS}+\beta L_{KL}$$

下図は上段が元画像、下段が再構成画像です。

Loss の推移について、低周波成分と 2 つの高周波成分群（群1・群2）ごとに見てみると、下図のようになりました。

ここで損失関数が共通であることにより、低周波の方が損失が大きく、高周波部分の損失は小さくなります。その結果、学習が低周波部分に着目してしまい、細部の高周波部分がおざなりになるという問題がありました（自己回帰のための attention mask を OFF にしても同様の結果でした）。

手法2: FAVAE

Organized AE の結果からのモチベーション

アーキテクチャによって流すデータを低周波と高周波に分離するだけでなく、損失関数も分離しなければ、Organized AE と同じ問題が起こってしまいます。
構造的な特徴（低周波部分）と詳細な特徴（高周波部分）に分離したままエンコーダから出力し潜在画像にする、という目標についても、引き続き離散ウェーブレット変換を使用します。

FAVAE を参考にしたアーキテクチャ

FAVAE の論文 (arXiv:2509.05441v1) のアーキテクチャや損失関数を、よりシンプルに再現したものです。

本記事で実装した FAVAE と本家 FAVAE の違い

本家 FAVAE は、離散ウェーブレット変換したデータ（LL, LH, HL, HH）のうち、低周波部分である LL を高周波部分とは別のエンコーダニューラルネットワークに通し、高周波部分（LH, HL, HH）はまとめて受け付けるエンコーダニューラルネットワークに通します。それらを潜在変数の部分で concat し、サンプリング処理を行ったうえで、同じ仕組みのデコーダを採用します。損失関数は DWT 直後と IDWT 直前のウェーブレット係数同士で取ることで、損失関数も分離するアプローチを採っています。また、低周波部分のみに VFM（Vision Foundation Model）アライメント損失を採用し、潜在空間における整列を行っています。

本記事で実装したコードは、（LL, LH, HL, HH）それぞれに対し完全に分離した 4 つのアーキテクチャを採用し、潜在画像の部分で concat します。損失関数は本家と同様に DWT 直後と IDWT 直前のウェーブレット係数同士で取り、損失関数を分離します。この実装にした理由は、VFM を使ったアライメントの実装方法が分からなかったこと、そして Diffusers ベースの VAE のニューラルネットワークを 4 つ分離した形で並べればよく、実装が容易だったことです。

なお、本家 FAVAE の方が論文上は性能が圧倒的に良いですが、以降では本記事で実装したものを「FAVAE」と呼ぶことにします。

FAVAE のデータの流れ

=== 入力データ ===

Input 'x' shape: torch.Size([2, 3, 256, 256]) -> (B=2, C=3, H=256, W=256)

=== エンコーダ (WaveletEncoder) ===
DWT 係数 'LL' shape: ([2, 3, 128, 128])
DWT 係数 'LH' shape: ([2, 3, 128, 128])
DWT 係数 'HL' shape: ([2, 3, 128, 128])
DWT 係数 'HH' shape: ([2, 3, 128, 128])

Encoded 'z_concat' shape: ([2, 4, 2, 32, 32])

=== 潜在空間のサンプリング ===
Flat 'z_flat' shape: ([8, 2, 32, 32])
Sampled 'z_sampled_flat' shape: ([8, 1, 32, 32])
Reshaped 'z' shape: ([2, 4, 1, 32, 32])

=== デコーダ ===
'rec_ll' shape: ([2, 3, 128, 128])
'rec_lh' shape: ([2, 3, 128, 128])
'rec_hl' shape: ([2, 3, 128, 128])
'rec_hh' shape: ([2, 3, 128, 128])
Reconstructed 'x_recon' shape: ([2, 3, 256, 256]) -> (B, C, H, W)

なお、 z_concat の 3 番目の次元 2 は、VAE のエンコーダが出力する平均 μ と対数分散 logσ² の対を表します。サンプリングによりこれが 1 になり、最終的な潜在表現は 4 グループ × 1 チャンネル = 合計 4 チャンネル（(B,4,32,32) 相当）になります。

FAVAE の損失関数

低周波（low-frequency）成分の損失は次式です。

$$L_{\text{low}}=L_{MAE}+L_{LPIPS}+\beta L_{KL}$$

高周波（high-frequency）成分の損失は次式です。

$$L_{\text{high}}=L_{rec}+\beta L_{KL}+{\lambda }_{gan}\cdot w_{gan}\cdot L_{gan}$$

ここで ${\lambda }_{gan}$ は、再構成損失と GAN 損失の勾配ノルムの比として動的に決定されます。

$${\lambda }_{gan}=\frac{\Vert {\mathrm{\nabla }}_{W_{last}}{L}_{rec}{\Vert }_2}{\Vert {\mathrm{\nabla }}_{W_{last}}{L}_{gan}{\Vert }_2+ϵ}$$

${\mathrm{\nabla }}_{W_{last}}{L}_{rec}$ は、再構成損失（MAE + LPIPS）の最終層の重みに関する勾配を表します。

損失関数における $w_{gan}$ の制御は、テクスチャの復元品質と生成の安定性を左右する最重要パラメータです。GAN Loss はリアルなテクスチャを生成しようとする反面、ターゲット画像との不一致を招き、VAE 本来の基盤ロスを巨大化・崩壊させるリスクがあります。この悪影響を抑え込み、学習を安定させるために、 $w_{gan}$ による厳密なコントロールが不可欠です。

FAVAE 学習における Loss や重みの推移

Total Loss	GAN Loss	λ_gan

圧縮 $(B,3,256,256)\to (B,4,32,32)$ 、 $w_{gan}=0.1$ 。

---

$w_{gan}$ による厳密なコントロールにより、結果に大きな差が出ています。 $w_{gan}=0.1$ では学習の発散がなくなりました。

FAVAE 学習における出力画像の比較

下表はいずれも上段が元画像、下段が再構成画像です。

潜在画像の shape / GAN Loss の重み係数	出力画像
$(B,3,256,256)\to (B,4,64,64)$ 、 $w_{gan}=1$
$(B,3,256,256)\to (B,4,32,32)$ 、 $w_{gan}=1$
$(B,3,256,256)\to (B,4,32,32)$ 、 $w_{gan}=0.1$

FAVAE の定量的・定性的評価

ここでの FID は、再構成画像と元画像の分布間で測る rFID（再構成 FID） であり、生成画像で測る gFID とは別物である点に注意してください。

評価指標	$w_{gan}=1,(B,4,64,64)$	$w_{gan}=1,(B,4,32,32)$	$w_{gan}=0.1,(B,4,32,32)$
LL MAE Average	0.111904	0.174042	0.173272
LH MAE Average	0.044668	0.063085	0.058838
HL MAE Average	0.042921	0.058615	0.056686
HH MAE Average	0.026653	0.042332	0.033365
LPIPS Average	0.245996	0.371687	0.378203
rFID（再構成FID）	28.4942	60.8554	60.1915

FAVAE は $(3,256,256)\to (4,64,64)$ に圧縮する場合は再構成誤差を小さくできましたが、 $(3,256,256)\to (4,32,32)$ に圧縮するタスクは改善できませんでした。ただし $w_{gan}$ の重みを調節することで、学習結果がわずかに良くなりました。
結局、LH・HL・HH に対して損失関数が下がりきらないという問題は継続しているといえます。

他の解像度・圧縮率における FAVAE の検証

下表はいずれも上段が元画像、下段が再構成画像です。

潜在画像の shape / GAN Loss の重み係数	出力画像
$(B,3,1024,1024)\to (B,4,128,128)$ 、 $w_{gan}=1$
$(B,3,1024,1024)\to (B,4,64,64)$ 、 $w_{gan}=1$
$(B,3,1024,1024)\to (B,4,32,32)$ 、 $w_{gan}=0.1$

FAVAE ではなぜ成功し、Organized AE ではなぜ失敗したのか

FAVAE は再構成品質が向上した一方で、Organized AE は再構成品質が改善しませんでした。
Organized AE はウェーブレット変換を複数回実行し、低周波のトークンに合わせて高周波部分を無理やり線形変換してまとめています。構造が複雑であるうえ、高周波部分の 急激な次元削減 が行われたことが、再構成品質の悪化につながったと考えています。

一般的に使われる VAE（Diffusers ベース）との比較

下図は上段が元画像、下段が再構成画像です。

損失関数は次式です。

$$L_{total}=L_{rec}+{\lambda }_{kl}{L}_{kl}+{\lambda }_{gan}\cdot w_{gan}\cdot L_{gan}$$

$${\lambda }_{gan}=\frac{\Vert {\mathrm{\nabla }}_{W_{last}}{L}_{rec}{\Vert }_2}{\Vert {\mathrm{\nabla }}_{W_{last}}{L}_{gan}{\Vert }_2+ϵ}$$

ここで ${\mathrm{\nabla }}_{W_{last}}{L}_{rec}$ は、最終層の重みに関する再構成損失（MAE + LPIPS）の勾配です。条件は $w_{gan}=0.1$ 、 $(B,3,256,256)\to (B,4,32,32)$ です。

評価指標 (Metric)	数値 (Value)
Average MAE	0.0775
Average LPIPS	0.1606
rFID（再構成FID）	18.2483

FAVAE は rFID 60.1915・平均 LPIPS 0.378203 であるのに対し、Diffusers ベース VAE は rFID 18.2483・平均 LPIPS 0.1606 です。再構成品質という観点では、現状の VAE に勝てていません。

Reconstruction-Generation Gap

Reconstruction-Generation Gap の視点における検証とモチベーション

VAE の再構成性能（rFID など）が高くても、その潜在空間を拡散モデルに使ったときの生成性能（gFID）が伴うとは限りません。この「再構成は良いのに生成は良くない」という乖離を Reconstruction-Generation Gap と呼びます。再構成誤差がいくら小さくても、生成性能が高くなければ意味がないため、この乖離をケアすることが重要です。

しかし、実際の生成プロセスまで検証するには膨大な時間と計算が必要になります。そこで、これを短縮するために変動係数・正規化エントロピー・ジニ係数を採用し、潜在画像を評価することで、現状の問題がどのように起こっているのかを分析することにしました。本節では、本記事で作成した FAVAE と通常の VAE について、潜在画像の分析を行います。

通常の VAE の損失関数だけでは、特徴容量を拡張した際に、再構成誤差を減らすために潜在空間を歪めてでも情報を無理やり詰め込む、というエンコーダの過剰適合を抑制できません。潜在空間の統計量を見ることは、その VAE が優れた生成性能（高い gFID）をもたらすかを事前に予測するための理論的アプローチとなりえます。

VA-VAE の論文 (arXiv:2501.01423) においても、特徴分布の均一性が向上するほど生成モデルのパフォーマンス（gFID）も向上すると結論付けています。

潜在画像における特徴分布の均一性の指標

これらの指標は、潜在ベクトルを t-SNE で 2 次元に圧縮し、カーネル密度推定（KDE）で推定した密度分布から算出します（算出方法の詳細は付録を参照）。

• density_cv（変動係数） : 潜在空間の確率密度分布の標準偏差を平均で割った値です。値が低いほど、分布の密度が場所によって偏らず均一に広がっていることを示します。値が高いことは、潜在空間内に「確率質量が異常に高い超過密地帯」と「確率密度がほぼゼロの過疎地帯」が併存することを意味します。
• normalized_entropy（正規化エントロピー） : 分布の情報量を表します。最大値 1 に近いほど特定の特徴に偏らず、空間を最大限効率的かつ一様に使えていることを意味します。1.0 に近いほど潜在空間の表現力をフルに活用できている状態です。
• gini_coefficient（ジニ係数） : 特定領域へのデータの集中度を表します。0 から 1 の値をとり、低いほど空間全体にデータが満遍なく分散していることを示します。

各指標の数学的な定義は、本記事末尾の「付録: 特徴分布の均一性指標の定義」にまとめています。

VA-VAE の論文に記載されていた各指標と gFID の比較

Tokenizer (仕様)	rFID	density_cv	gini_coefficient	normalized_entropy	gFID
f16d32 (32, 16, 16)	0.26	0.263	0.145	0.995	22.62
	0.28	0.193	0.101	0.997	19.89
	0.28	0.178	0.096	0.998	15.82
f16d64 (64, 16, 16)	0.17	0.296	0.166	0.994	36.83
	0.15	0.256	0.143	0.995	23.58
	0.14	0.251	0.141	0.996	24.00

※ 各トークナイザにつき複数行あるのは、論文中で報告された条件違いの複数モデルの結果です。

この表からわかるのは、特徴分布の均一性の指標と生成品質（gFID）の強い関係性です。

FAVAE と Diffusers ベース VAE の潜在画像特徴分布の均一性指標による比較

再構成品質は Diffusers ベース VAE の方が良いものの、ここで比較した 2 つの圧縮率（潜在 (4,32,32) と (4,64,64)）では、潜在画像の 3 つの指標すべてで FAVAE が上回りました。

基本となる Diffusers ベース VAE の検証

この検証が指し示す重要な点は 2 つあります。

• 通常の VAE では、潜在画像のチャンネル数が大きくなるほど潜在画像の質は悪化します。
• 通常の VAE では、同じ圧縮率であれば空間方向に潜在空間をとった方が、潜在画像の質の悪化が小さくなります。

複数の圧縮率における FAVAE の検証

より広い圧縮率でスイープすると、FAVAE の潜在画像の指標が（絶対値として）常に良いわけではありません。
ただし注目すべきは、通常の VAE が潜在空間の容量（チャンネル数）を増やすほど指標を悪化させるのに対し、FAVAE ではこの「容量増加に伴う悪化」という相関が現れない点です。つまり FAVAE は、容量を増やしても潜在画像の品質が劣化しにくいモデルだといえます。

まとめ

通常は潜在画像のチャンネル数（特徴次元）が大きくなるにつれて指標（density_cv・normalized_entropy・gini_coefficient）が悪化しますが、FAVAE はチャンネル数が大きくても潜在画像の品質が悪化しないモデルになりました。
ただし、ウェーブレット変換で詳細（高周波部分）と抽象（低周波部分）を完全に分離したとしても、それが潜在変数の分布密度の均一化・空間の有効活用・空間全体へのデータ分散に明確に寄与する、とまでは明言できません。
それでも、 $(3,256,256)\to (4,32,32)$ および $(3,256,256)\to (4,64,64)$ に圧縮するタスクにおいて、潜在画像の 3 つの指標では Diffusers ベース VAE より FAVAE の方が良い結果になっているといえます。

FAVAE では $(3,1024,1024)\to (4,128,128)$ に圧縮するタスクで、実際にサンプリングした再構成画像が目視でも悪くなく、潜在画像の品質も悪くありませんでした。
したがって、高解像度な画像を VAE で処理する際、潜在空間のキャパシティを十分に持たせることで、「多様で高品質なコンテンツの生成能力」と「元データへの高い忠実度（再構成能力）」をトレードオフにせず両立させた高性能モデルへと発展できる可能性があります。

現在ではさまざまな VAE が登場しています。特に、Vision Foundation Model などを使ってアライメントを行うモデルが性能を大きく向上させています。ここで、再構成についてのデータは論文などで確認できますが、潜在画像の構造や品質を把握することは困難でした。

潜在変数の均一性メトリクスの検証プロセスを行うことで、生成性能への悪影響を簡単に予測できると考えています。10 万枚ほど確かめても数分で終了できるため、価値は高いといえます。

---

付録: 特徴分布の均一性指標の定義

まず、手元にあるデータを次のように定義します。

• 全データ点数（サンプリングされた潜在ベクトルの数）: $n$
• 各データ点 $i$ における KDE の算出密度: $d_i(i=1,2,\dots ,n)$

ここでの密度とは、t-SNE によって 2 次元に圧縮された空間において、カーネル密度推定（KDE）を用いて算出された、各データ点における確率密度の値です。
与えられた $n$ 個の 2 次元座標の分布から、空間全体の連続的な確率密度関数 $f(\mathbf{x})$ を推測します。
2 次元空間上の任意の座標 $\mathbf{x}=[x,y{]}^T$ における確率密度関数 $f(\mathbf{x})$ は、以下で定義されます。

$$f(\mathbf{x})=\frac{1}{n\cdot h^2}\sum _{j=1}^{n}K\left(\frac{\mathbf{x}-{\mathbf{x}}_j}{h}\right)$$

カーネル関数 $K(\mathbf{u})$ には、標準 2 次元ガウス関数（平均 $\mathbf{0}$ 、分散共分散行列 $\mathbf{I}$ の正規分布の確率密度関数）を使用します。

$$K(\mathbf{u})=\frac{1}{2\pi }\exp (-\frac{1}{2}\Vert \mathbf{u}{\Vert }^2)$$

したがって、各データ点の密度 $d_i$ は次のように決まります。

$$d_i=f({\mathbf{x}}_i)=\frac{1}{2\pi \cdot n\cdot h^2}\sum _{j=1}^n\exp (-\frac{\Vert {\mathbf{x}}_i-{\mathbf{x}}_j{\Vert }^2}{2h^2})$$

ある点 ${\mathbf{x}}_i$ の近く（距離 $\Vert {\mathbf{x}}_i-{\mathbf{x}}_j\Vert$ が小さい範囲）に他のデータ点 ${\mathbf{x}}_j$ が多く存在するほど、指数関数 $\exp (\cdot )$ の値が 1 に近づき、それらの総和である $d_i$ （確率密度）の値は大きくなります。
逆に、周囲に他のデータ点が全く存在しない孤立した点である場合、距離が大きくなるため $\exp (\cdot )$ の値は 0 に近づき、 $d_i$ の値は極めて小さくなります。

density_cv（変動係数）

density_cv（Coefficient of Variation: 変動係数）は、「密度のばらつきの相対的な大きさ」を表す指標です。
密度の平均値 $\overline{d}$ と標準偏差 ${\sigma }_d$ を用いて、以下のように計算します。

$$\overline{d}=\frac{1}{n}\sum _{i=1}^n{d}_i$$

$${\sigma }_d=\sqrt{\frac{1}{n}\sum _{i=1}^n(d_i-\overline{d}{)}^2}$$

$$\text{density\_cv}=\frac{{\sigma }_d}{\overline{d}}$$

normalized_entropy（正規化エントロピー）

各データ点の密度 $d_i$ を、全データ点の密度の総和で割ることで、「合計が 1 になる確率 $p_i$ 」に変換します。

$$p_i=\frac{d_i}{\sum _{j=1}^n{d}_j}$$

次に、底を 2 とするエントロピー $H$ を計算します（ $ϵ$ はゼロ除算を防ぐ微小値 ${10}^{-10}$ ）。

$$H=-\sum _{i=1}^n{p}_i{\log }_2(p_i+ϵ)$$

エントロピーは、すべての確率が完全に等しい（ $p_i=1/n$ ）ときに最大値 $H_{max}={\log }_{2}n$ を取ります。これを用いて正規化します。

$$\text{normalized\_entropy}=\frac{H}{H_{max}}=\frac{-\sum _{i=1}^n{p}_i{\log }_2(p_i+ϵ)}{{\log }_{2}n}$$

• 完全均一なとき : すべての点が同じ密度（ $p_i=1/n$ ）のとき、分子のエントロピーは最大値 $H_{max}$ と一致するため、normalized_entropy は最大値の 1 になります。
• 偏りがあるとき : 特定の領域だけにデータが集中している（一部の $p_i$ が極端に大きく、他が 0 に近い）とき、エントロピーは小さくなり、normalized_entropy は 0 に近づきます。

gini_coefficient（ジニ係数）

gini_coefficient（ジニ係数）は、「データ密度が特定の点にどれだけ独占されているか」を評価するために使用します。

まず、各データ点の密度 $d_i$ を小さい順（昇順）にソートします。

$$d_{(1)}\le d_{(2)}\le \cdots \le d_{(n)}$$

ソートされた各データに順位（インデックス） $k=1,2,\dots ,n$ を割り振り、以下の数式でジニ係数を計算します。

$$\text{gini\_coefficient}=\frac{\sum _{k=1}^n(2k-n-1)d_{(k)}}{n\sum _{k=1}^n{d}_{(k)}}$$

• 順位 $k$ が低い（密度が低い）点に対しては、この係数は負の値になります。
• 順位 $k$ が高い（密度が高い）点に対しては、この係数は正の大きな値になります。
• ごく一部のデータ点だけが莫大な密度（高い $d_{(k)}$ ）を独占していると、正の大きな係数と掛け合わされるため、分子の合計値が非常に大きくなります。
• 完全均一なとき : すべてのデータ点の密度が全く同じ（ $d_{(k)}=\text{constant}$ ）場合、分子の $\sum (2k-n-1)$ の部分が綺麗に相殺して 0 になるため、gini_coefficient は 0 になります。
• 偏りがあるとき : ほとんどの場所がスカスカ（密度 0）で、ごく一握りの点だけが超高密度である場合、格差が最大となり gini_coefficient は 1 に近づきます。