Micro Diffusion - 学習のための小さなDiffusionモデル

こんにちは。約1か月前にテキストDiffusionモデルをゼロから実装して公開しました。最近GeekNewsを見ていてふと思い出し、今になって投稿してみます。

---

作ることになったきっかけ

Andrej KarpathyのMicroGPTを見て、「こんなに短いコードでGPTの核心を説明できるんだな」という印象を受けました。ちょうどDiffusionも勉強したいと思っていたところだったので、両方を組み合わせて、「Diffusionも同じやり方で理解できるコードを作ったら面白そうだ」 という考えで始めた教育用プロジェクトです。

---

AR vs Diffusion: 何が違うのか

LLM時代のテキスト生成は、ほとんどすべてAutoregressive（AR）方式です。左から右へトークンを1つずつ予測していく方式です。

Discrete Diffusionはその逆です。シーケンス全体を一度に置き、ノイズ（マスキング）から段階的に復元します。

名前 "emma" を例にすると:  
  
Forward (学習 - ノイズ追加):  
  t=0 :  e m m a   ← 原本  
  t=25:  e _ m a   ← 一部をマスク  
  t=50:  _ _ m _   ← さらにマスク  
  t=100: _ _ _ _   ← 完全にマスク  
  
Reverse (生成 - ノイズ除去):  
  t=100: _ _ _ _   ← 空の状態から開始  
  t=75:  _ m _ _   ← 確信度の高い位置から復元  
  t=50:  e m _ a  
  t=0 :  e m m a   ← 完成  

ARが「単語を1文字ずつ書くこと」だとすれば、Diffusionは「クロスワードを解くこと」に近いです。

---

実装構成

3つのバージョンがあり、難易度に応じて選んで見ることができます。

• train_minimal.py — 2-layer MLP / NumPyのみ使用（最も単純）
• train_pure.py — 3-layer MLP + skip connection / NumPyのみ使用
• train.py — 4-layer Transformer / PyTorch使用

3つのバージョンすべてが 同じdiffusionループ を共有しています。異なるのはデノイザーのアーキテクチャだけです。
学習データは32,000件の英語名で、パラメータ数は170K〜239K程度です。

すぐに実行してみます:

pip install numpy  # 最小バージョン  
python train_minimal.py  

---

フィードバックやPR歓迎です！