【要約】拡散モデルをMNISTで動かす：トイモデルから本物のU-Netへ [Zenn_Python] | Summary by TechDistill
> Source: Zenn_Python
Execute Primary Source
// Problem
開発者が、低次元のトイモデルから高次元の画像データを扱う際に、空間情報の欠落という課題に直面する。具体的には以下の問題が発生する。
- ・空間情報の喪失: MLPに画像を渡すと、ピクセル間の隣接関係が失われる。
- ・時刻情報の希薄化: 単一の数値として時刻tを渡すと、層が深くなるにつれ特徴が埋もれる。
- ・学習の不安定化: 高次元化に伴い、層を重ねるほど学習の進捗が困難になる。
// Approach
筆者は、画像の局所的な特徴を捉えつつ、時間経過を適切にモデルに伝えるために、以下の手法を導入した。
- ・U-Net構造の採用: エンコーダで特徴を抽出し、デコーダで解像度を復元する構造を実装。
- ・Skip Connection: 縮小過程で失われる位置情報をデコーダへ直接渡す。
- ・Time Embedding: 時刻tを複数の周波数の三角関数に展開し、各ResBlockへ注入。
- ・EMA (Exponential Moving Average): 重みの指数移動平均を保持し、生成品質を安定させる。
// Result
MNISTを用いた実験の結果、学習済みモデルはノイズから数字を生成する能力を示した。
- ・生成精度: 分類器による確信度0.9超えが約8割（本物は0.973）に達した。
- ・計算コスト: T4 GPUを使用し、6000ステップの学習に約10分を要した。
- ・生成特性: ラベル情報なし（Unconditional）でも、学習データの分布に基づいた多様な数字を生成可能。
Senior Engineer Insight
> 本実装は、拡散モデルの数学的本質とCNNの構造的利点を結びつける優れた教育的モデルである。実戦においては、Time Embeddingの設計やEMAの導入が生成品質の鍵となる。ただし、計算コストは次元数に比例して増大するため、スケーラビリティの確保にはさらなる最適化が不可欠である。