T2I 拡散モデルの設計メモ
目次
- テキストエンコーダー
- VAE
- GAN
- コンディショニング
- 位置埋め込み
- 目的関数
- Transformer アーキテクチャ
- ブロック図
- データセット
- キャプショニング
- 学習方法
- 高速化
テキストエンコーダー
CLIP や T5 が長らく使われてきたが、軽量 LLM を使うケースが増えている。
- phi-3:OmniGen
- Qwen2.5-VL-3B-Instruct:OmniGen2
- Gemma 2:Lumina-Image 2.0, SANA
- Llama3-8B:Playground v3
主な LLM と VLM
性能比較

Params(M) は間違いで、数値は FP32 精度の時のファイルサイズ
出典:SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers. Enze Xie et al. Table 9. https://openreview.net/forum?id=N8Oj1XhtYZ
むしろ T5-Large がコスパがよい。fp16(性能劣化なし)だと 1.6GB、GGUF Q8_0 量子化すると 900MB になる。
T5-XXL を使うよりは軽量 LLM の方が速いし性能もよい。
T5 と LLM を混ぜて使う場合は注意が必要だ。T5 は LLM に比べて出力の分散が小さい。なので LLM の出力に RMSNorm などの Normalization レイヤーが必要になる。
参考文献
Unified Multimodal Understanding via Byte-Pair Visual Encoding。画像にも byte-pair encoding を適用する。
VAE
なぜ KL VAE を使うのか
潜在空間の統計的制御(規則性)
KL loss を入れることで、潜在空間に意味的なスムーズさや構造性が生まれる。以下のような特徴を持つ潜在空間は拡散モデルにとって扱いやすい。
- 同じ意味を持つ画像は近い z にエンコードされる
- ノイズを加えてもデコーダーが画像として再構成しやすい
- 正規分布に従うことで、サンプリング可能な「全体的に埋まった」潜在空間になる
VAE の圧縮率
VAE の圧縮率を決めるパラメータは2つある。
- F:解像度の縮小率
- C:潜在空間のチャンネル数
拡散モデルの隠れ層の次元はさらにパッチサイズが関係している。
- P:パッチサイズ
VAE の隠れ層の次元数は以下のようになる。
\[\scriptsize{H \times W \times 3 \rightarrow \dfrac{H}{F} \times \dfrac{W}{F} \times C}\]拡散モデルの隠れ層の次元は以下のようになる。
\[\scriptsize{H \times W \times 3 \rightarrow \dfrac{H}{PF} \times \dfrac{W}{PF} \times (C \times P^2)}\]各 VAE の設定
モデル | F | C | P |
---|---|---|---|
SD1.5 | F32 | C4 | P2 |
SANA | F32 | C32 | P1 |
SDXL | F8 | C4 | P2 |
Dit-Air | F8 | C8 | P2 |
SD3・FLUX.1 | F8 | C16 | P2 |
Dit-Air の F8C8P2 = F16C32P1 が画質と処理効率とのバランスがいい。SANA は画質が悪すぎるし、SD3・FLUX は処理が重すぎる。
Transformer の隠れ層の次元 >> C なので、既存の VAE の P を大きくすれば、トークン数が減るので拡散モデルの処理速度を向上させられる。
SD-VAE・FLUX.1 VAE のレイヤー情報
Notes / Links about Stable Diffusion VAE
SDXL の VAE のレイヤー情報は sdxl-vae/config.json を参照。
SD3 の VAE や FLUX.1 の VAE にアクセスするには認証が必要。
Diffusers で様々な VAE の実装が見れる。
SD や FLUX は Diffusers の AutoencoderKL を使用しているので以下のコードで簡単にモデルを構築できる。学習するコードも 100 行程度あれば可能。block_out_channels は最終的には [128, 256, 512, 512] を使うことが多い。
AutoencoderKL を定義するコード
from diffusers.models import AutoencoderKL
F = 8
C = 16
resolution = 1024
model = AutoencoderKL(
sample_size=resolution,
in_channels=3,
out_channels=3,
latent_channels=C,
block_out_channels=[
int(resolution/F),
int(resolution/4),
int(resolution/2),
int(resolution/2)
],
layers_per_block=2,
down_block_types=[
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D",
"DownEncoderBlock2D"
],
up_block_types=[
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D",
"UpDecoderBlock2D"
],
act_fn="silu",
scaling_factor=0.18215,
mid_block_add_attention=True
)
VAE を学習させるコード
# train
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
model.train()
for i, (x, labels) in enumerate(train_loader):
x = x.to(device)
# encode
posterior = model.encode(x).latent_dist
latent_representation = posterior.sample()
# decode
# tanh で出力の範囲を [-1, 1] に制限
reconstructed_x = torch.tanh(model.decode(latent_representation).sample)
# loss
reconstruction_loss = F.mse_loss(reconstructed_x, x, reduction="sum")
kl_loss = torch.mean(posterior.kl())
vae_loss = reconstruction_loss + kl_loss
# back propagation
optimizer.zero_grad()
vae_loss.backward()
optimizer.step()
パッチサイズは小さい方が拡散モデルの性能が良い
パッチサイズは VAE の性能とは関係がない。パッチ処理があると拡散モデルが余計な処理をしないといけないので、拡散モデルの性能がわずかに低下する。
ただし MicroDiT はパッチマスクを行う前に Patch-mixer でテキスト Embedding を取り込むことで性能を向上させている。

パッチサイズは小さい方が拡散モデルの性能が良い
出典:SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers. Enze Xie et al. Figure 3. https://openreview.net/forum?id=N8Oj1XhtYZ
D2iT
D2iT は圧縮率の異なる2つの潜在空間表現にエンコード可能な VAE を使う。ディティールの細かい部分は低圧縮率、情報の少ない部分は高圧縮率と使い分ける。デノイザもディティールの細かい部分は多くのネットワークが通過する。
DC-AE Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
2025 年以降に発表された DiT でよく使われる。SANA や Nitro-T で使われている。space-to-channel, channel-to-space(要はパッチとアンパッチ)を学習可能な残差接続と、3段階の学習で性能を向上させる。
- 低解像度でエンコーダーとデコーダーを学習させる
- 高解像度でエンコーダーの終わりとデコーダーの入り口とのみを学習させる
- 高解像度でデコーダーの終わりのみを学習させる
DC-AE は dc-ae-f32c32-sana-1.1-diffusers からダウンロードできる。モデルの定義は efficientvit。
細部が大きく崩れるので、等倍では使い物にならない。Detailer の使用がほぼ必須で、顔なら解像度は 1024 x 1024 は欲しい。
VAE の学習
VAE の学習は簡単ではない。VAE 単体ではぼやけた画像にしかならない。
Stable DiffusionのVAEは、LPIPS、パッチベースのGAN損失、およびMSE損失を含む複雑な損失関数を使用して訓練される。この組み合わせにより、VAEは、ピクセル単位で正確 (MSE)、知覚的に整合 (LPIPS)、かつ全体的にリアル (GAN損失) な、高品質でぼやけのない再構成を生成できる。
taming-transformers は LPIPS と GAN とを使用して VQ-VAE を学習させており、VAE の作成の時に参考になる。
VIVAT: Virtuous Improving VAE Training through Artifact Mitigation は VAE の学習でよく起こるアーティファクトの対処法を解説している。
posterior collapse
VAE は同じ画像ばかり出力する posterior collapse(KL loss が急速に0に近づく)が起こりやすい。そこで以下のような戦略が良く使われる。
- KL Loss の逆伝播の量を調整する kl_beta を設定する。
- 最初の数万ステップは kl_beta を0にする。
- 次にウォームアップ区間を設ける。 kl_beta が 0→1 になる“ウォームアップ区間”は、実装・データ規模により 1 万〜10 万 SGD ステップ、あるいは 50〜200 エポック程度
- 各ピクセルの小さい loss を、ノイズとみなして、フィルターして無視する
- 各ピクセルの KL loss の最大値を設定する。max(KLi, threshold)
- ウォームアップ区間が短すぎると posterior collapse が起こりやすく、長すぎると訓練が停滞するため、近年では「1〜2% の総ステップ数をウォームアップに割く」か「 ELBO が安定するまでβを上げない」設計が主流
- 最低でも1万ステップのウォームアップ期間が必要。高解像度画像・映像用 VAE では 2.5〜5 万ステップ程度の長めのウォームアップが採用される傾向があり、テキスト VAE では 1〜2 万ステップ前後が多い
- posterior.mean を可視化して、すべて 0 近辺になっている場合、posterior collapse の疑いが強い
正規化
画像の表示時や保存時に逆正規化を忘れるのはよくやりがちだ。
transforms.Normalize([0.5], [0.5]) # 訓練画像の正規化 x = (x * 0.5 + 0.5).clamp(0, 1) # 逆正規化
MSE 単体では目的関数として力不足
The reasonable ineffectiveness of pixel metrics for future prediction (and what to do about it)
- ぼやけた画像は MSE が小さいことが多い。ぼやけた画像は位置の微妙な変化にも対応できる
- 最尤推定は訓練データとは著しく異なるデータを生成してもペナルティがない
正則化項
- KL。KL VAE はオートエンコーダーに KL 正則化項を追加したもの
- LPIPS
- GAN
- Masked AutoEncoders
- EQ-VAE
- Watsonの知覚モデル。離散フーリエ変換 (DFT) 等を使い特徴を抽出。LPIPS より6倍高速で、省メモリ
- 輪郭情報。Unleashing the Power of One-Step Diffusion based Image Super-Resolution via a Large-Scale Diffusion Discriminator
- そのほかは画像比較手法まとめを参照
LPIPS
LPIPS の損失の計算をするコード
import lpips
# LPIPS の準備('vgg' or 'alex')
lpips_fn = lpips.LPIPS(net='vgg').to(device)
lpips_fn.eval() # 推論モード
for epoch in range(num_epochs):
for i, (x, labels) in enumerate(train_loader):
x = x.to(device)
# encode
posterior = model.encode(x).latent_dist
latent_representation = posterior.sample().float()
# decode
# [-1, 1] に正規化
reconstructed_x = torch.tanh(model.decode(latent_representation).sample)
# loss: MSE
reconstruction_loss = F.mse_loss(reconstructed_x, x, reduction="sum")
# loss: KL
kl_loss = torch.mean(posterior.kl())
# loss: LPIPS
# LPIPS に入力する画像は [-1, 1] に正規化されている必要がある
lpips_loss = lpips_fn(x_norm, reconstructed_x).mean()
# AMP を使う場合は lpips の計算中は AMP を無効にする
# with autocast('cuda', enabled=False):
# lpips_loss = lpips_fn(x.float(), reconstructed_x.float()).mean()
# 総合 loss
vae_loss = reconstruction_loss + kl_loss + lpips_lambda * lpips_loss
optimizer.zero_grad()
vae_loss.backward()
optimizer.step()
lpips の動作は以下のコードで検証できる。> 0.1 なら lpips は正常。
x_noise = x + torch.randn_like(x) * 0.1
x_noise = torch.clamp(x_noise, 0.0, 1.0)
x_norm = x * 2 - 1
x_noise_norm = x_noise * 2 - 1
lpips_test = lpips_fn(x_norm, x_noise_norm)
EQ-VAE: Equivariance Regularized Latent Space for Improved Generative Image Modeling
EQ-VAE は SD-VAE が同変性(equivariance)を持たないことに注目し、スケール・回転変換を使った正則化で高画質を実現した。入力を変換したものと、潜在空間表現を変換したものをデコードしたものとの再構成 loss をとる。
以下の式で \(\scriptsize \tau\) はスケール・回転変換、\(\scriptsize \mathcal{E}\) はエンコーダー。\(\scriptsize \mathcal{L}_{rec}\) は再構成 loss。\(\scriptsize \mathcal{L}_{reg}\) は KL loss。エンコーダーで変換する画像は、スケール・回転変換しないことに注意。
\[ \small \mathcal{L}_{EQ-VAE}(x, \tau) = \mathcal{L}_{rec}(\tau \circ \mathbf{x}, \mathcal{D}(\tau \circ \mathcal{E}(\mathbf{x}))) + \lambda_{gan}\mathcal{L}_{gan}(\mathcal{D}(\tau \circ \mathcal{E}(\mathbf{x}))) + \lambda _{reg}\mathcal{L}_{reg} \]潜在空間表現の回転と拡大縮小するコードは以下のようになる。潜在空間表現は3次元 [C, H, W] である必要がある。
import random
# 普通はスケールしてから回転させるが、このコードでは回転させてからスケールしている
def rot_scale(x, angle, size):
# x.shape = [B, C, H, W]
x_s = torch.nn.functional.rotate(latent, angle)
return torch.nn.functional.interpolate(
x_s,
size=size,
mode='bicubic', # 'nearest', 'bilinear', 'bicubic' などが選択可能
align_corners=False # デフォルトはFalse。Trueにすると端のピクセルのアライメントが変わる
)
angle = 90 * random.randint(0, 3) # angle = 0, 90, 180, 270
size = (random.uniform(0.25, 1.0), random.uniform(0.25, 1.0)) # sx, sy = 0.25~1.0
x_rs = rot_scale(x, angle, size)
latent_rs = rot_scale(encoder(x), angle, size)
reconstructed_x = decoder(latent_rs)
reconstruction_loss = F.mse_loss(reconstructed_x, x_rs, reduction="sum")
Latent Diffusion Models with Masked AutoEncoders
Masked Autoencoders に KL 項を追加して、潜在空間のスムースさ・知覚的圧縮品質・再構成画像の品質の3つの指標を同時に達成。Loss は以下の4つ。
- マスクしていない部分の再構成 loss
- マスク部分の再構成 loss
- マスクしていない部分の LPIPS loss
- マスクしていない部分の KL loss
対称形の ViT ベースのエンコーダーとデコーダーとを採用している。SD-VAE がファイルサイズ 320MB に対して VMAE は 43MB で、性能で SD-VAE を上回る。
F4C16 のモデルの処理は以下のようになる。Transformer の次元は dmodel とする。Transformer は Transformerエンコーダの単一の層(Multi-Head Self-AttentionとFeed-Forward Network) のみの torch.nn.TransformerEncoderLayer が便利。
- 元画像(3 x 512 x 512)
- 画像のパッチ化によるダウンサンプル(3 x 512 x 512)->(64 x 128 x 128)-> (16,384 x 64)
- Linear で隠れ層の次元を調整 (16,384 x 64) ->(16,384 x dmodel)
- Transformer
- Linear 等で \(\scriptsize{\mu}\) と \(\scriptsize{\sigma}\) とを計算(それぞれ(16,384 x 16))
- z をサンプリング(z.shape =(16,384 x 16))
- Linear(16,384 x 16)->(16,384 x dmodel)
- Transformer
- Linear で隠れ層の次元を調整 (16,384 x dmodel) -> (16,384 x 64)
- アンパッチ (16,384 x 64) -> (64 x 128 x 128)->(3 x 512 x 512)
Transformer VAE の設計
Transformer VAE は出力がブロック状になりやすい。原因は、Transformer のパッチ化と KL loss との2つ。
KL loss 対策
- ステップ数が十分か確かめる。ブロック感をなくすには 50 万ステップは必要
- β-VAE を採用し、ベータを抑え気味にする
パッチ化対策
- 学習可能な位置エンコーディングを使う
- VIVAT: Virtuous Improving VAE Training through Artifact Mitigation
- トランスフォーマーブロックは最低でも8は必要
- パッチ化に畳み込みを使いオーバーラップさせる
- デコーダーの最後に refiner の畳み込みを追加する
オーバーラップのサンプルコード
self.image_to_token_cnn = nn.Conv2d(
in_channels=3,
out_channels=model_dim,
kernel_size=patch_size,
stride=patch_size//2, # stride < kernel_size → オーバーラップあり
padding=patch_size//4
)
Refiner のサンプルコード
self.reconstruction_refiner = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 3, kernel_size=3, padding=1)
)
# after rearrange
img = self.reconstruction_refiner(img)
参考文献
Self-Guided Masked Autoencoder
Masked Autoencoders Are Scalable Vision Learners
MaskGIT: Masked Generative Image Transformer
MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis
Latent Denoising Makes Good Visual Tokenizers
MVAE のように入力画像をマスクした上に、潜在空間表現にノイズを乗せてデコーダーを学習させる。MVAE の loss に加えて GAN も使う。
補間ノイズ \(\scriptsize x' = (1-\tau)x + \tau \epsilon\) と加算ノイズ \(\scriptsize x' = x + \tau \epsilon\) とでは補間ノイズの方が性能がよかった。\(\scriptsize \epsilon(\gamma) \sim \gamma \cdot \mathcal{N}(\mathbf{0},\mathbf{I})\) で \(\scriptsize \gamma = 3\) が最も性能がよかった。可算ノイズはオリジナルのシグナルを損なわないようなショートカットが作成される恐れがある。それにも関わらず VAR では性能が改善されたが、SiT では性能は改善されなかった。
マスク率は 70~90% が性能が高い。マスク率をランダム化した方が性能が良い。
マスクと潜在空間ノイズとを併用すると VAR では性能が上がったが SiT では性能が上がらなかった。SiT の場合は、マスクはしなくてもいい。
VAE 学習の参考文献
Community Training AutoencoderKL #894
Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing (pdf)
Summary: Beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework
beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework
Taming Transformers for High-Resolution Image Synthesis
The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
ALVAE Adversarial Latent Autoencoders
AS-VAE Adversarial Symmetric Variational Autoencoder
A Loss Function for Generative Neural Networks Based on Watson's Perceptual Model
参考文献
Conditional VAE Semi-Supervised Learning with Deep Generative Models
Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
VQ-VAE
VQ-VAE Neural Discrete Representation Learning
Generating Diverse High-Fidelity Images with VQ-VAE-2
SoftVQ-VAE: Efficient 1-Dimensional Continuous Tokenizer
Quantize-then-Rectify: Efficient VQ-VAE Training
Instella-T2I: Pushing the Limits of 1D Discrete Latent Space Image Generation。1次元バイナリ潜在空間に圧縮することで VQ-VAE 比で 32 倍の圧縮率を達成している。
VQ-Diffusion Vector Quantized Diffusion Model for Text-to-Image Synthesis。VQ-VAE は量子化されているので、小さなミスが大きな意味の変化へつながる恐れがある。VQ-Diffusion では Mask-and-replace という学習手法でその問題に対処している。
GAN
GAN は拡散モデルの主要技術ではないが、VAE の学習や蒸留のときに必要になる。
画像生成より真贋判定の方がタスクが簡単なので、Discriminator のネットワークはシンプルで浅いものを使うのが普通。Discriminator が強すぎると Generator の学習が進まなくなる。
Patch-based adversal loss を利用する Image-to-Image Translation with Conditional Adversarial Networks が基礎的な論文で、畳み込みを使う DCGAN でネットワークを作る。
tips
- GAN はバッチサイズを大きくしてはいけない。勾配が平均化されて生成される画像の多様性がなくなる。大きくても 32 程度にしておくのがよい。
- 学習率は 2e-4 以下。Generator より Discriminator の学習率を低くすることもある
- DCGAN は Adam の β1 を 0.5 にすることを推奨している
- Generator の入力に dropout や noise を追加する
- 潜在ベクトル z に情報を持たせる(例:カテゴリや style)ことでモード崩壊しにくくなる
- 定期的にランダムな潜在ベクトルで多様性を可視化して監視する
モード崩壊(Mode Collapse)
モード崩壊(Mode Collapse)とは Generator が同じ画像ばかり出力すること。以下の対策方法がある。
学習戦略・最適化の工夫
- Minibatch Discrimination。複数のサンプル間の類似度を導入して、生成器が同じようなサンプルを出すことを抑制
- One-sided label smoothing。本物のラベル(1.0)を 0.9 にするなどして、識別器を強くしすぎない
- ノイズの追加(Input noise / instance noise)。入力画像にランダムノイズを加えて、識別器の判別を難しくしすぎない
- 識別器を強くしすぎない。学習率の調整、早めの停止など。Discriminator が強すぎると Generator が collapse しやすい。
ロス関数の変更
- WGAN。JS divergence の代わりに Earth Mover 距離を使うことで、学習が安定
- WGAN-GP。WGAN の改良。勾配クリッピングの代わりに勾配ノルムのペナルティ。
- LSGAN。二乗誤差を使うことで、識別器が飽和しにくくなる
- Relativistic GAN。「偽物が本物よりもリアルかどうか」を評価する相対的な損失関数を用いる
モデル構造の工夫
- モード追跡型の Generator(e.g. VAE-GAN, InfoGAN)。
- InfoGAN: 潜在変数に意味を持たせることで多様性を向上
- VAE-GAN: 潜在空間の連続性・多様性を保証
- 複数の Generator を使う(e.g. MGAN, MAD-GAN)。Generator を複数用意し、それぞれ異なるモードを学習させる
- Progressive Growing(PGGAN)。低解像度から徐々に高解像度にしていくことで、学習の安定性向上
正則化や多様性の強化
- Mode-seeking Regularization(Mode-seeking Loss)。潜在変数の違いに応じて出力が大きく変わるようにする正則化項
- Feature Matching。Discriminator の中間特徴を一致させるような loss を導入し、モード崩壊を抑制
- Diversity Sensitive Loss。潜在ベクトルの距離と出力の違いを関連付けるような損失項を導入
そのほか
- Unrolled GAN。Discriminator のパラメータ更新を数ステップ先まで「展開(unroll)」して Generator を学習
- PacGAN。識別器に複数の生成画像(例:2個ずつ)を同時に入力することで、多様性がないことを見抜けるようにする
- Mini-batch Standard Deviation Layer(StyleGANなどで使用)。Generator に統計量(標準偏差)を追加して多様性を強化
- スペクトル正規化
- Mixture of GANs(Discriminator や Generator を複数使う)
- Diversity Promoting Regularization
導入順
- WGAN-GP ・ LSGAN
- Minibatch Discrimination ・ Feature Matching
- InfoGAN ・ Mode-seeking loss
- Unrolled GAN ・ PacGAN
LSGAN
JS Divergenceの最小化に基づくGANでは、勾配消失が起きやすい。LSGAN はGANの損失関数を最小二乗誤差にすることで、この問題の解決する。WGAN より実装が単純なので、まずは LSGAN から試すのがよい。
LSGAN(Least Square Generative Adversarial Networks)を試してみた
Wasserstein GAN
Wasserstein 距離に基づいて loss を計算する。モード崩壊が軽減され、学習の安定化と品質の改善が期待できる。通常 Discriminator は入力が本物である確率を計算するが、wgan では Discriminator(wgan では Critic と呼ぶ)は Wasserstein 距離の近似値を計算する。Critic は入力された画像がどれだけ現実離れしているかを数値化する。
詳細な解説は今さら聞けないGAN(4) WGAN を参照。通常は Grgdient penalty を採用する。
Improved Training of Wasserstein GANs
Weight Clip を使用する古いコード例
import torch
import torch.nn as nn
import torch.optim as optim
# 損失関数の定義 (Weight Clipping を使用)
def wgan_loss(critic, generator, real_images, fake_images, clip_value):
"""
WGAN の Critic と Generator の損失を計算する関数。
Args:
critic: Critic モデル。
generator: Generator モデル。
real_images: 本物の画像データ (Tensor)。
fake_images: Generator が生成した偽の画像データ (Tensor)。
clip_value: Weight Clipping に使用する値。
Returns:
tuple: (Critic loss, Generator loss) のタプル。
"""
# Critic の損失計算
critic_real_output = critic(real_images)
critic_fake_output = critic(fake_images.detach()) # detach() で勾配計算を停止させる
critic_loss = torch.mean(critic_real_output - critic_fake_output) # Wasserstein distance の近似値を最大化
# Generator の損失計算
generator_output = critic(fake_images)
generator_loss = torch.mean(-generator_output) # Critic の評価を下げる (Wasserstein distance を最小化)
return critic_loss, generator_loss
# 簡易的な Critic と Generator の定義 (例)
class Critic(nn.Module):
def __init__(self):
super(Critic, self).__init__()
self.model = nn.Sequential(
nn.Linear(784, 256), # 例えば MNIST サイズの場合
nn.ReLU(),
nn.Linear(256, 1)
)
def forward(self, x):
x = x.view(x.size(0), -1) # Flatten the image
return self.model(x)
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(100, 256), # ノイズ次元は例として100
nn.ReLU(),
nn.Linear(256, 784) # MNIST サイズの場合
)
def forward(self, z):
z = z.view(z.size(0), -1) # Flatten the noise vector
return self.model(z).view(-1, 28, 28) # Reshape to image size
# ハイパーパラメータの設定 (例)
latent_dim = 100
batch_size = 64
learning_rate_critic = 0.0001
learning_rate_generator = 0.0001
clip_value = 0.01 # Weight Clipping の値
# モデルとオプティマイザの作成
critic = Critic()
generator = Generator()
optimizer_critic = optim.Adam(critic.parameters(), lr=learning_rate_critic)
optimizer_generator = optim.Adam(generator.parameters(), lr=learning_rate_generator)
# ダミーデータ生成 (テスト用)
real_data = torch.randn(batch_size, 1, 28, 28) # MNIST サイズのダミーデータ
noise = torch.randn(batch_size, latent_dim)
# 損失計算と学習ループの一部
for epoch in range(10): # 例として10エポック
optimizer_critic.zero_grad()
optimizer_generator.zero_grad()
fake_data = generator(noise)
critic_loss, generator_loss = wgan_loss(critic, generator, real_data, fake_data, clip_value)
# Critic の学習 (勾配の計算と適用)
critic_loss.backward()
# Weight Clipping を適用 (重要!)
for p in critic.parameters():
torch.nn.utils.clip_grad_value_(p, -clip_value, clip_value) # 重要な制約
optimizer_critic.step()
# Generator の学習 (勾配の計算と適用)
generator_loss.backward()
optimizer_generator.step()
print(f"Epoch {epoch+1}: Critic Loss = {critic_loss:.4f}, Generator Loss = {generator_loss:.4f}")
Gradient penalty を使用するコード例
- lambda_gp: Gradient Penalty の係数。この値を調整することで、Lipschitz 制約への適合度合いを制御する。 大きすぎる値は学習を不安定にする可能性がある
- 線形補間: real_images と fake_images の間で線形補間を行い、Critic が入力データ空間全体で滑らかであるかを評価する
- 学習率: Gradient Penalty を使用した場合、従来の WGAN よりも高い学習率を使用する
import torch
import torch.nn as nn
import torch.optim as optim
# 損失関数の定義 (Gradient Penalty を使用)
def wgan_loss_gp(critic, generator, real_images, fake_images, device, lambda_gp=10):
"""
WGAN と Gradient Penalty の損失を計算する関数。
Args:
critic: Critic モデル。
generator: Generator モデル。
real_images: 本物の画像データ (Tensor)。
fake_images: Generator が生成した偽の画像データ (Tensor)。
device: 計算デバイス (CPU または GPU)。
lambda_gp: Gradient Penalty の係数。
Returns:
tuple: (Critic loss, Generator loss) のタプル。
"""
# Critic の損失計算
critic_real_output = critic(real_images)
critic_fake_output = critic(fake_images.detach()) # detach() で勾配計算を停止させる
critic_loss = torch.mean(critic_real_output - critic_fake_output)
# Gradient Penalty の計算
with torch.no_grad(): # 重要な部分: 勾配の計算をしない
duals = real_images + (torch.rand_like(real_images, device=device) - 1) * (fake_images - real_images).sum(dim=(1,2,3), keepdim=True) # 線形補間
critic_dual_output = critic(duals)
gradient_penalty = torch.mean((critic_dual_output - torch.zeros_like(critic_dual_output)).pow(2)) * lambda_gp
critic_loss += gradient_penalty
# Generator の損失計算
generator_output = critic(fake_images)
generator_loss = torch.mean(-generator_output)
return critic_loss, generator_loss
# 簡易的な Critic と Generator の定義 (例)
class Critic(nn.Module):
def __init__(self):
super(Critic, self).__init__()
self.model = nn.Sequential(
nn.Linear(784, 256), # 例えば MNIST サイズの場合
nn.LeakyReLU(0.2), # ReLU の代わりに Leaky ReLU を使用するのが一般的
nn.Linear(256, 256),
nn.LeakyReLU(0.2),
nn.Linear(256, 1)
)
def forward(self, x):
x = x.view(x.size(0), -1) # Flatten the image
return self.model(x)
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(100, 256), # ノイズ次元は例として100
nn.ReLU(),
nn.Linear(256, 784) # MNIST サイズの場合
)
def forward(self, z):
z = z.view(z.size(0), -1) # Flatten the noise vector
return self.model(z).view(-1, 28, 28) # Reshape to image size
# ハイパーパラメータの設定 (例)
latent_dim = 100
batch_size = 64
learning_rate_critic = 0.0001
learning_rate_generator = 0.0001
lambda_gp = 10 # Gradient Penalty の係数
# モデルとオプティマイザの作成
critic = Critic()
generator = Generator()
optimizer_critic = optim.Adam(critic.parameters(), lr=learning_rate_critic)
optimizer_generator = optim.Adam(generator.parameters(), lr=learning_rate_generator)
# ダミーデータ生成 (テスト用)
real_data = torch.randn(batch_size, 1, 28, 28).cuda() # MNIST サイズのダミーデータ
noise = torch.randn(batch_size, latent_dim).cuda() # GPU を使用する場合
# 計算デバイスの設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
critic.to(device)
generator.to(device)
# 損失計算と学習ループの一部
for epoch in range(10): # 例として10エポック
optimizer_critic.zero_grad()
optimizer_generator.zero_grad()
fake_data = generator(noise)
critic_loss, generator_loss = wgan_loss_gp(critic, generator, real_data, fake_data, device, lambda_gp)
# Critic の学習 (勾配の計算と適用)
critic_loss.backward()
optimizer_critic.step()
# Generator の学習 (勾配の計算と適用)
generator_loss.backward()
optimizer_generator.step()
print(f"Epoch {epoch+1}: Critic Loss = {critic_loss:.4f}, Generator Loss = {generator_loss:.4f}")
StyleGAN
現在の画像生成 GAN の主流は StyleGAN で、StyleGAN によって学習の安定化と生成画像の操作が可能になった。
StyleGAN2 Analyzing and Improving the Image Quality of StyleGAN
StyleGAN3 Alias-Free Generative Adversarial Networks
Real-ESRGAN
LPIPS(VGG)loss と GAN loss と Content Loss とを損失関数として採用している。
Content Loss
\[ \begin{split} \normalsize L_{1} &= E_{x_{i}}||G(x_{i})-y||_{1} \\ \scriptsize{G(x_{i})} &\scriptsize{: 生成画像} \\ \scriptsize{y} &\scriptsize{: 本物画像} \end{split} \]pytorch-CycleGAN-and-pix2pix
安定化
Spectral Normalization (スペクトル正規化)
Discriminatorの各層の重みの Lipschitz 定数を制約する(要は Normalize する)ことで学習が安定する。Spectrally Normalized Generative Adversarial Networks (SN-GAN)。
Label Smoothing (ラベルスムージング)
Discriminatorのターゲットラベルを0または1のようなバイナリ値ではなく、0.9や0.1のような滑らかな値(小数値)に置き換える。Discriminatorが訓練データに対して過度に自信を持つことを防ぎ、過学習を抑制。これにより、Generatorに与えられる勾配信号がより適切になり、敵対的サンプルの生成リスクを低減する効果も期待できる。What is the intuition behind the Label Smoothing in GANs? - AI Stack Exchange
one_labels = torch.full((BATCH_SIZE,1),fill_value=0.9).to(device)
zero_labels = torch.full((BATCH_SIZE,1),fill_value=0.1).to(device)
One-sided Label Smoothing (片側ラベルスムージング)
従来のラベルスムージングとは異なり、Discriminatorにおいて「本物」の画像に対するラベルのみをスムージングし、「偽物」の画像に対するラベルはそのまま(通常0)にする。Label Smoothing でも Discriminator が強い場合に使う。Label smoothing should be one-sided? · Issue #10 · soumith/ganhacks - GitHub。Application of Smoothing Labels to Alleviate Overconfident of the GAN's Discriminator - IIETA
one_labels = torch.full((BATCH_SIZE,1),fill_value=0.9).to(device)
zero_labels = torch.zeros((BATCH_SIZE, 1)).to(device)
コード例
GAN の構築自体は容易だ。ResNet や VGG の最終層を取り換えて使えば、Discriminator として使える。VAE の学習ならば、Generator は VAE。GAN のシンプルな実装は【Pytorch】MNISTのGAN(敵対的生成ネットワーク)を実装する等を参照。ResNet を Discriminator に使うコードは以下のようになる。
import torch
import torch.nn as nn
import torchvision.models as models
# 訓練済みResNet-18をロード
resnet_model = models.resnet18(pretrained=True)
# 既存の最終層の入力特徴量サイズを取得
# ResNet-18の場合、fc.in_featuresは 512
num_ftrs = resnet_model.fc.in_features
# 新しい最終層を定義
# 本物/偽物の2クラス分類のため、出力は1
resnet_model.fc = nn.Linear(num_ftrs, 1)
# Sigmoidアクティベーションは、通常GANのDiscriminatorの出力では直接適用せず、
# BCEWithLogitsLossなど、ロジット値を直接受け取る損失関数と組み合わせます。
# もし0-1の確率として明示的に出力したい場合は、以下のように追加します。
# resnet_model.fc = nn.Sequential(
# nn.Linear(num_ftrs, 1),
# nn.Sigmoid()
# )
訓練コード
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# --- 1. Discriminator (ResNet) の定義 ---
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
# 訓練済みResNet-18をロード
resnet = models.resnet18(pretrained=True)
# 最終層をGANの Discriminator 用に変更
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 1) # 出力は1(ロジット)
self.model = resnet
def forward(self, x):
return self.model(x)
# --- 2. Generator のダミー定義 (簡単のため) ---
# 実際にはもっと複雑なネットワークになります
class Generator(nn.Module):
def __init__(self, latent_dim, img_channels, img_size):
super(Generator, self).__init__()
self.img_size = img_size
self.main = nn.Sequential(
nn.Linear(latent_dim, 256 * (img_size // 16) ** 2), # 適当な中間層
nn.LeakyReLU(0.2, inplace=True),
nn.Unflatten(1, (256, img_size // 16, img_size // 16)),
nn.ConvTranspose2d(256, 128, 4, 2, 1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2, inplace=True),
nn.ConvTranspose2d(128, 64, 4, 2, 1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
nn.ConvTranspose2d(64, img_channels, 4, 2, 1),
nn.Tanh() # 画像出力のためTanh
)
def forward(self, x):
return self.main(x)
# --- 3. ハイパーパラメータとデバイス設定 ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
latent_dim = 100 # Generatorへの入力ノイズの次元
image_size = 64 # 生成/入力画像のサイズ
batch_size = 64
num_epochs = 50
lr = 0.0002
# --- 4. データローダー (例としてCIFAR-10) ---
transform = transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
# 訓練データセットの準備
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# --- 5. モデル、オプティマイザ、損失関数の初期化 ---
discriminator = Discriminator().to(device)
generator = Generator(latent_dim, 3, image_size).to(device) # CIFAR-10は3チャンネル
# オプティマイザ
optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.999))
# 損失関数 (Sigmoidと交差エントロピーを組み合わせたもの)
criterion = nn.BCEWithLogitsLoss()
# 真偽ラベル
real_label = 1.0
fake_label = 0.0
# --- 6. 訓練ループ ---
print("Starting Training Loop...")
for epoch in range(num_epochs):
for i, data in enumerate(train_loader, 0):
# --- Discriminatorの訓練 ---
discriminator.zero_grad()
real_images = data[0].to(device)
b_size = real_images.size(0)
# 1. 本物画像を Discriminator に入力
label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
output = discriminator(real_images).view(-1)
errD_real = criterion(output, label)
errD_real.backward()
D_x = torch.sigmoid(output).mean().item() # Sigmoidで0-1に変換して表示
# 2. 偽物画像を Generator で生成し、Discriminator に入力
noise = torch.randn(b_size, latent_dim, device=device)
fake_images = generator(noise)
label.fill_(fake_label)
output = discriminator(fake_images.detach()).view(-1) # detach()でGeneratorへの勾配計算を停止
errD_fake = criterion(output, label)
errD_fake.backward()
D_G_z1 = torch.sigmoid(output).mean().item()
errD = errD_real + errD_fake
optimizer_D.step()
# --- Generatorの訓練 ---
generator.zero_grad()
label.fill_(real_label) # Generatorは偽物画像を本物とDに誤認させたい
output = discriminator(fake_images).view(-1)
errG = criterion(output, label)
errG.backward()
D_G_z2 = torch.sigmoid(output).mean().item()
optimizer_G.step()
if i % 100 == 0:
print(f'[{epoch}/{num_epochs}][{i}/{len(train_loader)}] '
f'Loss_D: {errD.item():.4f} Loss_G: {errG.item():.4f} '
f'D(x): {D_x:.4f} D(G(z)): {D_G_z1:.4f} / {D_G_z2:.4f}')
# エポックごとにモデルを保存するなどの処理を追加することもできます
# if (epoch + 1) % 10 == 0:
# torch.save(generator.state_dict(), f'generator_epoch_{epoch+1}.pth')
# torch.save(discriminator.state_dict(), f'discriminator_epoch_{epoch+1}.pth')
print("Training finished!")
コンディショニング
拡散トランスフォーマーではコンディショニングとして、タイムステップやテキストをトークン化せずに入力する。
しかしコンディショニングをトークンにして入力する方法が 2025 年現在のトレンドだ。OmniGen はコンディショニングをすべてトークンにして入力している。 しかし OmniGen2 では Time step をトークン化しない方法を選択している。

出典:Qi Qin et al. Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Figure 3.
https://arxiv.org/abs/2503.21758
図の位置の改変は筆者
コンディショニングの方法はいくつかある。
- トークン化してノイズやプロンプトと結合(OmniGen)
- Cross Attention(PixArt-α, Hunyuan-DiT, SANA)
- zero-Initialized gated cross-attention(Lumina-Image 2.0)
- Mod(Stable Diffusion 3)
- Scale, adaLN-Zero(DiT, Lumina-Image 2.0)
しかしテキストのコンディショニングはどのモデルも、トークン化してノイズ画像と結合するか、クロスアテンションで取り込むかしている。DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation によると、クロスアテンションよりノイズ画像と結合する方がパラメータ効率が良い。
結合
モデルの隠れ層の次元が dmodel で、テキストエンコーダーの出力 T が [nt, dt] で、潜在空間の画像の次元 IMG が [ni, di] とする。
これらを結合するには Linear などで、T と IMG の隠れ層の次元を dmodel に変換する必要がある。
t_hidden = self.fc_t(T) # [n_t, d_t] -> [n_t, d_model]
img_hidden = self.fc_img(IMG) # [n_i, d_i] -> [n_i, d_model]
input = torch.cat([img_hidden, t_hidden], dim=0) # 実際はバッチサイズがあるので dim=1 になる
print(input.shape) # [n_t + n_i, d_model]
# 分離
img = [:n_i, :] # [n_i, d_model]
t = [n_i:, :] # [n_t, d_model]
Mod
\[x_{out} = x_{in} \odot (1 + \gamma) + \beta\]adaLN
Layer Nomalization のガンマとベータは学習可能なパラメータだ。adaLN はこれを線形変換などを使用して動的に変更するもの。つまり
c = タイムステップなど γ = Linear(c) β = Linear(c)\[ \displaylines{ LN(x) = \gamma \odot \dfrac{x - \mu}{\sigma} + \beta \\ \scriptsize{Layer Normalization の数式} } \]\[ \displaylines{ AdaLN(x, c) = \gamma (c) \odot \dfrac{x - \mu}{\sigma} + \beta (c) \\ \scriptsize{AdaLN の数式。c はプロンプトやタイムステップ等} } \]
ブロック図での表記法
- Scale & Shift:Scale & Shift とだけ書かれていても、AdaLN の場合がある
- LayerNorm の後に Scale, Shift
- LayerNorm の後に Mod
Unveiling the Secret of AdaLN-Zero in Diffusion Transformer
ゼロではなくガウス分布を利用した初期化をする adaLN-Gaussian を提唱している。
位置埋め込み
SD3 以前は Sinusoidal Encoding 使われてきたが、FLUX 以降は RoPE を使うのが一般的。ViT の RoPE については Rotary Position Embedding for Vision Transformer が詳しい。
Timestep を Sinusoidal Encoding するコード
import torch
import math
def get_timestep_embedding(
timesteps: torch.Tensor,
embedding_dim: int,
flip_sin_to_cos: bool = False,
downscale_freq_shift: float = 1,
scale: float = 1,
max_period: int = 10000,
) -> torch.Tensor:
"""
This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
Args:
timesteps (torch.Tensor): a 1-D Tensor of N indices, one per batch element
embedding_dim (int): the dimension of the output
flip_sin_to_cos (bool): Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
downscale_freq_shift (float): Controls the delta between frequencies between dimensions
scale (float): Scaling factor applied to the embeddings
max_period (int): Controls the maximum frequency of the embeddings
"""
assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
half_dim = embedding_dim // 2
exponent = -math.log(max_period) * torch.arange(
start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
)
exponent = exponent / (half_dim - downscale_freq_shift)
emb = torch.exp(exponent)
emb = timesteps[:, None].float() * emb[None, :]
# scale embeddings
emb = scale * emb
# concat sine and cosine embeddings
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
# flip sine and cosine embeddings
if flip_sin_to_cos:
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
# zero pad
if embedding_dim % 2 == 1:
emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
return emb
SANA は RoPE ではなく Mix-FFN を採用している。
RoFormer: Enhanced Transformer with Rotary Position Embedding
Rotary Position Embedding for Vision Transformer
mRoPE(Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution)
Mix-FFN
3x3 の畳み込みによって RoPE なしでも相対位置情報を取得できる。SegFormer の Mix-FFN は以下のようになる。
xout = MLP(GELU(Conv3x3(MLP(xin)))) + xin
SANA の Mix-FFN は GELU ではなく SiLU を使っている。SiLU(Swish) は ReLU に形が似ていて、(0, 0) を通る。0を境界に導関数が急に変化する ReLU と違い、SiLU は導関数が滑らかに変化する。
GELU も SiLU も ReLU のような形で導関数が滑らかに変化するように設計されたもの。GELU の方が性能が良いが、SiLU の方が計算が速く、どちらも ReLU より性能が良い。
SANA が Mix-FFN を使用しているが、その隠れ層の次元が 5,600 もあり、RoPE を使った方が速い可能性が高い。
目的関数
ノイズを予測させる epsilon prediction から、ノイズ差分を予測させる v prediction へトレンドが移り、2025 年では Rectified Flow や Flow matching を使うのが一般的になっている。
今後は、拡散モデルとフローマッチングと自己回帰モデルとを一般化した Transition Matching が主流になるかもしれない。
Transition Matching
DTM(Difference Transition Matching)がフローマッチングに相当する。目的関数の部分はフローマッチングとほとんど同じ。
エンコーダー・デコーダーモデル。backbone と呼ばれるエンコーダーが巨大。小規模なデコーダーは head と呼ばれている。
参考文献
Improving and generalizing flow-based generative models with minibatch optimal transport
Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow
SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers
Transition Matching: Scalable and Flexible Generative Modeling
Transformer アーキテクチャ
Transformer のトレンドは以下のように変化した
- sinusoidal positional embedding -> RoPE
- Multi-Head Attention -> Grouped-Query Attention
- GELU -> SwiGLU
- LayerNorm -> RMSNorm
- QK-Norm
Grouped-Query Attention
Grouped-Query Attention は Multi-Head Attention の K, V を複数のクエリで共有する。 GQA-4 のケースでは、K と V との数が 1/4 になり、推論速度が 1.3 倍高速になり、メモリ使用量 50% 減少、性能の低下は1~3%。
Dit-Air ではヘッドだけではなく、QKVO をすべてのトランスフォーマーブロックで共有している。
Multi-Head Latent Attention
DeepSeek-V2 で採用された、KV キャッシュを圧縮する手法。
QK-Norm
Multi-Head Attention の Q, K の RoPE 適用前に RMSNorm を入れる。
参考文献
The Big LLM Architecture Comparison
ブロック図
トランスフォーマーブロックの内部に以下のパターンが2~3回出現するのが基本。
- AdaLN
- Attention・MLP ・FFN
- Scale
DiT
FFN
Pointwise Feedforward Network と Position-wise Feed-Forward Network とはほぼ同じ意味。バッチ・系列長・埋め込み次元のテンソルにおいて、各位置 i にあるベクトル x_i に対して、まったく同じ重みの MLP を適用する。数式は以下のようになる。
\[FFN(x) = max(0, xW_1 + b_1)W_2 + b_2\]Patching(Patchify)
VAE の設定が F8C4P2 とする。元の画像が 3 x 1024 x 1024 だとする。その潜在空間表現は 4 x 128 x 128。パッチサイズが2ということは4ピクセル(2 x 2)を一つの埋め込み次元にするので、16 x 64 x 64。16 = チャンネル数 * パッチサイズ^2。幅と高さを1次元化して転置すると、最終的に 4096 x 16 になる。埋め込み次元は 16 になる。
その後で Linear によって任意の埋め込み次元数に拡張する。
from einops import rearrange
# x = [B, C, H, W] -> [B, H*W, C*P*P]
patches = rearrange(x, 'b c (h p) (w p) -> b (h w) (c p p)', p=self.patch_size)
# [B, H*W, C*P*P] -> [B, H*W, D_model]
hidden_states = self.projection(patches)
# [B, H*W, D_model] -> [B, H*W, C*P*P]
patches = self.reverse_proj(hidden_states)
rearrange(patches, 'b (h w) (c f1 f2) -> b c (h f1) (w f2)', f1=P, f2=P, h=int(H/P), w=int(W/P))
# 畳み込みの場合
image_to_patches = nn.Conv2d(input_channel, model_dim, kernel_size=(P, P), stride=P, bias=True)
patches = image_to_patches(x)
tokens = rearrange(patches, 'b c h w -> b (h w) c')
# unpatch
patches_to_image_deconv = nn.ConvTranspose2d(model_dim, input_channel, kernel_size=(P, P), stride=P, bias=True)
patches = rearrange(tokens, 'b (h w) c -> b c h w', h=int(H/P), w=int(W/P))
patches_to_image_deconv(patches)

2022年
出典:Scalable Diffusion Models with Transformers. William Peebles et al. Figure 3. https://arxiv.org/abs/2212.09748
PixArt-α
位置エンコーディングには DiffFit: Unlocking Transferability of Large Diffusion Models via Simple Parameter-Efficient Fine-Tuning で使用された、 sinusoidal interpolation を使っている。実装は単純で、DiT の解像度は 256 だが DiffFit は 512 なので、DiT の sinusoidal encoding の (i, j) を (i/2, j/2) にしただけだ。
MLP
トランスフォーマーブロックの隠れ層と同じ次元を持つ、2レイヤーの MLP。活性化関数は SiLU。

2023年
出典:PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis. Junsong Chen et al. Figure 4. https://arxiv.org/abs/2310.00426
Stable Diffusion 3
テキストと画像とを別々に処理するダブルストリーム方式。ダブルストリームはパラメータ効率が悪く、現在では選択肢に入らない。
Stable Diffusion 3(SD3)の理論やモデル構造について解説

2024年
出典:Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. Patrick Esser et al. Figure 2. https://stability.ai/news/stable-diffusion-3-research-paper
Playground v3
発表は 2024/09/16、設計の終了と学習の開始は Stable Diffusion 3(2024/02/22)の論文の発表前。
EDM スケジューラーを使っている。

2024年
出典:Playground v3: Improving Text-to-Image Alignment with Deep-Fusion Large Language Models. Bingchen Liu et al. Figure 2. https://arxiv.org/abs/2409.10695
FLUX.1
AuraFlow
Hunyuan-DiT

2024年
出典:Hunyuan-DiT: A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding. Zhimin Li et al. Figure 7. https://arxiv.org/abs/2405.08748
MicroDiT
パッチマスクはパッチのランダムドロップより性能が良いのでよく使われる(BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, Masked Autoencoders Are Scalable Vision Learners)。Fast Training of Diffusion Models with Masked Transformers ではマスクパッチの再構成能力を上げるためにオートエンコーダー loss を利用している。これらの従来手法はパッチのマスク率が 50% を超えると性能が低下し始め、75% を超えると大幅に性能が低下する。しかし MicroDiT は 75% のマスク率でもわずかな性能低下に抑えられる。
MicroDiT はパッチマスクを行う前に Patch-mixer でテキスト Embedding を取り込むことで性能を向上させた。Patach-mixer はアテンションと FFN で構成された数レイヤーの Transformer。

2024年
出典:Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget Vikash Sehwag et al. Figure 2. https://arxiv.org/abs/2407.15811
パッチマスクはブロックサイズを小さくしてランダムにした方が性能低下を抑えられる。
クロスアテンションでキャプションを取り込み、位置エンコーディングとタイムステップは Sinusoidal Embedding。
損失関数。マスクした部分は、元の画像+ノイズを予測させる。
\[ \small \begin{split} \mathcal{L}_{diff}&=\mathbb{E}_{(\mathbf{x},\mathbf{c})\sim\mathcal{D}}\mathbb{E}_{\mathbf{\epsilon}\sim\mathcal{N}(\mathbf{0},\sigma(t)^2\mathbf{I})}\left|\left|\big(\bar{F}_\theta((\mathbf{x}+\mathbf{\epsilon})\odot(1-m);\sigma,\mathbf{c})-\mathbf{x}\big)\odot(1-m)\right|\right|^2_2\\ \mathcal{L}_{mae}&=\mathbb{E}_{(\mathbf{x},\mathbf{c})\sim\mathcal{D}}\mathbb{E}_{\mathbf{\epsilon}\sim\mathcal{N}(\mathbf{0},\sigma(t)^2\mathbf{I})}\left|\left|\big(\bar{F}_\theta((\mathbf{x}+\mathbf{\epsilon})\odot(1-m);\sigma,\mathbf{c})-(\mathbf{x}+\mathbf{\epsilon})\big)\odot m\right|\right|^2_2\\ \mathcal{L} &= \mathcal{L}_{diff} + \gamma\mathcal{L}_{mae} \end{split} \]1.16B のモデルの学習に $1,890 のコストしかかかってない。

2024年
出典:Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget Vikash Sehwag et al. Figure 3. https://arxiv.org/abs/2407.15811
AI 生成画像を学習に使う
学習画像の 40% が JourneyDB や DiffusionDB の画像データ。これらの画像を使ってもベンチのスコアは変わらないが、主観的な品質は明らかに向上する。
そこで ChatGPT にどちらの画像が質が高いか質問した。プロンプトは「Which image do you prefer, Image A or Image B, considering factors like image details, quality, realism, and aesthetics? Respond with 'A' or 'B' or 'none' if neither is preferred.」DrawBench と PartiPrompts で生成したプロンプトで、AI 生成画像を学習に含めたものと含めなかったものとで画像を生成させて ChatGPT に評価させたところ、AI 生成画像を学習に含めたモデルの生成した画像が圧勝している。

出典:Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget Vikash Sehwag et al. Figure 9. https://arxiv.org/abs/2407.15811
Lumina-Image 2.0

2025年
Element-wise Addition と Element-wise Multiplication は記号が逆
出典:Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Qi Qin et al. Figure 2. https://arxiv.org/abs/2503.21758

Image Text Attention は右の FFN と等価で、キャプションの量を増やせばネットワークのパラメータ数を増やすのと同じ効果がある
出典:Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Qi Qin et al. Figure 5. https://arxiv.org/abs/2503.21758
RMSNorm
RMSNorm は Layer Normalization の計算コストを削減したもの。具体的には各層の活性化に対して、平均の計算を省略し、二乗平均平方根 (Root Mean Square: RMS) のみを用いて正規化を行う。RMSNorm は LayerNorm と同等の性能を達成しつつ、計算時間を 7% 〜 64% 削減できる。
RMSNorm の実装例
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float=1e-5):
"""
基本的な役割はLayerNormと同じだが計算量が7%から64%少ない。ただし大規模なネットワークでは効果は小さい
Args:
dim (int): 入力次元
eps (float): 0 除算エラーを避けるバイアス
"""
super().__init__()
self.eps = eps
self.dim = dim
self.weight = nn.Parameter(torch.ones(dim))
self.bias = nn.Parameter(torch.zeros(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True)) + self.eps
x = x / rms
return self.weight * x + self.bias
ゲート
Gate は以下の式で表現される。ただし Normalization の後にゲートを入れる場合は、b1, b2 は、出力平均が0に近いため省略されることが多い。
\[ \displaylines{ Gate(\mathbf{x}) = (\mathbf{x}W_1 + b_1) \odot \sigma(\mathbf{x}W_2 + b_2)\\ \begin{split} \scriptsize{\odot} &\scriptsize{: アダマール積(要素ごとの積)}\\ \scriptsize{\sigma} &\scriptsize{: 非線形活性化関数(sigmoid, GELU など)} \end{split} } \]nn.Lienar は bias=False でバイアスを省略できる。
nn.Linear(in_features, out_features, bias=False) # バイアスなし
GEGLU の実装例
import torch
import torch.nn as nn
import torch.nn.functional as F
class GEGLU(nn.Module):
def __init__(self, input_dim, hidden_dim):
super().__init__()
self.proj = nn.Linear(input_dim, hidden_dim * 2)
def forward(self, x):
x_proj = self.proj(x) # [B, T, 2*hidden_dim]
x1, x2 = x_proj.chunk(2, dim=-1) # 分割
return x1 * F.gelu(x2) # Gated
# 使用例
batch_size, seq_len, input_dim, hidden_dim = 32, 128, 512, 2048
x = torch.randn(batch_size, seq_len, input_dim)
geglu = GEGLU(input_dim, hidden_dim)
out = geglu(x) # -> [B, T, hidden_dim]
HiDream-I1
- パッチ使用
- Playground v3 のように LLM の中間出力を取り込んでいく
- FLUX.1 のようにデュアルストリームの後にシングルストリームになる
- SD3 と同じ3つのテキストエンコーダーに LLM を加え、合計4つのテキストエンコーダーを使用
- 位置エンコーディングは RoPE
- 拡散モデルでは珍しく MoE を採用している

2025年
出典:HiDream-I1: A High-Efficient Image Generative Foundation Model with Sparse Diffusion Transformer. Qi Cai et al. Figure 3. https://arxiv.org/abs/2505.22705
SwiGLU
ゲートの活性化関数に Swish を使ったもの。SwiGLU は入力の二乗や乗算のような多項式近似をエミュレートできるのが強み。LLM では事実上のスタンダードになっている。表現力の向上、収束の高速化(学習の高速化)、大規模モデルでの性能向上が見られることが複数の論文で実証されているが、性能向上の理由は不明。
従来の FFN は2つの重み行列(W1, W2)を持つのに対し、GLU 派生は3つ(W, V, W2)持つ。なので SwiGLU を含む GLU の派生形のパラメータ数は、FFN と比較して、隠れ層の次元が 2/3 に削減される。
SwiGLU は外れ値同士が掛け算されて値が増幅されることがある。Smooth-SwiGLU によってその問題は解決され、LLM では FP8 での学習も可能になっている。
ReLUと比較して、SwiGLUは固有のスパース性がほとんどないため、従来のスパース性を活用する技術の効果が薄い。
DiT においては収束を 21 倍加速させた事例が報告されているが、性能が向上した論文はない。
参考文献
SANA

2025年
出典:SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers. Enze Xie et al. Figure 5. https://openreview.net/forum?id=N8Oj1XhtYZ

2025年
出典:SANA 1.5: Efficient Scaling of Training-Time and Inference-Time Compute in Linear Diffusion Transformer. Enze Xie et al. Figure 10. https://arxiv.org/abs/2501.18427
Linear Attention
アテンションは (QK)V の順で計算するが Linear Attention は Q(KV) の順で計算する。(QK) = n x n だが、(KV) = dmodel * dmodel 。通常のアテンションはトークン数が n とすると計算量とメモリ使用量とが O(n2) になるが、Linear Attention では O(n) になる。SD3 なら画像サイズが 1024 だと n は 4,096。
Linear Attention は高速だが性能は低下する。その低下分を Mix-FFN でカバーしている。Mix-FFN は RoPE を排除できるが、隠れ層の次元が dmodel = 2,240、dFFN = 5,600 と FFN の隠れ層の次元が巨大になっている。
畳み込み
3x3 の前の 1x1 の畳み込みによって計算量を削減している。1x1 の畳み込みはチャンネル数の次元を削減するのに使われることが多い。例えば (C, H, W) = (3, 8, 8) に 1x1 の畳み込みを実行して (1, 8, 8) にする。その後で計算コストの高い 3x3 の畳み込みを実行し、1x1 の畳み込みで (3, 8, 8) へチェンネル数を復元する。これはピクセル単位の MLP を実行しているとみることもできる。
畳み込みの出力次元数
\[ \scriptsize \begin{split} OH &= \dfrac{H + 2*padding - filter size}{stride} + 1\\ OW &= \dfrac{W + 2*padding - filter size}{stride} + 1 \end{split} \]Exploring 1x1 Convolutions in Deep Learning
パラメータ
dmodel = 2,240、dFFN = 5,600。
4.8B は 60 レイヤー。1.6B は 20 レイヤー。モデルサイズで隠れ層の次元は変わらない。
DiT-Air
ネットワークは Lumina-Image 2.0 とほぼ同じ。
DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation によると、テキストコンディショニングはクロスアテンションで取り込むのではなくノイズ画像と結合し、Normalize に AdaLN を使うのが最強らしい。
AdaLN のパラメータをレイヤー間で共有した場合、性能を変えずに 2.7 億個のパラメータを削減できる。アテンションレイヤー(QKVO)をレイヤー間で共有した場合、すべてのベンチでわずかに性能が低下するが、9,000 万パラメータの削減が可能。
結論としては AdaLN とアテンションレイヤーのウェイトを共有して、MLP のみ独立させるのが最もパラメータ効率が良い。
DiT-Air/L-Lite は総パラメータ数 1.15B で内訳は、CLIP/H(335M), DiT(700M), 8ch VAE(84M)。
隠れ層の次元は 64 * レイヤー数。
モデル | レイヤー数 | 隠れ層の次元 |
---|---|---|
S | 12 | 768 |
B | 18 | 1,152 |
L | 24 | 1,536 |
XL | 30 | 1,920 |
XXL | 38 | 2,432 |
マルチヘッドアテンションの数はトランスフォーマーの深さと同じ。
Pixart と同じ MLP なら、トランスフォーマーブロックの隠れ層と同じ次元を持つ、2レイヤーの MLP。活性化関数は SiLU。

2025年
出典:DiT-Air: Revisiting the Efficiency of Diffusion Model Architecture Design in Text to Image Generation. Chen Chen et al. Figure 4. https://arxiv.org/abs/2503.10618
これは、DiT-Air/XXLでは、できる限り最も優れたモデルを提供したいと考えているためです。この場合、特に複雑なプロンプトや長文の文字表現においては、CLIP単独の場合よりも、CLIPとLLMを組み合わせることで、さらに良い性能が得られます。
— Chen Chen (@alex8937) March 14, 2025
Nitro-T
0.6B は PixArt-α に似たアーキテクチャ、1.2B は SD3 に似たアーキテクチャで、ブロック図はない。
MicroDiT で使用された Deferred Patch Masking を採用しパッチ化の性能が上がっている。
使用技術的に SANA に近い。SANA はパッチを使用しないが Nitro-T はパッチを使うところが相違点。
- テキストエンコーダーに軽量 LLM(Llama 3.2 1B)を使用
- テキストエンコーダーの出力に RMSNorm を通すことで loss を安定させている。
- VAE に 32 倍の圧縮率の DC-AE を使用
35M 枚数の学習画像の過半数が AI 生成(DiffusionDB と JourneyDB)の画像。
データセット
データとパラメータ数はセットで考える。パラメータ数だけ増やしても性能は上がらない。データに多様性がなければデータを増やしても性能は上がらない。
最近では多くても 50M 程度の教師画像枚数で学習させることが多い。仕上げのファインチューニングは1~3M枚。
MicroDiT では、生成した画像の品質評価をするために ChatGPT にどちらの画像が質が高いか質問した。プロンプトは「Which image do you prefer, Image A(first) or Image B(second), considering factors like image details, quality, realism, and aesthetics? Respond with 'A' or 'B' or 'none' if neither is preferred.」これは、ローカルの Gemma3 でも実行できる。意味があるかどうかはわからないが「Tell me the reason.」で理由も聞ける。データセットの画像の厳選にも使える。
Sub-Scaling Laws: On the Role of Data Density and Training Strategies in LLMs
LLM ではモデルのパラメータ数が増えると性能向上率が悪化する。その原因はデータの多様性の不足と、学習のさせすぎだ。
論文ではデータをクラスタ分けしてデータの多様性を計測する方法や、Over-Training Ratio の概念を紹介している。
NoHumansRequired: Autonomous High-Quality Image Editing Triplet Mining
画像・修正指示・修正後の画像のデータセット、の公開と、データセットの作成方法についての解説。
著作権について
平成30年著作権法改正により学習は合法。ただしモデルが著作権侵害コンテンツを高頻度で生成する場合、モデルの公開・販売や画像生成サービスの提供は複製権侵害になる。
便利ツール
データセット
データセットは画像 URL とキャプションとのペアで提供される。画像のライセンスは画像の所有者ごとに異なる。以下のリストのライセンスが書いていないデータセットの画像のライセンスは画像所有者依存。
商用利用可能な大規模なデータセットは flicker ぐらいしかない。flicker は CC0 で画像を提供すれば無限のストレージが利用できる。
よく使われるデータセット
ImageNet:商用不可
Segment Anything 1B(Meta):Apache 2.0 License だが商用利用不可
AI 生成画像(多様性があり学習に使うと性能が上がる)
JourneyDB:商用不可。商用利用したければプロンプトから自作すればよい。$100万を超える売り上げのある企業は有料プランに加入しないと資産を所持できない。Midjourney は著作権侵害コンテンツを平気で生成してくるので、プロンプトと生成した画像との厳選は必要になる。
CC0
soa-full-florence2 はスミソニアン協会が公開している情報をもとにmadebyollin氏が作成したCC-0の絵画などを集めた画像リンク集。古いものが多く、メインのデータセットにはならない。
flicker
megalith-10m(flicker 画像のキュレーション)
次点
Open Images Dataset V7(Google):画像は CC BY 2.0 だが個々の画像のライセンスには Google は関知しない、キャプションは CC BY 4.0
DOCCI Descriptions of Connected and Contrasting Images
Unsplash:Unsplash ライセンス(商用可能、転売は不可)
LAION-5B:低品質な画像が多く使われなくなってきている
データの水増し
OmniGen2 は動画から画像を切り出して、学習に使用している。
Magpie: Alignment Data Synthesis from Scratch by Prompting Aligned LLMs with Nothing。LLM のケースだが、AI にデータを作成させる方法の紹介。
データの分類
OmniGen2 ではデータを以下のように分類し、均等になるようにしている。
- キャラ
- 男
- 女
- 子供
- 老人
- アニメキャラ
- オブジェクト
- バーチャル
- 食べ物
- 人工物
- 植物
- 動物
- シーン
- コマーシャル
- ランドマーク
- 地方(Domestic)
- 都市(Urban)
- 自然(Natural)
キャプショニング
データセットが提供するキャプションを使うのではなく、LLM で画像のキャプション付けを行うのが一般化している。ローカルで実行する場合はパラメータ数 70B のような高性能なものを使用することが多い。VILA-3B/13B や InternVL2-8B/26B もよく使われる。
解像度を下げる場合は、キャプションを作り直す必要はない。しかし画像をクロップした場合は、キャプションを作り直す。
LLM で以下の情報を抽出する。単に「車」と記述するのではなく、「通りを走る車」のようにアクションを記述させる。
- 人
- 人数
- 人の表情・感情
- ポーズ・姿勢:例:「走っている子供」、「ベンチに座っている人」
- 身体属性:例:「長い金髪の女性」、「ふわふわの毛の犬」
- 年齢・身長・性別
- 相互作用:例:「握手している二人」、「ボールで遊んでいる犬」
- 物
- 物の画面内での位置
- オブジェクト間の相対的位置
- 物色、パーツの色
- 形状
- 材質
- サイズ
- 時間帯
- 照明条件:例:「明るい昼光」、「薄暗い夕方の光」
- 季節
- 天候
- 背景の詳細: 例:「都市の景観」、「穏やかな森」
- 服の種類
- 動物の種類
- 室内か屋外か
- 全体的な雰囲気/ムード:例:「賑やかな通り」、「平和な庭」
- 文字の位置と文字
キャプションが複数ある場合は、キャプションをクリップスコアで評価し、正規化することで確率的にキャプションを選択する手法がある。
\[P(c_i) = \frac{exp(c_i/\tau)}{\sum_{j=1}^N exp(c_j/\tau)} \]c は CLIP スコアで、t は温度。0に近い値にすると常に最も高いスコアのキャプションが選択される。
長すぎるプロンプト問題
DetailMaster: Can Your Text-to-Image Model Handle Long Prompts? では長すぎるプロンプトが性能を低下させることを報告している。
Collage Prompting
従量課金の MLLM は画像枚数あたりで課金されるため、Collage Prompting では複数の画像を1枚にまとめてキャプションを作成する方法を紹介している。
プロンプト戦略
戦略カテゴリ | 説明 | プロンプト |
---|---|---|
直接キャプショニング | 画像全体を包括的に記述する。MLLMのゼロショット能力を活用。 | 「この画像を詳細に記述してください。」 「画像の内容を包括的な段落で説明してください。」 |
きめ細かい指示生成(Q&A形式) | 特定のオブジェクト、属性、関係性について質問と回答のペアを生成させる | 「画像内のオブジェクトの色、サイズ、素材、形状について質問と回答のペアを生成してください。」 「画像内の人物の行動と、彼らが他のオブジェクトとどのように相互作用しているかについて質問と回答を生成してください。」 |
属性ベースの拡張 | 既存の簡潔なキャプションに、きめ細かい属性(色、サイズなど)や空間的関係を追加して詳細化させる。 | 「このキャプションを、画像内のすべてのオブジェクトの具体的な色、サイズ、素材、形状、およびそれらの正確な空間的配置を含めて拡張してください。」 |
シーン・文脈の豊かさ | 画像の全体的な雰囲気、照明、時間帯、背景の詳細などを記述させる。 | 「この画像の全体的なムード、時間帯、照明条件、および背景の環境を詳細に記述してください。」 |
人物・動物のインタラクション強調 | 画像内の人物や動物の行動、感情、相互作用に焦点を当てて記述させる。 | 「画像内の人物の表情、ポーズ、および彼らが互いに、または周囲のオブジェクトとどのように相互作用しているかを詳しく説明してください。」 |
T2Iモデル向け最適化 | T2Iモデルの訓練データ分布に合わせた構造や表現でキャプションを生成させる。 | 「Text-to-Imageモデルのプロンプトとして最適化された形式で、この画像を記述してください。タグ形式または詳細な段落形式で、可能な限り多くの視覚的詳細を含めてください。」 |
コスト効率化(コラージュプロンプティング) | 複数の画像を1つの入力として処理し、それぞれの詳細を記述させる。 | 「提供されたコラージュ内の各画像を個別に識別し、それぞれについて詳細なキャプションを生成してください。」 |
参考文献
Unleashing Text-to-Image Diffusion Prior for Zero-Shot Image Captioning
To See is to Believe: Prompting GPT-4V for Better Visual Instruction Tuning
DetailMaster: Can Your Text-to-Image Model Handle Long Prompts?
学習方法
低解像度で学習した後、高解像度の学習を行い、最後に高品質な画像だけを使った仕上げを行う。
8倍圧縮の VAE の場合は 256 x 256 の解像度で事前学習を行うが、32 倍圧縮の VAE を使う SANA や Nitro-T では 512 x 512 から事前学習を始めている。
Scaling Laws For Diffusion Transformers によると、事前学習の loss もスケーリング則に従う。スケーリング則から、最適なモデルサイズと必要なデータセットの量が計算できる。
Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits
エポックごとにモデルのコピーを保存しておいて、EMA を使ってモデルの平滑化を行う。たいていは学習の最終段階のファインチューンの時にのみ行われる。
サンプルコードは Model EMA (Exponential Moving Average) を参照。
On the Scalability of Diffusion-based Text-to-Image Generation
隠れ層の次元を増やすよりトランスフォーマーブロックを増やした方が性能が上がる。データセットは質と多様性が重要で、量は重要ではない(現状十分な量の画像が手に入る)。
Classifier Free Guidance
一定確率(10%前後)でキャプションなしで学習させることで、CFG が使えるようになる。CFG は品質を大きく向上させるが、蒸留モデルでは使えないことが多い。蒸留モデルのみを公開する場合は CFG 学習をしない選択もある。
学習率スケジューラー
CosineAnnealingLR のような学習率が上下するスケジューラーは局所的最適解を避けられる。
from torch.optim import lr_scheduler
# T_max: 学習率が最小値に達するまでのエポック数。総エポック数の 1/5 程度が目安
# eta_min: 学習率の最小値。0に設定するコード例が多いが、計算資源の無駄でしかない
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-5)
# optimizer.step() の代わりに scheduler.step() を呼び出す
scheduler.step()
タイムステップスケジューラー
タイムステップは通常 0~1,000 の間でランダムに選択される。
A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Training によると、タイムステップの後半は Loss が小さく学習が遅いので、前半(250 前後)に集中して学習させることで3倍学習が高速になる。
ノイズスケジューラー
Stable Diffusion 1.5 や XL はスケジューラーにバグがあり、色の再現性の低下や高解像度での分裂などの問題を引き起こしていた。Stable Diffusion 3 ではタイムステップを以下の式で変換することで高解像度での分裂等に対処している。色の再現性は ZTSNR や Rectified Flow を採用することで改善されている。
\[ \displaylines{ \begin{split} t_m &= \dfrac{\sqrt{\frac{m}{n}} t_n}{1 + (\sqrt{\frac{m}{n}}-1)t_n}\\ \scriptsize{t_m} &\scriptsize{= 変換後のタイムステップ}\\ \scriptsize{t_n} &\scriptsize{= 変換前のタイムステップ}\\ \scriptsize{n} &\scriptsize{= ベース画素数。SD3 は 1,024 \times 1,024}\\ \scriptsize{m} &\scriptsize{= 生成する画像の画素数} \end{split} } \]
\(\small{x_t = \sqrt{\gamma}\cdot x_0 + \sqrt{1-\gamma}\cdot \epsilon : \gamma = 0.7}\)
高解像度の画像の情報を破壊するにはより多くのノイズが必要になる
出典:On the Importance of Noise Scheduling for Diffusion Models. Ting Chen. Figure 2
Common Diffusion Noise Schedules and Sample Steps are Flawed では、T=1000 まで学習すべきなのに T=999 で止まるノイズスケジューラーのバグを指摘している。
ZTSNR(Zero Terminal SNR)
Common Diffusion Noise Schedules and Sample Steps are Flawed では、T=1000 まで学習すべきなのに T=999 で止まるノイズスケジューラーのバグを指摘している。
Input Perturbation Reduces Exposure Bias in Diffusion Models
学習時、ネットワークは実際の画像にノイズを乗せたものを使って学習するが、推論時には自分の推論結果からノイズを徐々に除去していく。これは exposure bias 問題に似ていて、エラーが蓄積していく。
そこで学習時に追いノイズによってこのバイアスを除去する。以下の式の \(\scriptsize \sqrt{1+\gamma^2}\) が追いノイズの部分。
\[ \small \begin{split} \mathbf{y} &= \sqrt{\bar{\alpha}_t \mathbf{x}_0} + \sqrt{1-\bar{\alpha}_t}\sqrt{1+\gamma^2}\mathbf{\epsilon}\\ \gamma &= 0.1 \end{split} \]A Comprehensive Review on Noise Control of Diffusion Model
以下のノイズスケジューラーのレビュー。
- Linear
- Fibonacci
- Cosine
- Sigmoid
- Exponential
- Cauchy distribution
- Laplace distribution
- Logistic
- Monotonic neural network
高解像度では Cosine より Sigmoid の方が良い。
Improved Noise Schedule for Diffusion Training
ノイズスケジューラ―の Laplace, Cauchy は logSNR=0(つまり純粋なノイズ)の時の学習効率と品質がよい。
参考文献
Common Diffusion Noise Schedules and Sample Steps are Flawed
On the Importance of Noise Scheduling for Diffusion Models
Diffusion Noise Optimization (DNO)
損失関数の重みづけ
Debiased Estimation Loss
モデルの再構成誤差に \(\scriptsize \dfrac{1}{\sqrt{SNR(t)}}\) を掛けるだけ。v prediction の場合は、\(\scriptsize \dfrac{1}{SNR(t) + 1} \)。
\[ \small \begin{split} L &= \sum_t \mathbb{E}_{x_0,\epsilon} \left[ \frac{1}{\sqrt{SNR(t)}} ||\epsilon - \epsilon_{\theta}(x_t, t)||^2 \right]\\ \alpha :&= \prod^t_{s=1}(1-\beta_s)\\ SNR(t) &=\dfrac{\alpha_t}{1-\alpha_t} \end{split} \]Rectified Flow の場合
SD3 では以下の式で損失関数を重みづけしてる。ノイズが大きい部分のウェイト大きくなる。T=0, 1 の時ウェイトが0に近い。\(\scriptsize \lambda^{\prime}\) は SN 比の微分。
\[\small \begin{split} logit(t) &= log\dfrac{t}{1-t}\\ \pi_{ln}(t;m,s) &= \dfrac{1}{s\sqrt{2\pi}}\dfrac{1}{t(1-t)}exp \left (-\dfrac{(logit(t)-m)^2}{2s^2}\right ) \end{split} \]
横軸はタイムステップ、縦軸は\(\pi_{ln}(t; 0.00, 1.00)(\dfrac{\sigma_t\lambda^{\prime}_t}{2})^2\)
\(\lambda^{\prime} = 2(\frac{a^{\prime}_t}{a_t} - \frac{b^{\prime}_t}{b_t})\)
\(\sigma_t = a = t\)
\(b = 1 - a\)
a' と b' はそれぞれ a と b の微分
プロットコード
import math
import matplotlib.pyplot as plt
num_sample = 1000
#f = lambda x: x
old_a = 0
old_b = 1
dt = 1/num_sample
def f(t,m=0,s=1):
if t==0 or t==1:
return 0
global old_a, old_b
a = t # [0, 1]
b = 1 - a
sigma_t = t
logit = lambda x: math.log(x) - math.log(1-x)
coeff = 1/(s*math.sqrt(2*math.pi) * t * (1-t))
logit_coeff = coeff * math.exp(-(logit(t)-m)**2 / 2*s*s)
diff_a = (a - old_a)/dt
diff_b = (b - old_b)/dt
diff_lambda = 2*(diff_a/a - diff_b/b)
old_a = a
old_b = b
return logit_coeff * (sigma_t * diff_lambda / 2)**2
inputs = [i/num_sample for i in range(num_sample+1)]
outputs = [f(x) for x in inputs]
title = ''
label_x = 'Timestep'
label_y = 'Loss scale'
flg = plt.figure()
ax = flg.add_subplot()
ax.plot(inputs, outputs)
ax.set_title(title)
ax.set_xlabel(label_x)
ax.set_ylabel(label_y)
# グリッド線
ax.grid(axis='x') # グリッド線の表示:X軸
ax.grid(axis='y') # グリッド線の表示:Y軸(点線)
ax.set_yticks(range(0, 50, 5))
flg.show()
plt.show()
ファインチューン
マルチモーダルモデルのファインチューンは ShareGPT-4o-Image が詳しい。
小さいモデルを先に訓練する
SANA 1.5 では 1.6B を作成した後、学習済みの 1.6B にブロックを追加して 4.8B にしてから学習させることで、4.8B の学習時間を 60% 削減している。
蒸留
性能向上の面から見れば小さいモデルは、大きなモデルから蒸留した方が性能はよくなる。
参考文献
基盤技術と初期研究
Progressive Distillation for Fast Sampling of Diffusion Models。教師モデルのサンプリングステップを段階的に半減させる反復的蒸留アプローチ。この手法は8192ステップの教師モデルを4ステップまで蒸留する。
Consistency Models。ノイズから画像を直接予測する方法を学ばせる。
SCott: Accelerating Diffusion Models with Stochastic Consistency Distillation。ODE ではなく SDE を使って蒸留する。
DMD1 One-step Diffusion with Distribution Matching Distillation。ノイズとノイズから生成された画像とから学習するが、Distribution Matching Gradient も同時に計算し、その loss も学習させる。
LoRA-Enhanced Distillation on Guided Diffusion Models。蒸留+LoRA
理論的基盤と分析
Theory of Consistency Diffusion Models: Distribution Estimation Meets Fast Sampling
Improved Techniques for Training Consistency Models
Optimal Stepsize for Diffusion Sampling
Simple and Fast Distillation of Diffusion Models。従来手法の最適化。
スコア蒸留と敵対的学習(GAN)
Score Distillation Sampling with Learned Manifold Corrective。Score Distillation Sampling (SDS) の損失関数の問題点の解説とその解決方法の説明。
Adversarial Diffusion Distillation。スコア蒸留+GAN。
UFOGen: You Forward Once Large Scale Text-to-Image Generation via Diffusion GANs。拡散モデル+GAN。蒸留ではない。
Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation。蒸留、GAN、分布マッチング。
Latent Consistency Models (LCM)
Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference。Consistency Model を潜在空間で実行。
LCM-LoRA: A Universal Stable-Diffusion Acceleration Module。蒸留をファインチューニングととらえて、LoRA に蒸留内容を学習させる。
Hyper-SD: Trajectory Segmented Consistency Model for Efficient Image Synthesis
フローマッチングと軌道最適化
InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generation。Rectified Flow を使った蒸留。
One Diffusion Step to Real-World Super-Resolution via Flow Trajectory Distillation。多段階フローマッチングモデルをワンステップ超解像モデルに蒸留。
ワンステップ生成の特殊化
SwiftBrush: One-Step Text-to-Image Diffusion Model with Variational Score Distillation。Score Distillation Samplingを活用し、訓練画像データなしで蒸留。
SwiftBrush v2: Make Your One-step Diffusion Model Better Than Its Teacher。LoRA の採用と最適化。
高度な蒸留技術
Target-Driven Distillation: Consistency Distillation with Target Timestep Selection and Decoupled Guidance。蒸留に使うタイムステップを厳選。学習中に Guidance Scale を分離して、CFG が使えるようにする。
DMD2 Improved Distribution Matching Distillation for Fast Image Synthesis。回帰 loss を廃止して、データセット作成コストを削減。ノイズから生成された画像以外の画像も使うことで教師を超える性能を実現。
LoRA-Loop: Closing the Synthetic Replay Cycle for Continual VLM Learning。VLM を使用して品質を向上させる方法。シードガチャをして、質のいいものだけを学習させる。
Synthetic Data is an Elegant GIFT for Continual Vision-Language Models
知識蒸留の拡張
Knowledge Diffusion for Distillation。学生モデルの出力の feature をデノイズするモデルを作る?
Your Student is Better Than Expected: Adaptive Teacher-Student Collaboration for Text-Conditional Diffusion Models。学生のデノイズ結果が教師を上回った場合は学習しないようにすることで、蒸留モデルが教師を上回る性能になる。
新興技術
Distilling ODE Solvers of Diffusion Models into Smaller Steps。ODE ソルバーを蒸留する?
Accelerating Diffusion Sampling with Optimized Time Steps。ODE ソルバーと学習に使う最適なタイムステップとの選択方法の解説。
Learning Few-Step Diffusion Models by Trajectory Distribution Matching。Distribution Matching + Trajectory Matching。
動画と超解像
DOLLAR: Few-Step Video Generation via Distillation and Latent Reward Optimization
Acceleration Method for Super-Resolution Based on Diffusion Models by Intermediate Step Prediction
RL Dreams: Policy Gradient Optimization for Score Distillation based 3D Generation
大きいモデルのブロックを削減して小さいモデルを作成する
SANA 1.5 では大きなモデルの推論時のブロックの入力と出力との類似度を計算し、性能に貢献していないブロックを削除する。その後でファインチューニングすることで小さいモデルを作成している。ブロック削除後のファインチューニングはたった 100 ステップ程度で十分。
SLEB: Streamlining LLMs through Redundancy Verification and Elimination of Transformer Blocks ではトランスフォーマーブロックの入力と出力との類似度をコサイン類似度で計算している。
Do Language Models Use Their Depth Efficiently? はコサイン類似度と残差ストリームとを使用して分析している。
省メモリ学習
- 8-bit Optimizers(AdamW, CAME)
- Gradient Checkpointing
- Gradient Accumulation(バッチサイズによるが、最もメモリ削減効果が大きい)
- Freezing embeddings
- AMP(Automatic Mixed Precision)(NVIDIA GPU のみ。Tensor コア使用。学習が最大2倍高速化する効果もある)
8-bit Optimizers
8bit AdamW は bitsandbytes をインストールすると使える。ただし CNN を多用する VAE では効果は小さい。
bitsandbytes
枯れたバージョンの CUDA を使用しているなら普通にインストールする。
!python -m pip install bitsandbytes
使用している CUDA のバージョンに対応したものをインストール場合。CUDA のバージョンチェックは venv を有効にした状態で nvcc -V。
!python -m pip install bitsandbytes-cuda111
使えるかどうかは python -m bitsandbytes で以下のメッセージが出るか調べる。
PyTorch settings found: CUDA_VERSION=124, Highest Compute Capability: (8, 6). Checking that the library is importable and CUDA is callable... SUCCESS!
使用方法
import bitsandbytes optimizer = bitsandbytes.optim.AdamW8bit(model.parameters(), lr=learning_rate)
Gradient Checkpointing
Training Deep Nets with Sublinear Memory Cost
勾配の計算には各ブロックのフォワードパスの結果を保存しておく必要がある。しかしこれはメモリ使用量が大きい。そこで、チェックポイントに指定したブロックのみフォワードパスの結果を保存しておき、バックプロパゲーション時にフォワードパスの結果が必要になると、一番近いチェックポイントからフォワードパスの結果を再計算する。
メモリ使用量が減るが、学習時間が増加する。
バッチサイズ
大きい場合
利点
- 学習時間の短縮
- 勾配推定の安定性。勾配が均されるため、ノイズが減り学習が安定する
欠点
- 汎化性能の低下。 勾配のノイズが少ないため、局所最適解・sharp minimaに陥りやすい
- メモリ使用量の増加
- 学習率の調整が難しい
小さい場合
利点
- 汎化性能の向上。flat minima に収束しやすく、これは汎化性能の向上に効果があると考えられている
- 局所最適解からの脱却
- メモリ使用量が少ない
欠点
- 学習時間の増加
- 勾配推定が不安定になる
Gradient Accumulation
バッチサイズを小さなミニバッチの集積で近似する。目的のバッチサイズにメモリが足りない場合の選択肢。分散学習時にも使われる。バックプロパゲーションの回数が減るので学習が高速化することもある。
コードはGPUメモリ使用量を削減する を参照。
AMP
16bit でデータを保存するが、精度が必要な部分は 32bit を使う。メモリ使用量が減少し、計算が速くなる。
AMP は精度の低下を抑えるために loss をスケールして 16bit の精度の範囲内になるように調整している。なので逆伝播の際にスケールを元に戻す必要がある。
torch.cuda.amp ではなく torch.amp を使う。
from torch.amp import autocast, GradScaler
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = GradScaler('cuda')
# 学習ループ
for data, target in dataloader:
optimizer.zero_grad()
with autocast('cuda'):
# autocastコンテキスト内で計算
output = model(data)
loss = nn.functional.cross_entropy(output, target)
# スケーラーで勾配をスケーリング
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
分散学習
The Basics of Distributed Training in Deep Learning: Speed, Efficiency, and Scalability
Distributed Deep Learning: Training Method for Large-Scale Model
DeepSpeed Training Overview and Features Overview
データ並列
モデルが各ノードのメモリに収まる場合、まず各モデルでバッチを処理して勾配を計算する。各ノードの勾配を集計して、各ノードのモデルを更新する。この方法だと、バッチの処理の一番遅かったノードに律速される。そこで、各ノードでデータを学習させて後でモデルをマージする方法がある。
モデル並列
LLM のようなモデルが1つのノードに収まらない巨大なモデルは、モデル自体を複数のノードに配置する。
実例
PixArt-α
学習に $28,400(2025年のドル円レートで、およそ 400 万円)しかかかっていない安価なモデル。学習にかかった時間は 64 台の V100 で 26 日。ちなみに Stable Diffusion 1.5 の学習コストは $320,000。

出典:PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis. Junsong Chen et al. Table 4. https://arxiv.org/abs/2310.00426
V100 は 14 TFLOPS、RTX5090 は 105 TFLOPS なので、RTX5090 1台だと学習に 160 日かかる。400 万円で RTX5090 を 10 台買えば 16 日でできる(契約アンペアを 60A 以上にする必要がある)。
オプティマイザーは AdamW。
Stable Diffusion 3
データセットは conceptual-12m と ImageNet。
バッチサイズは 1,024 でオプティマイザーは AdamW、lr = 1e-4 で、1,000 ステップの線形ウォームアップ。
Lumina-Image 2.0
32 台の A100(FP32 は 19.5 TFLOPS) を使用。3段階の訓練でそれぞれ、191, 176, 224 GPU*Days。191 + 176 + 224 = 591GPU*Days、32 台で割ると学習日数は 18.5 日。オプティマイザーは AdamW。

出典:Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Qi Qin et al. Table 3. https://arxiv.org/abs/2503.21758
HiDream-I1
事前学習
256 x 256 を 600,000 ステップ。24 バッチ/GPU。
512 x 512 を 200,000 ステップ。8 バッチ/GPU。
1,024 x 1,024 を 200,000 ステップ。 2 バッチ/GPU。
AdamW, lr=1e-4, 1,000 ステップの線形ウォームアップ。mixed-precision と gradignt checkpointing を使用。
ファインチューニング
手動でキャプションを付けたハイクオリティな画像で 20,000 ステップ, lr=1e-5, グローバルバッチサイズ 64。
SANA
Pytorch DistributedDataParallel を使用し、64 台の A100 で、かかった時間は非公開。512 x 512 の解像度で 1e-4 の学習率で 200K ステップ以上事前学習を行い、1024 から 4096 へ解像度を増やしながら、2e-5 の学習率で最大 10K ステップの教師ありファインチューニングを行う。オプティマイザーは 8bit CAME。1.6B のモデルで、CAME 8bit は AdamW 8bit よりわずかに性能がよく、メモリ使用量は若干少ない(43GB vs 45GB)。
1/3 のレイヤーだけを使って事前学習させた後、乱数で初期化した残りの 2/3 のレイヤーを追加して学習させる。
DiT-Air
目的関数フローマッチング、オプティマイザーは AdaFactor、1e-4 の固定学習率。グローバルバッチサイズ 4,096 で100 万ステップ。
Nitro-T
学習期間は 32 AMD Instinct MI300X GPU で1日。
512px を 100k ステップ、1,024px を 20k ステップ。詳細は GitHub を参照。
REPA(DINO v2)はマスクしてないトークンにのみ適用。
データセット。()は使用枚数。
- Segment Anything 1B(10M)
- CC12m(7M)
- DiffusionDB(14M)
- JourneyDB(4M)
AI 生成画像を使いすぎると、画風がカートゥーン化する。JourneyDB は品質がよいので、ファインチューニングには DiffusionDB ではなく JourneyDB を使う。
Illustrious: an Open Advanced Illustration Model
バージョン | ステップ数 | バッチサイズ | データ枚数 | プロンプトスタイル | アノテーション手法 | 解像度 |
---|---|---|---|---|---|---|
0.1 | 781,250 | 192 | 7.5M | タグベース | オリジナルプロンプト + 手動フィルタリング/再構成 | 1,024 |
1.0 | 625,000 | 128 | 10M | タグベース | オリジナルプロンプト + 手動フィルタリング/再構成 | 1,536 |
1.1 | 93,750 | 512 | 12M | タグベース | マルチレベルキャプション | 1,536 |
2.0 | 78,125 | 512 | 20M | タグベース | マルチレベルキャプション | 1,536 |
- タグベース:Danbooru形式などのタグによるプロンプト設計
- マルチレベルキャプション:タグと自然言語キャプションを併用した注釈方式
- 手動フィルタリング/再構成:クリーンデータの選別やタグ整備の工程を含む
Danbooru 画像とタグは以下の問題がある
- 同一トークンに複数の意味が重なることが多く、曖昧性や誤解の原因になる
- 例:「doctor」というタグは職業名としてもキャラクター名としても使われる
- 画像にタグが極端に少ないものが存在する
- 極端に高解像度な画像がある
- 極端なアスペクト比の画像がある
- コミック形式の画像がある
キャプション構造
NovelAI のようなタグオーダーを採用している。
人数 ||| キャラ名 ||| レーティング ||| 一般タグ ||| アーティスト名 ||| パーセンタイルベースのレーティング ||| 年代
最初はパーセンタイルベースではなく、スコアレンジベースのレーティングだった。しかしスコアレンジは年代やカテゴリでばらつきが大きかったので、以下のようなパーセンタイルベースのレーティングを採用することにした。
レーティングタグ | パーセンテージ |
---|---|
worst quality | ~8% |
bad quality | ~20% |
average quality | ~60% |
good quality | ~82% |
best quality | ~92% |
masterpiece | ~100% |
解像度
v2.0 では 0.15MP(500px * 300px)から 2.25MiP(1536 px * 1536px / (1024*1024) = 2.25)。
訓練方法
- No Dropout Token。性的(provocative)なトークンを含む NSFW トークンはドロップアウトしない。加えて NSFW トークンを CFG で制御(つまりネガティブプロンプトに NFSW トークンを入れる)することで 100% 性的なコンテンツを生成しないモデルを作成できた。
- コサインアニーリング
- 疑似レジスタートークン。データ数が少なかったりモデルが理解できない概念を吸わせるトークンを入れる。シーケンス長を調整するためのパディングトークンがレジスタートークンとして機能してしまうことがあるので注意が必要。
- キャラ名とアーティスト名をドロップアウトしない Contrastive Learning(ラベルなしで学習させる手法)。キャラの特徴とアーティストのスタイルを効率よく学習できる。ただし、アーティストとキャラとの結びつきが強くなる欠点がある。
- キャプションの言い換え。"1girl" を "one girl", "single woman" に一定確率で入れ替える。
- マルチレベルドロップアウト。No Dropout Token 以外のトークンは以下ドロップアウトのどれかが確率的に適用される。
- 30%: max(0.3 * total tokens, 10)
- 20%: max(0.4 * total tokens, 15)
- 10%: min(total tokens, 6)
- 4%: min(total tokens, 4)
- 36%: ドロップアウトなし?(論文に記載なし)
- Input Perturbation Noise Augmentation(0 < ε < 0.1)(Ning et al. [2023])
- Debiased Estimation Loss(Yu et al. [2024])
訓練設定
バージョン | データ枚数 | バッチサイズ | 学習率 | テキストエンコーダー学習率 | エポック | 解像度 | プロンプトスタイル | ドロップアウト | レジスタトークン | マルチキャプション |
---|---|---|---|---|---|---|---|---|---|---|
0.1 | 7.5M | 192 | 4.5e-6 | 20 | 1,024 | タグ | ✖ | ✖ | ✖ | |
1.0 | 10M | 128 | 1e-5 | 6e-6 | 8 | 1,536 | タグ | 〇 | 〇 | ✖ |
1.1 | 12M | 512 | 3e-5 | 4e-6 | 4 | 1,536 | タグ+ | 〇 | 〇 | ✖ |
2.0 | 20M | 512 | 4e-5 | 3e-6 | 2 | 1,536 | タグ+ | 〇 | 〇 | 〇 |
ファインチューニングの問題点
ユーザーの評価をフィードバックさせたファインチューニングは生成される画像の多様性がなくなる。流行の画風が高く評価され、AI が上手く書けない手や背景などを描かなくなる。
Illustrious がファインチューニング前のモデルを公開しているのはこれが理由だ。ファインチューニング版を LoRA 形式で配布する方法も考えられる。
参考文献
The Delta Learning Hypothesis。品質差のあるモデルの出力を DPO を利用して学習させることで、性能が向上できる可能性がある。
性能検証
論文で使われるベンチは基本的にプロンプトと生成された画像との忠実性を計測する。以下の要素があっても減点されることはない。
- 指の本数を間違える
- 細部が溶けている
- 画質が悪い
- 直線がゆがんでいる
- 文字がおかしい
- オブジェクトの細部がおかしい
- 水平線がずれている
ベンチ
高速化
DataLoader
persistent_workders
Windows 環境ではエポックの開始のプロセスの生成と破棄とに時間がかかる。persistent_workers を True に設定するとこれを短縮できる。
train_loader = torch.utils.data.DataLoader(
...
persistent_workers=(os.name == 'nt'),
)
BatchSampler
DataLoader のデフォルトで設定される BatchSampler はバッチサイズが巨大な場合は遅い。バッチサンプラーを自作する場合は自分でバッチサンプラーをつくる (PyTorch)や PytorchのDataloaderとSamplerの使い方を参照。
学習高速化
torch.compile
How to use torch.compile on Windows CPU/XPU で C++ 開発環境を入れたら、以下の python コードでコンパイルして学習・推論可能。速度向上率は 10 % 前後。
# mode は 'default', 'reduce-overhead', 'max-autotune'
model=torch.compile(model, mode="default")
GLU Variants Improve Transformer
FFN の代わりに SwiGLU を使うと表現力の向上、収束の高速化(学習の高速化)、大規模モデルでの性能向上が見られる。
REPA Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think
pretrained visual encoder の出力を正則化項に追加することで、高速化と生成品質の向上を実現。学習速度が最大 17.5 倍向上する。
DINO v2 がよく使われる。

出典:Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think. Sihyun Yu et al. Figure 1. https://arxiv.org/abs/2410.06940
Learning Diffusion Models with Flexible Representation Guidance
REPA の発展形の REED は、REPA の4倍速く学習できる。
ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
パラメータを共有することで、性能低下を抑えてパラメータ数を削減。
Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models
パッチ単位で評価すると同時に、パッチサイズを変更することで2倍以上の学習速度を達成。この手法の一番の利点は、学習画像の枚数が最低 5,000 枚から学習できること。
ただし、切り出した画像とキャプションとを一致させるのが面倒。
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
アテンションブロックの行列演算の最適化。
Lightning Attention-1: TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer
アテンションのスケールと softmax をなくして、アテンションを O = Norm(Q(KTV)) に置き換える。
Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models
HBM のあるハードウェアが対象。Linear Attention は cumulative summation(cumsum) のせいで性能の理論値が出せない。Lightning Attention-2 では LLaMA で理論値が出せることを実証した。
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
分散環境でデータの冗長性を減らす技術。
PixArt-Σ
PVT v2 に似た key/value の圧縮によって高速化する。
Transformers without Normalization
Normalization Layer を \(\scriptsize{DyT(\mathbf{x}) = \gamma \times tanh(\alpha \mathbf{x}) + \beta}\) で置き換えると計算が速くなる。性能は-1~1ポイント上下する。レイヤー単体では 40~50% 高速になり、モデル全体では8%程度高速になる。ただし torch.compile した場合両者に速度差はない。
# x の次元は [B, T, C]
# B: バッチサイズ, T: トークン数, C: 隠れ層の次元
class DyT(nn.Module):
def __init__(self, channels: int, init_alpha: float=1.0):
"""
tanh を使った Normalization。精度を維持して Normalization を高速化
Args:
channels (int): 隠れ層の次元
init_alpha (float): x にかかる係数
"""
super().__init__()
self.alpha = nn.Parameter(torch.tensor([init_alpha]))
self.beta = nn.Parameter(torch.zeros(channels))
self.gamma = nn.Parameter(torch.ones(channels))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.gamma * torch.tanh(self.alpha * x) + self.beta
Data Efficacy for Language Model Training
データセットをスコアリングし厳選することで、データ量を半減させても同じ性能を出す手法の解説。
Implicit Reward as the Bridge: A Unified View of SFT and DPO Connections
LLM のケースだが、SFT と DPO が暗黙的報酬学習の一種であることを数学的に証明し、SFT の学習率を下げることで性能を向上させられることを実証した。
推論高速化
Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference
U-Net のエンコーダーをキャッシュすることで高速化する。
Delta-DiT: A Training-Free Acceleration Method Tailored for Diffusion Transformers
デノイズの初期段階では DiT ブロックの後方をキャッシュし、後半では前方をキャッシュする事で高速化する。これは DiT ブロックが画像の概形(前方のブロック)と詳細(後方のブロック)を生成する役割の違いに基づいている。
Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model
TeaCache。一定間隔でキャッシュを破棄するのではなく、キャッシュがしきい値を超えるまでキャッシュを使い続ける。
直にタイムステップを減らすより TeaCache の方が有利。直にタイムステップを減らした場合はデノイズに使うタイムステップを選択できない。TeaCache は画質に影響が出そうになるとキャッシュを破棄するので、画質の劣化を抑えられる。TeaCache は(出力 - 入力)の差分をキャッシュしているので、モデルを動作させなくてもキャッシュを利用したデノイズが行われるが、タイムステップを減らした場合はそうではない。
量子化と枝刈り
Joint Pruning, Quantization and Distillation for Efficient Inference of Transformers
訓練済みモデルから重要度の低い重みやフィルターを削減する枝刈りの解説。
SANA 1.5: Efficient Scaling of Training-Time and Inference-Time Compute in Linear Diffusion Transformer
- 16K 以上のパラメータの Linear と 1x1 畳み込みを8bit 整数で量子化
SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models
ウェイトには外れ値がある。この外れ値が量子化の精度を低下させる。そこで、行列を外れ値の部分 W と量子化しやすい部分 X とに分け、元の行列を W + X で表現する。量子化しやすい X は4bit 量子化、外れ値の部分は SVD で分解、16bit の精度で LoRA(rank=32)化する。
SDXL, PixArt, FLUX.1 で動作検証をしている。
LittleBit: Ultra Low-Bit Quantization via Latent Factorization
SVDQuant と同様に、行列を LoRA のように分解してから量子化することで、最大 0.1 BPWを達成している。
そのほか
Understanding Transformer from the Perspective of Associative Memory。活性化関数には速度と記憶容量とのトレードオフが存在する。ReLU は高速だがパラメータあたりの記憶量は小さくなる。Softmax は遅いがパラメータあたりの記憶容量が多くなる。
Accelerate High-Quality Diffusion Models with Inner Loop Feedback
DeepCache: Accelerating Diffusion Models for Free
A Simple Early Exiting Framework for Accelerated Sampling in Diffusion Models
サンプラー
PFDiff: Training-Free Acceleration of Diffusion Models Combining Past and Future Scores
参考文献
Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation
SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers GitHub
CogView3: Finer and Faster Text-to-Image Generation via Relay Diffusion
Playground v3: Improving Text-to-Image Alignment with Deep-Fusion Large Language Models
SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers
Flow Matching for Generative Modeling
Unveiling the Secret of AdaLN-Zero in Diffusion Transformer
GELU の論文:Bridging Nonlinearities and Stochastic Regularizers with Gaussian Error Linear Units
GELU Activation Function in Deep Learning: A Comprehensive Mathematical Analysis and Performance
DualFast: Dual-Speedup Framework for Fast Sampling of Diffusion Models。離散化誤差と近似誤差という2つのサンプリング誤差を同時に考慮することで、DPMのサンプリング速度を向上。
Fast Sampling of Diffusion Models via Operator Learning
An Interpretable Framework for Drug-Target Interaction with Gated Cross Attention
Understanding Batch Normalization
Normalization はそれ自身に性能を上げる効果はない。Normalization を入れることで高い学習率でも学習が安定するようになり、その結果、速い学習と局所的最適解の回避とが実現している。
拡散モデルのミニ実装
A Gentle Introduction to Diffusion
特異値分解(SVD)
scipy.sparse.linalg.svds(アルゴリズムはARPACK(Augmented Implicitly Restarted Lanczos Bidiagonalization Algorithm)を使用)で疎行列の特異値分解ができる。
次元削減目的なら sklearn.decomposition.TruncatedSVD も使える。
そのほか
LLM推論に関する技術メモ
LLM を利用するサービスを提供する場合の技術的注意点について網羅的に解説されている。
マスクテンソルの作成
block_count は1辺のブロック個数。
import torch
def create_mask(image_size:int, patch_size:int, hide_rate:float, batch_size:int, image_channel:int=3):
"""
Args:
image_size (int): 解像度
patch_size (int): パッチ解像度。F。
hide_rate (float): 数値が増えるとより多くマスクされる。範囲は [0, 1]
"""
assert image_size % patch_size == 0
block_count = image_size//patch_size # 1辺のブロック個数
scale_factor = int(image_size/block_count)
mini_mask = (torch.rand(block_count*block_count) >= hide_rate).float().view(-1,1,block_count,block_count)
mask = torch.nn.functional.interpolate(mini_mask, size=None, scale_factor=scale_factor, mode='nearest-exact')
return mask.expand(batch_size, image_channel, image_size, image_size)