近年の拡散モデルは高品質な画像生成が可能ですが、その計算コストは依然として大きく、特に高解像度データセットでは学習・生成の負担が課題になります。そこで今回は、MNISTとCIFAR-10の2つのデータセットを題材に、潜在空間（latent space）で拡散処理を行う軽量U-Net実装を試してみました。
潜在空間を使う理由
通常の拡散モデルはピクセル空間で直接ノイズ除去を学習しますが、潜在空間を経由することで次のようなメリットがあります。
- 計算量の大幅削減（解像度が下がるため演算コストが減る）
- 学習速度の向上（小さなU-Netで処理可能）
- モデルサイズの縮小（低リソース環境にも対応しやすい）
PCAによる潜在空間比較（VAEなし）
実験の第一歩として、まずはVAEを使わずに画像をPCAで二次元圧縮し、その潜在空間を直接DDPMで学習させてみました。
この方法では、元データ（例：MNIST）をPCAで2次元の潜在ベクトルに変換し、その分布を拡散モデルで再現します。
学習後、生成された潜在分布を元のPCA潜在と比較すると、実データはクラスごとに偏りや非対称性を持った複雑な形状をしている一方、DDPMの生成結果はほぼガウス分布に近く、全体的に中心へ集まった形になりました。
これはモデルの学習容量や学習ステップ数の影響が大きく、短時間学習では分布の細部まで再現するのが難しいことを示しています。
この比較により、潜在空間生成の基礎挙動をシンプルな形で確認でき、以降のVAE＋潜在拡散実験に向けた有用なベースラインとなりました。
実装概要
- VAEで画像を低次元の潜在ベクトルに圧縮
- 潜在空間で軽量U-NetベースのDDPMを学習
- 推論時は潜在から画像を復元
今回はMNIST（28×28, グレースケール）とCIFAR-10（32×32, RGB）の両方で試し、軽量U-Netを採用することで学習・生成速度を大きく向上できました。
結果と印象
- MNIST：数字の形状はある程度再現され、速度は非常に快適。ややぼけた印象はあるものの、全体的なクオリティは十分。こちらはdiffusersは使わず、pythonフルスクラッチに。
もう少しCNN版も提案されたのでやってみて多少改良されたかな？？
結構輪郭はっきりするなど改良が見れます。
- CIFAR-10：カラーデータでも高速に動作。小さなU-Netでも十分動き、試行回数を重ねてパラメータ調整する余裕がある。
- 以前試した「ピクセル空間で直接学習」した場合は輪郭がよりシャープでしたが、潜在空間方式の軽快さは魅力的。
CIFAR-10でもPCAで直接マッピングをしてみました。先ほどのMNISTよりはしっかりしてそう。
画像自体は雰囲気はわかるもののぼやっとしてます。
高性能にしたと思ったらいきなりノイズに。
改善したらちょっとは良くなりました。
VAE再構築も程々
まとめ
潜在空間を活用した軽量U-Netによる拡散生成は、クオリティと速度のバランスが良く、試行実験やプロトタイプ開発に最適だと感じました。次のステップとしては、条件付き生成やより高次元の潜在表現を試し、さらなる画質改善を目指したいと思います。
結構色々やってみましたが、画像が小さくてよくわからないですね。もちょい大きいのでチャレンジしてみます。