dskjal
広告
広告

T2I 拡散モデルの設計メモ

カテゴリ:deeplearning

目次

テキストエンコーダー

CLIP や T5 が長らく使われてきたが、軽量 LLM を使うケースが増えている。

主な LLM と VLM

性能比較

Params(M) は間違いで、数値は FP32 精度の時のファイルサイズ<br/>出典:SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers. Enze Xie et al. Table 9. https://openreview.net/forum?id=N8Oj1XhtYZ

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 を入れることで、潜在空間に意味的なスムーズさや構造性が生まれる。以下のような特徴を持つ潜在空間は拡散モデルにとって扱いやすい。

VAE の圧縮率

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

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

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 の設定

モデルFCP
SD1.5F32C4P2
SANAF32C32P1
SDXLF8C4P2
Dit-AirF8C8P2
SD3・FLUX.1F8C16P2

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 の VAEFLUX.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 を取り込むことで性能を向上させている。

パッチサイズは小さい方が拡散モデルの性能が良い<br/>出典:SANA: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers. Enze Xie et al. Figure 3. https://openreview.net/forum?id=N8Oj1XhtYZ

パッチサイズは小さい方が拡散モデルの性能が良い
出典: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段階の学習で性能を向上させる。

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

GitHub

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

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

オリジナル

オリジナル

VAE(dc-ae-f32c32-sana-1.1-diffusers)通過後

VAE(dc-ae-f32c32-sana-1.1-diffusers)通過後

差分

差分

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に近づく)が起こりやすい。そこで以下のような戦略が良く使われる。

  1. KL Loss の逆伝播の量を調整する kl_beta を設定する。
  2. 最初の数万ステップは kl_beta を0にする。
  3. 次にウォームアップ区間を設ける。 kl_beta が 0→1 になる“ウォームアップ区間”は、実装・データ規模により 1 万〜10 万 SGD ステップ、あるいは 50〜200 エポック程度
  4. 各ピクセルの小さい loss を、ノイズとみなして、フィルターして無視する
  5. 各ピクセルの KL loss の最大値を設定する。max(KLi, threshold)
正規化

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

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. 最尤推定は訓練データとは著しく異なるデータを生成してもペナルティがない

正則化項

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

GitHub

Latent Diffusion Models with Masked AutoEncoders

Masked Autoencoders に KL 項を追加して、潜在空間のスムースさ・知覚的圧縮品質・再構成画像の品質の3つの指標を同時に達成。Loss は以下の4つ。

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

F4C16 のモデルの処理は以下のようになる。Transformer の次元は dmodel とする。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 dmodel
  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 dmodel
  8. Transformer
  9. Linear で隠れ層の次元を調整 (16,384 x dmodel) -> (16,384 x 64)
  10. アンパッチ (16,384 x 64) -> (64 x 128 x 128)->(3 x 512 x 512)
Transformer VAE の設計

Transformer VAE は出力がブロック状になりやすい。原因は、Transformer のパッチ化と KL loss との2つ。

KL loss 対策

パッチ化対策

オーバーラップのサンプルコード
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

mae/models_mae.py

patch_embed.py(畳み込みパッチ)

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 の場合は、マスクはしなくてもいい。

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

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

モード崩壊(Mode Collapse)

モード崩壊(Mode Collapse)とは Generator が同じ画像ばかり出力すること。以下の対策方法がある。

学習戦略・最適化の工夫
ロス関数の変更
モデル構造の工夫
正則化や多様性の強化
そのほか
導入順
  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

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 - GitHubApplication 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.<br/>https://arxiv.org/abs/2503.21758<br/>図の位置の改変は筆者

出典:Qi Qin et al. Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Figure 3.
https://arxiv.org/abs/2503.21758
図の位置の改変は筆者

コンディショニングの方法はいくつかある。

しかしテキストのコンディショニングはどのモデルも、トークン化してノイズ画像と結合するか、クロスアテンションで取り込むかしている。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 はプロンプトやタイムステップ等} } \]

ブロック図での表記法

Unveiling the Secret of AdaLN-Zero in Diffusion Transformer

ゼロではなくガウス分布を利用した初期化をする adaLN-Gaussian を提唱している。

位置埋め込み

SD3 以前は Sinusoidal Encoding 使われてきたが、FLUX 以降は RoPE を使うのが一般的。ViT の RoPE については Rotary Position Embedding for Vision Transformer が詳しい。

SD3 は画像の位置埋め込みに絶対位置を使っている

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)

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) を通る。0を境界に導関数が急に変化する ReLU と違い、SiLU は導関数が滑らかに変化する。

GELU も SiLU も ReLU のような形で導関数が滑らかに変化するように設計されたもの。GELU の方が性能が良いが、SiLU の方が計算が速く、どちらも ReLU より性能が良い

SANA が Mix-FFN を使用しているが、その隠れ層の次元が 5,600 もあり、RoPE を使った方が速い可能性が高い。

目的関数

ノイズを予測させる epsilon prediction から、ノイズ差分を予測させる v prediction へトレンドが移り、2025 年では Rectified FlowFlow 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 のトレンドは以下のように変化した

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回出現するのが基本。

  1. AdaLN
  2. Attention・MLP ・FFN
  3. 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年<br/>出典:Scalable Diffusion Models with Transformers. William Peebles et al. Figure 3. https://arxiv.org/abs/2212.09748

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年<br/>出典:PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis. Junsong Chen et al. Figure 4. https://arxiv.org/abs/2310.00426

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年<br/>出典:Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. Patrick Esser et al. Figure 2. https://stability.ai/news/stable-diffusion-3-research-paper

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年<br/>出典: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

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年<br/>出典: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

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年<br/>出典:Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget Vikash Sehwag et al. Figure 2. https://arxiv.org/abs/2407.15811

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} \]

GitHub

1.16B のモデルの学習に $1,890 のコストしかかかってない。

2024年<br/>出典:Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget Vikash Sehwag et al. Figure 3. https://arxiv.org/abs/2407.15811

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

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

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 と等価で、キャプションの量を増やせばネットワークのパラメータ数を増やすのと同じ効果がある<br/>出典:Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Qi Qin et al. Figure 5. 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

RMSNormLayer 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

2025年<br/>出典: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

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年<br/>出典: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: Efficient High-Resolution Text-to-Image Synthesis with Linear Diffusion Transformers. Enze Xie et al. Figure 5. https://openreview.net/forum?id=N8Oj1XhtYZ

2025年<br/>出典: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

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 * レイヤー数。

モデルレイヤー数隠れ層の次元
S12768
B181,152
L241,536
XL301,920
XXL382,432

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

Pixart と同じ MLP なら、トランスフォーマーブロックの隠れ層と同じ次元を持つ、2レイヤーの MLP。活性化関数は SiLU。

2025年<br/>出典: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

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

Nitro-T

0.6B は PixArt-α に似たアーキテクチャ、1.2B は SD3 に似たアーキテクチャで、ブロック図はない。

MicroDiT で使用された Deferred Patch Masking を採用しパッチ化の性能が上がっている。

使用技術的に SANA に近い。SANA はパッチを使用しないが Nitro-T はパッチを使うところが相違点。

35M 枚数の学習画像の過半数が AI 生成(DiffusionDB と JourneyDB)の画像。

GitHub

データセット

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

最近では多くても 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年著作権法改正により学習は合法。ただしモデルが著作権侵害コンテンツを高頻度で生成する場合、モデルの公開・販売や画像生成サービスの提供は複製権侵害になる。

便利ツール

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 生成画像(多様性があり学習に使うと性能が上がる)

DiffusionDBCC1.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

UnsplashUnsplash ライセンス(商用可能、転売は不可)

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

データの水増し

torchvision.Transform のリスト

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

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

データの分類

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

キャプショニング

データセットが提供するキャプションを使うのではなく、LLM で画像のキャプション付けを行うのが一般化している。ローカルで実行する場合はパラメータ数 70B のような高性能なものを使用することが多い。VILA-3B/13BInternVL2-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}\)  <br/>高解像度の画像の情報を破壊するにはより多くのノイズが必要になる<br/>出典:On the Importance of Noise Scheduling for Diffusion Models. Ting Chen. Figure 2

\(\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

以下のノイズスケジューラーのレビュー。

高解像度では 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\)<br/>\(\lambda^{\prime} = 2(\frac{a^{\prime}_t}{a_t} - \frac{b^{\prime}_t}{b_t})\)<br/>\(\sigma_t = a = t\)<br/>\(b = 1 - a\)<br/>a' と b' はそれぞれ a と b の微分

横軸はタイムステップ、縦軸は\(\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

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メモリ使用量を削減する

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

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

バッチサイズ
大きい場合

利点

欠点

小さい場合

利点

欠点

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

出典: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-12mImageNet

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

出典: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)はマスクしてないトークンにのみ適用。

データセット。()は使用枚数。

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

Illustrious: an Open Advanced Illustration Model

バージョンステップ数バッチサイズデータ枚数プロンプトスタイルアノテーション手法解像度
0.1781,2501927.5Mタグベースオリジナルプロンプト + 手動フィルタリング/再構成1,024
1.0625,00012810Mタグベースオリジナルプロンプト + 手動フィルタリング/再構成1,536
1.193,75051212Mタグベースマルチレベルキャプション1,536
2.078,12551220Mタグベースマルチレベルキャプション1,536

Danbooru 画像とタグは以下の問題がある

キャプション構造

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)。

訓練方法
訓練設定
バージョンデータ枚数バッチサイズ学習率
テキストエンコーダー学習率
エポック解像度プロンプトスタイルドロップアウトレジスタトークンマルチキャプション
0.17.5M1923.5e-54.5e-6201,024タグ
1.010M1281e-56e-681,536タグ
1.112M5123e-54e-641,536タグ+自然言語
2.020M5124e-53e-621,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

出典: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
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

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 TransformersGitHub

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

拡散モデルのミニ実装

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 も使える。

そのほか

拡散モデルの時系列

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)


広告
広告

カテゴリ