T2I 拡散モデルの設計メモ

目次

• テキストエンコーダー
• VAE
• VAE の学習
• LPIPS
• GAN
• コンディショニング
• 位置埋め込み
• 目的関数
• ブロック図
• データセット
• キャプショニング
• 学習方法
• タイムステップスケジューラー
• ノイズスケジューラー
• 蒸留
• 省メモリ学習
• 8-bit Optimizers
• Gradient Checkpointing
• バッチサイズ
• Gradient Accumulation
• AMP
• 性能検証
• 高速化
• 学習高速化
• 推論高速化
• 量子化と枝刈り

テキストエンコーダー

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

• T5
• OpenVision
• BLIP3-o
• Qwen3
• phi-4-mini-instruct
• Gemma2-2b-it-gguf
• Gemma3

性能比較

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

VAE の圧縮率を決めるパラメータは２つある。

• F：解像度の縮小率
• C：潜在空間のチャンネル数

拡散モデルの隠れ層の次元はさらにパッチサイズが関係している。

• P：パッチサイズ

VAE の隠れ層の次元数は以下のようになる。

拡散モデルの隠れ層の次元は以下のようになる。

各 VAE の設定

Dit-Air の F8C8P2 = F16C32P1 が画質と処理効率とのバランスがいい。SANA は画質が悪すぎるし、SD3・FLUX は処理が重すぎる。

パッチサイズは小さい方が拡散モデルの性能が良い

パッチサイズは 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

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] を使うことが多い。

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
)

# 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()

VQ-Diffusion Vector Quantized Diffusion Model for Text-to-Image Synthesis

VQ-VAE は量子化されているので、小さなミスが大きな意味の変化へつながる恐れがある。VQ-Diffusion では Mask-and-replace という学習手法でその問題に対処している。

D²iT

D²iT は圧縮率の異なる２つの潜在空間表現にエンコード可能な VAE を使う。ディティールの細かい部分は低圧縮率、情報の少ない部分は高圧縮率と使い分ける。デノイザもディティールの細かい部分は多くのネットワークが通過する。

SANA

SANA は高圧縮の VAE を利用することで通常の VAE 比で 10 倍高速化している。従来の VAE は８倍の圧縮率だったが、SANA は 32 倍の圧縮率の VAE を採用している。

SANA の VAE は dc-ae-f32c32-sana-1.1-diffusers からダウンロードできる。モデルの定義は efficientvit。

細部が大きく崩れるので、等倍では使い物にならない。Detailer の使用がほぼ必須で、顔なら解像度は 1024 x 1024 は欲しい。

オリジナル

VAE（dc-ae-f32c32-sana-1.1-diffusers）通過後

差分

Latent Diffusion Models with Masked AutoEncoders

Masked Autoencoders に KL 項を追加して、潜在空間のスムースさ・知覚的圧縮品質・再構成画像の品質の３つの指標を同時に達成。再構成 loss はマスクしていない部分だけで、マスク部分の loss・LPIPS loss・KL 項を正則化に使っている。GAN は未使用。

対称形の ViT ベースのエンコーダーとデコーダーとを採用している。SD-VAE がファイルサイズ 320MB に対して VMAE は 43MB で、性能で SD-VAE を上回る。

F4C16 のモデルの処理は以下のようになる。Transformer の次元は d_model とする。Transformer は Transformerエンコーダの単一の層（Multi-Head Self-AttentionとFeed-Forward Network） のみの torch.nn.TransformerEncoderLayer が便利。

1. 元画像（3 x 512 x 512）
2. 画像のパッチ化によるダウンサンプル（3 x 512 x 512）->（64 x 128 x 128）-> (16,384 x 64)
3. Linear で隠れ層の次元を調整 (16,384 x 64) ->（16,384 x d_model）
4. Transformer
5. Linear 等で \(\scriptsize{\mu}\) と \(\scriptsize{\sigma}\) とを計算（それぞれ（16,384 x 16））
6. z をサンプリング（z.shape =（16,384 x 16））
7. Linear（16,384 x 16）->（16,384 x d_model）
8. Transformer
9. Linear で隠れ層の次元を調整 (16,384 x d_model) -> (16,384 x 64)
10. アンパッチ (16,384 x 64) -> （64 x 128 x 128）->（3 x 512 x 512）

参考文献

Masked Autoencoders Are Scalable Vision Learners

MaskGIT: Masked Generative Image Transformer

MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis

DC-AE Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models

2025 年以降に発表された DiT でよく使われる。SANA や Nitro-T で使われている。space-to-channel, channel-to-space（要はパッチとアンパッチ）を学習可能な残差接続と、３段階の学習で性能を向上させる。

1. 低解像度でエンコーダーとデコーダーを学習させる
2. 高解像度でエンコーダーの終わりとデコーダーの入り口とのみを学習させる
3. 高解像度でデコーダーの終わりのみを学習させる

GitHub

VAE の学習

VAE の学習は簡単ではない。VAE 単体ではぼやけた画像にしかならない。

Stable DiffusionのVAEは、LPIPS、パッチベースのGAN損失、およびMSE損失を含む複雑な損失関数を使用して訓練される。この組み合わせにより、VAEは、ピクセル単位で正確 (MSE)、知覚的に整合 (LPIPS)、かつ全体的にリアル (GAN損失) な、高品質でぼやけのない再構成を生成できる。

taming-transformers は LPIPS と GAN とを使用して VQ-VAE を学習させており、VAE の作成の時に参考になる。

posterior collapse

VAE は同じ画像ばかり出力する posterior collapse（KL loss が急速に０に近づく）が起こりやすい。そこで以下のような戦略が良く使われる。

1. KL Loss の逆伝播の量を調整する kl_beta を設定する。
2. 最初の数万ステップは kl_beta を０にする。
3. 次にウォームアップ区間を設ける。 kl_beta が 0→1 になる“ウォームアップ区間”は、実装・データ規模により 1 万〜10 万 SGD ステップ、あるいは 50〜200 エポック程度

• ウォームアップ区間が短すぎると posterior collapse が起こりやすく、長すぎると訓練が停滞するため、近年では「1〜2% の総ステップ数をウォームアップに割く」か「 ELBO が安定するまでβを上げない」設計が主流
• 最低でも１万ステップのウォームアップ期間が必要。高解像度画像・映像用 VAE では 2.5〜5 万ステップ程度の長めのウォームアップが採用される傾向があり、テキスト VAE では 1〜2 万ステップ前後が多い
• posterior.mean を可視化して、すべて 0 近辺になっている場合、posterior collapse の疑いが強い
• print(f'posterior.mean mean:{torch.mean(posterior.mean):0.4f}, var: {torch.var(posterior.mean): 0.4f}')

正規化

画像の表示時や保存時に逆正規化を忘れるのはよくやりがちだ。

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)

Your VAE Sucks

1. ぼやけた画像は MSE が小さいことが多い。ぼやけた画像は位置の微妙な変化にも対応できる
2. 最尤推定は訓練データとは著しく異なるデータを生成してもペナルティがない

正則化項

• KL。KL VAE はオートエンコーダーに KL 正則化項を追加したもの
• LPIPS
• GAN
• Watsonの知覚モデル。離散フーリエ変換 (DFT) 等を使い特徴を抽出。LPIPS より６倍高速で、省メモリ
• 輪郭情報。Unleashing the Power of One-Step Diffusion based Image Super-Resolution via a Large-Scale Diffusion Discriminator
• そのほかは画像比較手法まとめを参照

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 をとる。

以下の式で \(\small{\tau}\) はスケール・回転変換、\(\small{\mathcal{E}}\) はエンコーダー。\(\small{\mathcal{L}_{rec}}\) は再構成 loss。\(\small{\mathcal{L}_{reg}}\) は KL loss。

潜在空間表現の回転と拡大縮小するコードは以下のようになる。潜在空間表現は３次元 [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")

GitHub

VAE 学習の参考文献

taming-transformers

Community Training AutoencoderKL #894

How to train your VAE

Cyclical KLAnnealing Schedule

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

参考文献

【徹底解説】VAEをはじめからていねいに

Adversarial Autoencoders

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 倍の圧縮率を達成している。

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 の β₁ を 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

導入順

1. WGAN-GP ・ LSGAN
2. Minibatch Discrimination ・ Feature Matching
3. InfoGAN ・ Mode-seeking loss
4. Unrolled GAN ・ PacGAN

LSGAN

JS Divergenceの最小化に基づくGANでは、勾配消失が起きやすい。LSGAN はGANの損失関数を最小二乗誤差にすることで、この問題の解決する。WGAN より実装が単純なので、まずは LSGAN から試すのがよい。

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

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}")

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

VGG の出力と Adversarial loss と Content Loss とを損失関数として採用している。

Content Loss

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 によると、クロスアテンションよりノイズ画像と結合する方がパラメータ効率が良い。

結合

モデルの隠れ層の次元が d_model で、テキストエンコーダーの出力 T が [n_t, d_t] で、潜在空間の画像の次元 IMG が [n_i, d_i] とする。

これらを結合するには Linear などで、T と IMG の隠れ層の次元を d_model に変換する必要がある。

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

adaLN

Layer Nomalization のガンマとベータは学習可能なパラメータだ。adaLN はこれを線形変換などを使用して動的に変更するもの。つまり

c = タイムステップなど
γ = Linear(c)
β = Linear(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 が詳しい。

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）

Omni-RoPE

RPE-2D Boosting Resolution Generalization of Diffusion Transformers with Randomized Positional Encodings

Mix-FFN

3x3 の畳み込みによって RoPE なしでも相対位置情報を取得できる。SegFormer の Mix-FFN は以下のようになる。

xout = MLP(GELU(Conv3x3(MLP(xin)))) + xin

SANA の Mix-FFN は GELU ではなく SiLU を使っている。SiLU（Swish） は ReLU に形が似ていて、(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

ブロック図

トランスフォーマーブロックの内部に以下のパターンが２～３回出現するのが基本。

1. AdaLN
2. Attention・MLP ・FFN
3. Scale

DiT

FFN

Pointwise Feedforward Network と Position-wise Feed-Forward Network とはほぼ同じ意味。バッチ・系列長・埋め込み次元のテンソルにおいて、各位置 i にあるベクトル x_i に対して、まったく同じ重みの MLP を適用する。数式は以下のようになる。

Patching（Patchify）

VAE の設定が F8C4P2 とする。元の画像が 3 x 1024 x 1024 だとする。その潜在空間表現は 4 x 128 x 128。パッチサイズが２ということは４ピクセル（2 x 2）を一つの埋め込み次元にするので、16 x 64 x 64。16 = チャンネル数 * パッチサイズ^2。幅と高さを１次元化して転置すると、最終的に 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)

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

トランスフォーマーブロックの隠れ層と同じ次元を持つ、２レイヤーの 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

テキストと画像とを別々に処理するダブルストリーム方式。ダブルストリームはパラメータ効率が悪く、現在では選択肢に入らない。

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

• Flux.1-Architecture-Diagram
• Demystifying Flux Architecture

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。

画像をマスクしてノイズ付与、マスクごとデノイズして、マスクで損失関数を分けるタイプ。

損失関数

GitHub

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

RMSNorm は Layer Normalization の計算コストを削減したもの。具体的には各層の活性化に対して、平均の計算を省略し、二乗平均平方根 (Root Mean Square: RMS) のみを用いて正規化を行う。RMSNorm は LayerNorm と同等の性能を達成しつつ、計算時間を 7% 〜 64% 削減できる。

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

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

ゲート

Gate は以下の式で表現される。ただし Normalization の後にゲートを入れる場合は、b₁, b₂ は、出力平均が0に近いため省略されることが多い。

nn.Lienar は bias=False でバイアスを省略できる。

nn.Linear(in_features, out_features, bias=False) # バイアスなし

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 と同じ３つのテキストエンコーダーに LLM を加え、合計４つのテキストエンコーダーを使用
• 位置エンコーディングは 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 は２つの重み行列（W1, W2）を持つのに対し、GLU 派生は３つ（W, V, W2）持つ。なので SwiGLU を含む GLU の派生形のパラメータ数は、FFN と比較して、隠れ層の次元が 2/3 に削減される。

SwiGLU は外れ値同士が掛け算されて値が増幅されることがある。Smooth-SwiGLU によってその問題は解決され、LLM では FP8 での学習も可能になっている。

ReLUと比較して、SwiGLUは固有のスパース性がほとんどないため、従来のスパース性を活用する技術の効果が薄い。

DiT においては収束を 21 倍加速させた事例が報告されているが、性能が向上した論文はない。

参考文献

• GLU Variants Improve Transformer

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) = d_model * d_model 。通常のアテンションはトークン数が n とすると計算量とメモリ使用量とが O(n²) になるが、Linear Attention では O(n) になる。SD3 なら画像サイズが 1024 だと n は 4,096。

Linear Attention は高速だが性能は低下する。その低下分を Mix-FFN でカバーしている。Mix-FFN は RoPE を排除できるが、隠れ層の次元が d_model = 2,240、d_FFN = 5,600 と FFN の隠れ層の次元が巨大になっている。

畳み込み

3x3 の前の 1x1 の畳み込みによって計算量を削減している。1x1 の畳み込みはチャンネル数の次元を削減するのに使われることが多い。例えば (C, H, W) = (3, 8, 8) に 1x1 の畳み込みを実行して (1, 8, 8) にする。その後で計算コストの高い 3x3 の畳み込みを実行し、1x1 の畳み込みで (3, 8, 8) へチェンネル数を復元する。これはピクセル単位の MLP を実行しているとみることもできる。

Exploring 1x1 Convolutions in Deep Learning

パラメータ

d_model = 2,240、d_FFN = 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 * レイヤー数。

マルチヘッドアテンションの数はトランスフォーマーの深さと同じ。

Pixart と同じ MLP なら、トランスフォーマーブロックの隠れ層と同じ次元を持つ、２レイヤーの 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）の画像。

GitHub

データセット

データとパラメータ数はセットで考える。パラメータ数だけ増やしても性能は上がらない。データに多様性がなければデータを増やしても性能は上がらない。

最近では多くても 50M 程度の教師画像枚数で学習させることが多い。仕上げのファインチューニングは１～３M枚。

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.」で理由も聞ける。データセットの画像の厳選にも使える。

著作権について

平成30年著作権法改正により学習は合法。ただしモデルが著作権侵害コンテンツを高頻度で生成する場合、モデルの販売や画像生成サービスの提供は複製権侵害になる。

便利ツール

aesthetic-predictor

CLIP-based-NSFW-Detector

LAION-5B-WatermarkDetection

データセット

データセットは画像 URL とキャプションとのペアで提供される。画像のライセンスは画像の所有者ごとに異なる。以下のリストのライセンスが書いていないデータセットの画像のライセンスは画像所有者依存。

商用利用可能な大規模なデータセットは flicker ぐらいしかない。flicker は CC0 で画像を提供すれば無限のストレージが利用できる。

よく使われるデータセット

conceptual-12m

ImageNet：商用不可

Segment Anything 1B（Meta）：Apache 2.0 License だが商用利用不可

AI 生成画像（多様性があり学習に使うと性能が上がる）

DiffusionDB：CC1.0

JourneyDB：商用不可。商用利用したければプロンプトから自作すればよい。$100万を超える売り上げのある企業は有料プランに加入しないと資産を所持できない。Midjourney は著作権侵害コンテンツを平気で生成してくるので、プロンプトと生成した画像との厳選は必要になる。

CC0

soa-full-florence2 はスミソニアン協会が公開している情報をもとにmadebyollin氏が作成したCC-0の絵画などを集めた画像リンク集。古いものが多く、メインのデータセットにはならない。

flicker

flicker Creative Commons

megalith-10m（flicker 画像のキュレーション）

次点

Open Images Dataset V7（Google）：画像は CC BY 2.0 だが個々の画像のライセンスには Google は関知しない、キャプションは CC BY 4.0

LAION-AESTHETICS

Recap-DataComp-1B

DOCCI Descriptions of Connected and Contrasting Images

LLaVA-OneVision-Data

SAM-LLaVA-Captions10M

ALLaVA-4V

DenseFusion-1M

BLIP3o

Unsplash：Unsplash ライセンス（商用可能、転売は不可）

LAION-5B：低品質な画像が多く使われなくなってきている

データの水増し

torchvision.Transform のリスト

OmniGen2 は動画から画像を切り出して、学習に使用している。

Magpie: Alignment Data Synthesis from Scratch by Prompting Aligned LLMs with Nothing。LLM のケースだが、AI にデータを作成させる方法の紹介。

データの分類

OmniGen2 ではデータを以下のように分類し、均等になるようにしている。

• キャラ
• 男
• 女
• 子供
• 老人
• アニメキャラ
• オブジェクト
• バーチャル
• 食べ物
• 人工物
• 植物
• 動物
• シーン
• コマーシャル
• ランドマーク
• 地方（Domestic）
• 都市（Urban）
• 自然（Natural）

The Delta Learning Hypothesis

品質差のあるモデルの出力を DPO を利用して学習させることで、性能が向上できる可能性がある。

キャプショニング

データセットが提供するキャプションを使うのではなく、LLM で画像のキャプション付けを行うのが一般化している。ローカルで実行する場合はパラメータ数 70B のような高性能なものを使用することが多い。VILA-3B/13B や InternVL2-8B/26B もよく使われる。

解像度を下げる場合は、キャプションを作り直す必要はない。しかし画像をクロップした場合は、キャプションを作り直す。

LLM で以下の情報を抽出する。単に「車」と記述するのではなく、「通りを走る車」のようにアクションを記述させる。

• 人
• 人数
• 人の表情・感情
• ポーズ・姿勢：例：「走っている子供」、「ベンチに座っている人」
• 身体属性：例：「長い金髪の女性」、「ふわふわの毛の犬」
• 年齢・身長・性別
• 相互作用：例：「握手している二人」、「ボールで遊んでいる犬」
• 物
• 物の画面内での位置
• オブジェクト間の相対的位置
• 物色、パーツの色
• 形状
• 材質
• サイズ
• 時間帯
• 照明条件：例：「明るい昼光」、「薄暗い夕方の光」
• 季節
• 天候
• 背景の詳細: 例：「都市の景観」、「穏やかな森」
• 服の種類
• 動物の種類
• 室内か屋外か
• 全体的な雰囲気/ムード：例：「賑やかな通り」、「平和な庭」
• 文字の位置と文字

キャプションが複数ある場合は、キャプションをクリップスコアで評価し、正規化することで確率的にキャプションを選択する手法がある。

c は CLIP スコアで、t は温度。０に近い値にすると常に最も高いスコアのキャプションが選択される。

長すぎるプロンプト問題

DetailMaster: Can Your Text-to-Image Model Handle Long Prompts? では長すぎるプロンプトが性能を低下させることを報告している。

Collage Prompting

従量課金の MLLM は画像枚数あたりで課金されるため、Collage Prompting では複数の画像を１枚にまとめてキャプションを作成する方法を紹介している。

プロンプト戦略

参考文献

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?

学習方法

低解像度で学習した後、高解像度の学習を行い、最後に高品質な画像だけを使った仕上げを行う。

８倍圧縮の 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: 学習率の最小値。０に設定するコード例が多いが、計算資源の無駄でしかない
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 前後）に集中して学習させることで３倍学習が高速になる。

ZTSNR（Zero Terminal SNR）

Common Diffusion Noise Schedules and Sample Steps are Flawed では、T=1000 まで学習すべきなのに T=999 で止まるノイズスケジューラーのバグを指摘している。

ノイズスケジューラー

Stable Diffusion 1.5 や XL はスケジューラーにバグがあり、色の再現性の低下や高解像度での分裂などの問題を引き起こしていた。Stable Diffusion 3 ではタイムステップを以下の式で変換することで高解像度での分裂等に対処している。色の再現性は ZTSNR や Rectified Flow を採用することで改善されている。

\(\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 で止まるノイズスケジューラーのバグを指摘している。

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)

ファインチューン

マルチモーダルモデルのファインチューンは 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 を廃止して、データセット作成コストを削減。ノイズから生成された画像以外の画像も使うことで教師を超える性能を実現。

知識蒸留の拡張

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 コア使用。学習が最大２倍高速化する効果もある）

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

GPUメモリ使用量を削減する

勾配の計算には各ブロックのフォワードパスの結果を保存しておく必要がある。しかしこれはメモリ使用量が大きい。そこで、チェックポイントに指定したブロックのみフォワードパスの結果を保存しておき、バックプロパゲーション時にフォワードパスの結果が必要になると、一番近いチェックポイントからフォワードパスの結果を再計算する。

メモリ使用量が減るが、学習時間が増加する。

バッチサイズ

大きい場合

利点

• 学習時間の短縮
• 勾配推定の安定性。勾配が均されるため、ノイズが減り学習が安定する

欠点

• 汎化性能の低下。 勾配のノイズが少ないため、局所最適解・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 のようなモデルが１つのノードに収まらない巨大なモデルは、モデル自体を複数のノードに配置する。

実例

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 １台だと学習に 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） を使用。３段階の訓練でそれぞれ、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 で１日。

512px を 100k ステップ、1,024px を 20k ステップ。詳細は GitHub を参照。

REPA（DINO v2）はマスクしてないトークンにのみ適用。

データセット。（）は使用枚数。

• Segment Anything 1B（10M）
• CC12m（7M）
• DiffusionDB（14M）
• JourneyDB（4M）

AI 生成画像を使いすぎると、画風がカートゥーン化する。JourneyDB は品質がよいので、ファインチューニングには DiffusionDB ではなく JourneyDB を使う。

性能検証

論文で使われるベンチは基本的にプロンプトと生成された画像との忠実性を計測する。以下の要素があっても減点されることはない。

• 指の本数を間違える
• 細部が溶けている
• 画質が悪い
• 直線がゆがんでいる
• 文字がおかしい
• オブジェクトの細部がおかしい
• 水平線がずれている

ベンチ

• GenEval
• DPG-Bench（ELLA）
• ImageReward
• VisionReward：ImageReward の動画版
• T2I-CompBench
• T2I-CompBench++

高速化

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 の４倍速く学習できる。

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

パラメータを共有することで、性能低下を抑えてパラメータ数を削減。

Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models

パッチ単位で評価すると同時に、パッチサイズを変更することで２倍以上の学習速度を達成。この手法の一番の利点は、学習画像の枚数が最低 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(K^TV)) に置き換える。

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% 高速になり、モデル全体では８%程度高速になる。ただし 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 畳み込みを８bit 整数で量子化

SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models

ウェイトには外れ値がある。この外れ値が量子化の精度を低下させる。そこで、行列を外れ値の部分 W と量子化しやすい部分 X とに分け、元の行列を W + X で表現する。量子化しやすい X は４bit 量子化、外れ値の部分は 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

PyTorchでの学習・推論を高速化するコツ集

サンプラー

PFDiff: Training-Free Acceleration of Diffusion Models Combining Past and Future Scores

参考文献

The Illustrated Transformer

OmniGen

OmniGen2

Lumina-Image 2.0

Janus: Decoupling Visual Encoding for Unified Multimodal Understanding and Generation

SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers GitHub

SANA 1.5: Efficient Scaling of Training-Time and Inference-Time Compute in Linear Diffusion Transformer

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

SegFormer

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。離散化誤差と近似誤差という２つのサンプリング誤差を同時に考慮することで、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 を入れることで高い学習率でも学習が安定するようになり、その結果、速い学習と局所的最適解の回避とが実現している。

拡散モデルのミニ実装

tiny-diffusion

transformer_latent_diffusion

A Gentle Introduction to Diffusion

特異値分解（SVD）

『ゼロから作る Deep Learning 2』斎藤 康毅

SVD（特異値分解）解説

scipy.sparse.linalg.svds（アルゴリズムはARPACK（Augmented Implicitly Restarted Lanczos Bidiagonalization Algorithm）を使用）で疎行列の特異値分解ができる。

次元削減目的なら sklearn.decomposition.TruncatedSVD も使える。

そのほか

拡散モデルの時系列

マスクテンソルの作成

block_count は１辺のブロック個数。

import torch
import torch.nn.functional as F
image_size = 512 # block_count で割り切れる必要がある
block_count = 16 # １辺のブロック個数
hide_rate = 0.5 # 数値が増えるとより多くマスクされる
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 = F.interpolate(mini_mask, size=None, scale_factor=scale_factor, mode='nearest-exact')

B = 1   # バッチサイズ
C = 3   # 任意のチャンネル数
mask = mask.expand(B, C, image_size, image_size)