T2I 拡散モデルの設計メモ

目次

• テキストエンコーダー
• VAE
• VAE の学習
• LPIPS
• GAN
• コンディショニング
• 位置埋め込み
• 目的関数
• Transformer アーキテクチャ
• ブロック図
• データセット
• キャプショニング
• 学習方法
• タイムステップスケジューラー
• ノイズスケジューラー
• 蒸留
• 省メモリ学習
• 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

なぜ KL VAE を使うのか

潜在空間の統計的制御（規則性）

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

• 同じ意味を持つ画像は近い z にエンコードされる
• ノイズを加えてもデコーダーが画像として再構成しやすい
• 正規分布に従うことで、サンプリング可能な「全体的に埋まった」潜在空間になる

VAE の圧縮率

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

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

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

• P：パッチサイズ

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

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

各 VAE の設定

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

Transformer の隠れ層の次元 >> C なので、既存の VAE の P を大きくすれば、トークン数が減るので拡散モデルの処理速度を向上させられる。

SD-VAE・FLUX.1 VAE のレイヤー情報

Notes / Links about Stable Diffusion VAE

SDXL の VAE のレイヤー情報は sdxl-vae/config.json を参照。

SD3 の VAE や FLUX.1 の VAE にアクセスするには認証が必要。

Diffusers で様々な VAE の実装が見れる。

SD や FLUX は Diffusers の AutoencoderKL を使用しているので以下のコードで簡単にモデルを構築できる。学習するコードも 100 行程度あれば可能。block_out_channels は最終的には [128, 256, 512, 512] を使うことが多い。

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

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

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

D²iT

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

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

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

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

オリジナル

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

差分

CoVAE: Consistency Training of Variational Autoencoders

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

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

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

prior hole problem

VAE の潜在空間（latent space）において、事前分布（通常は標準正規分布）に従ってサンプリングした点が、まともな出力を生成できないことがあるという問題。エンコーダーが潜在空間 z の全体を使用していないことが原因。

対策

• VAE with a VampPrior
• KL loss の重みを適切に調整（例: KL annealing）
• beta VAE の β を高める（GAN や LPIPS を使わないと画像がぼける）
• デコーダーの能力を上げる（ただし posterior collapse が起こりやすくなる）
• Two-stage VAE などの構造的改良（例: hierarchical VAE）

正規化

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

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
• Masked AutoEncoders
• EQ-VAE
• 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 をとる。

以下の式で \( \tau\) はスケール・回転変換、\( \mathcal{E}\) はエンコーダー。\( \mathcal{L}_{rec}\) は再構成 loss。\( \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

Latent Diffusion Models with Masked AutoEncoders

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

• マスクしていない部分の再構成 loss
• マスク部分の再構成 loss
• マスクしていない部分の LPIPS loss
• マスクしていない部分の KL loss

対称形の 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 等で \({\mu}\) と \({\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）

Transformer VAE の設計

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

KL loss 対策

• ステップ数が十分か確かめる。ブロック感をなくすには 50 万ステップは必要
• β-VAE を採用し、ベータを抑え気味にする

パッチ化対策

• 学習可能な位置エンコーディングを使う
• VIVAT: Virtuous Improving VAE Training through Artifact Mitigation
• トランスフォーマーブロックは最低でも８は必要
• パッチ化に畳み込みを使いオーバーラップさせる
• デコーダーの最後に refiner の畳み込みを追加する

self.image_to_token_cnn = nn.Conv2d(
    in_channels=3,
    out_channels=model_dim,
    kernel_size=patch_size,
    stride=patch_size//2,  # stride < kernel_size → オーバーラップあり
    padding=patch_size//4
)

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 も使う。

補間ノイズ \( x' = (1-\tau)x + \tau \epsilon\) と加算ノイズ \( x' = x + \tau \epsilon\) とでは補間ノイズの方が性能がよかった。\( \epsilon(\gamma) \sim \gamma \cdot \mathcal{N}(\mathbf{0},\mathbf{I})\) で \( \gamma = 3\) が最も性能がよかった。加算ノイズはオリジナルのシグナルを損なわないようなショートカットが作成される恐れがある。それにも関わらず VAR では性能が改善されたが、SiT では性能は改善されなかった。

マスク率は 70～90% が性能が高い。マスク率をランダム化した方が性能が良い。

マスクと潜在空間ノイズとを併用すると VAR では性能が上がったが SiT では性能が上がらなかった。SiT の場合は、マスクはしなくてもいい。

エンコーダーをフリーズしてデコーダーを訓練させたところ性能が劣化した。つまり、エンコーダーの能力強化が VAE の性能を上げるうえで重要。

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 の学習や蒸留のときに必要になる。

GAN の学習についてはGAN の学習を参照。

コンディショニング

拡散トランスフォーマーではコンディショニングとして、タイムステップやテキストをトークン化せずに入力する。

しかしコンディショニングをトークンにして入力する方法が 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)

import torch
import torch.nn as nn

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

class AdaLN(nn.Module):
    def __init__(self, embed_dim, style_dim):
        super().__init__()
        self.norm = RMSNorm(embed_dim)
        self.proj = nn.Linear(style_dim, 2 * embed_dim) # Scale and Shift

    def forward(self, x, style):
        scale, shift = self.proj(style).chunk(2, dim=-1)
        scale = scale.unsqueeze(1)
        shift = shift.unsqueeze(1)
        return self.norm(x) * (1 + scale) + shift

ブロック図での表記法

• 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 が詳しい。Flash Attention や xformers と RoPE の組み合わせは self_attention.py を参照。

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

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 は以下のようになる。

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 を使うのが一般的になっている。

今後は Energy-Based Transformers や、拡散モデルとフローマッチングと自己回帰モデルとを一般化した Transition Matching が主流になるかもしれない。

Transition Matching

DTM（Difference Transition Matching）がフローマッチングに相当する。目的関数の部分はフローマッチングとほとんど同じ。

エンコーダー・デコーダーモデル。backbone と呼ばれるエンコーダーが巨大。小規模なデコーダーは head と呼ばれている。

参考文献

Improving and generalizing flow-based generative models with minibatch optimal transport

Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers

Transition Matching: Scalable and Flexible Generative Modeling

Transformer アーキテクチャ

Transformer のトレンドは以下のように変化した

• sinusoidal positional embedding -> RoPE
• Multi-Head Attention -> Grouped-Query Attention
• GELU -> SwiGLU
• LayerNorm -> RMSNorm
• QK-Norm

Grouped-Query Attention

Grouped-Query Attention は Multi-Head Attention の K, V を複数のクエリで共有する。 GQA-4 のケースでは、K と V との数が 1/4 になり、推論速度が 1.3 倍高速になり、メモリ使用量 50% 減少、性能の低下は１～３%。

Dit-Air ではヘッドだけではなく、QKVO をすべてのトランスフォーマーブロックで共有している。

Multi-Head Latent Attention

DeepSeek-V2 で採用された、KV キャッシュを圧縮する手法。

QK-Norm

Multi-Head Attention の Q, K の RoPE 適用前に RMSNorm を入れる。

参考文献

The Big LLM Architecture Comparison

ブロック図

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

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

DiT

FFN

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

Patching（Patchify）

Transformer はトークン数の２乗で計算量が増えるので、パッチ化によってトークン数をチャンネル数に変換することで計算負荷を下げられる。d_model > C * P * P のときにこの戦略は機能する。

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)

# [B, H*W, D_model] -> [B, H*W, C*P*P]
patches = self.reverse_proj(hidden_states)
rearrange(patches, 'b (h w) (c f1 f2) -> b c (h f1) (w f2)', f1=P, f2=P, h=int(H/P), w=int(W/P))

# 畳み込みの場合
image_to_patches = nn.Conv2d(input_channel, model_dim, kernel_size=(P, P), stride=P, bias=True)
patches = image_to_patches(x)
tokens = rearrange(patches, 'b c h w -> b (h w) c')

# unpatch
patches_to_image_deconv = nn.ConvTranspose2d(model_dim, input_channel, kernel_size=(P, P), stride=P, bias=True)
patches = rearrange(tokens, 'b (h w) c -> b c h w', h=int(H/P), w=int(W/P))
patches_to_image_deconv(patches)

2022年
出典：Scalable Diffusion Models with Transformers. William Peebles et al. Figure 3. https://arxiv.org/abs/2212.09748

PixArt-α

位置エンコーディングには DiffFit: Unlocking Transferability of Large Diffusion Models via Simple Parameter-Efficient Fine-Tuning で使用された、 sinusoidal interpolation を使っている。実装は単純で、DiT の解像度は 256 だが DiffFit は 512 なので、DiT の sinusoidal encoding の (i, j) を (i/2, j/2) にしただけだ。

MLP

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

2023年
出典：PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis. Junsong Chen et al. Figure 4. https://arxiv.org/abs/2310.00426

Stable Diffusion 3

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

Stable Diffusion 3(SD3)の理論やモデル構造について解説

2024年
出典：Scaling Rectified Flow Transformers for High-Resolution Image Synthesis. Patrick Esser et al. Figure 2. https://stability.ai/news/stable-diffusion-3-research-paper

Playground v3

発表は 2024/09/16、設計の終了と学習の開始は Stable Diffusion 3（2024/02/22）の論文の発表前。

EDM スケジューラーを使っている。

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

FLUX.1

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

AuraFlow

Hunyuan-DiT

2024年
出典：Hunyuan-DiT: A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding. Zhimin Li et al. Figure 7. https://arxiv.org/abs/2405.08748

MicroDiT

パッチマスクはパッチのランダムドロップより性能が良いのでよく使われる（BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, Masked Autoencoders Are Scalable Vision Learners）。Fast Training of Diffusion Models with Masked Transformers ではマスクパッチの再構成能力を上げるためにオートエンコーダー loss を利用している。これらの従来手法はパッチのマスク率が 50% を超えると性能が低下し始め、75% を超えると大幅に性能が低下する。しかし MicroDiT は 75% のマスク率でもわずかな性能低下に抑えられる。

MicroDiT はパッチマスクを行う前に Patch-mixer でテキスト Embedding を取り込むことで性能を向上させた。Patach-mixer はアテンションと FFN で構成された数レイヤーの Transformer。

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

パッチマスクはブロックサイズを小さくしてランダムにした方が性能低下を抑えられる。

クロスアテンションでキャプションを取り込み、位置エンコーディングとタイムステップは Sinusoidal Embedding。

損失関数。マスクした部分は、元の画像＋ノイズを予測させる。

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

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

Image Text Attention は右の FFN と等価で、キャプションの量を増やせばネットワークのパラメータ数を増やすのと同じ効果がある
出典：Lumina-Image 2.0: A Unified and Efficient Image Generative Framework. Qi Qin et al. Figure 5. https://arxiv.org/abs/2503.21758

RMSNorm

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

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 の後にゲートを入れる場合は、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.」で理由も聞ける。データセットの画像の厳選にも使える。

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

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）

キャプショニング

データセットが提供するキャプションを使うのではなく、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

隠れ層の次元を増やすよりトランスフォーマーブロックを増やした方が性能が上がる。データセットは質と多様性が重要で、量は重要ではない（現状十分な量の画像が手に入る）。

マスク

MaskGIT: Masked Generative Image Transformer ノイズ＋マスクで学習させる。

MADI: Masking-Augmented Diffusion with Inference-Time Scaling for Visual Editing ノイズ＋マスクで学習し、推論時にはパディングトークンを追加することで性能を上げる。

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 前後）に集中して学習させることで３倍学習が高速になる。

ノイズスケジューラー

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

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

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 問題に似ていて、エラーが蓄積していく。

そこで学習時に追いノイズによってこのバイアスを除去する。以下の式の \( \sqrt{1+\gamma^2}\) が追いノイズの部分。

A Comprehensive Review on Noise Control of Diffusion Model

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

• Linear
• Fibonacci
• Cosine
• Sigmoid
• Exponential
• Cauchy distribution
• Laplace distribution
• Logistic
• Monotonic neural network

高解像度では Cosine より Sigmoid の方が良い。

Improved Noise Schedule for Diffusion Training

ノイズスケジューラ―の Laplace, Cauchy は logSNR=0（つまり純粋なノイズ）の時の学習効率と品質がよい。

参考文献

Common Diffusion Noise Schedules and Sample Steps are Flawed

On the Importance of Noise Scheduling for Diffusion Models

Diffusion Noise Optimization (DNO)

損失関数の重みづけ

Debiased Estimation Loss

モデルの再構成誤差に \( \dfrac{1}{\sqrt{SNR(t)}}\) を掛けるだけ。v prediction の場合は、\( \dfrac{1}{SNR(t) + 1} \)。

Rectified Flow の場合

SD3 では以下の式で損失関数を重みづけしてる。ノイズが大きい部分のウェイト大きくなる。T=0, 1 の時ウェイトが０に近い。\( \lambda^{\prime}\) は SN 比の微分。

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

小さいモデルを先に訓練する

SANA 1.5 では 1.6B を作成した後、学習済みの 1.6B にブロックを追加して 4.8B にしてから学習させることで、4.8B の学習時間を 60% 削減している。

蒸留

性能向上の面から見れば小さいモデルは、大きなモデルから蒸留した方が性能はよくなる。

詳細は蒸留技術まとめを参照。

大きいモデルのブロックを削減して小さいモデルを作成する

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 を使う。Tensor コアが使われる。メモリ使用量が減少し、計算が速くなる。

AMP は精度の低下を抑えるために loss をスケールして 16bit の精度の範囲内になるように調整している。なので逆伝播の際にスケールを元に戻す必要がある。

最適化方法は NVIDIA の AMP が遅くなる理由を参照。

分散学習

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 を使う。

Illustrious: an Open Advanced Illustration Model

• タグベース：Danbooru形式などのタグによるプロンプト設計
• マルチレベルキャプション：タグと自然言語キャプションを併用した注釈方式
• 手動フィルタリング／再構成：クリーンデータの選別やタグ整備の工程を含む

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

• 同一トークンに複数の意味が重なることが多く、曖昧性や誤解の原因になる
• 例：「doctor」というタグは職業名としてもキャラクター名としても使われる
• 画像にタグが極端に少ないものが存在する
• 極端に高解像度な画像がある
• 極端なアスペクト比の画像がある
• コミック形式の画像がある

キャプション構造

NovelAI のようなタグオーダーを採用している。

人数 ||| キャラ名 ||| レーティング ||| 一般タグ ||| アーティスト名 ||| パーセンタイルベースのレーティング ||| 年代

最初はパーセンタイルベースではなく、スコアレンジベースのレーティングだった。しかしスコアレンジは年代やカテゴリでばらつきが大きかったので、以下のようなパーセンタイルベースのレーティングを採用することにした。

解像度

v2.0 では 0.15MP（500px * 300px）から 2.25MiP（1536 px * 1536px / (1024*1024) = 2.25）。

訓練方法

• No Dropout Token。性的（provocative）なトークンを含む NSFW トークンはドロップアウトしない。加えて NSFW トークンを CFG で制御（つまりネガティブプロンプトに NFSW トークンを入れる）することで 100% 性的なコンテンツを生成しないモデルを作成できた。
• コサインアニーリング
• 疑似レジスタートークン。データ数が少なかったりモデルが理解できない概念を吸わせるトークンを入れる。シーケンス長を調整するためのパディングトークンがレジスタートークンとして機能してしまうことがあるので注意が必要。
• キャラ名とアーティスト名をドロップアウトしない Contrastive Learning（ラベルなしで学習させる手法）。キャラの特徴とアーティストのスタイルを効率よく学習できる。ただし、アーティストとキャラとの結びつきが強くなる欠点がある。
• キャプションの言い換え。"1girl" を "one girl", "single woman" に一定確率で入れ替える。
• マルチレベルドロップアウト。No Dropout Token 以外のトークンは以下ドロップアウトのどれかが確率的に適用される。
• 30%: max(0.3 * total tokens, 10)
• 20%: max(0.4 * total tokens, 15)
• 10%: min(total tokens, 6)
• 4%: min(total tokens, 4)
• 36%: ドロップアウトなし？（論文に記載なし）
• Input Perturbation Noise Augmentation（0 < ε < 0.1）（Ning et al. [2023]）
• Debiased Estimation Loss（Yu et al. [2024]）

訓練設定

ファインチューニングの問題点

ユーザーの評価をフィードバックさせたファインチューニングは生成される画像の多様性がなくなる。流行の画風が高く評価され、AI が上手く書けない手や背景などを描かなくなる。

Illustrious がファインチューニング前のモデルを公開しているのはこれが理由だ。ファインチューニング版を LoRA 形式で配布する方法も考えられる。

参考文献

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

性能検証

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

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

ベンチ

• 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

NVIDIA の AMP が遅くなる理由を参照。

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 を \({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 は（出力 - 入力）の差分をキャッシュしているので、モデルを動作させなくてもキャッシュを利用したデノイズが行われるが、タイムステップを減らした場合はそうではない。

Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

トークンごとに FFN の通過回数を動的に変更する。MoE は横展開だが、MoR は縦展開。

量子化と枝刈り

SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

FlashAttention-2 を INT 8 bit 量子化する。FlashAttention3 や xformer より高速。Softmax の精度は FP32。

Attention の行列の量子化には以下の課題があった

1. 行列 K はチャンネル単位（channel-wise）の外れ値がある
2. 行列 P, V を単純に INT8 に量子化すると精度がばらつく

SageAttention では以下のようにこれらの課題を解決した。

1. K にスムージングを適用する
2. P, V は INT8 で量子化し、その行列積の計算は FP16 を使う

K の外れ値にはパターンがあり、バイアスが大きいだけで分散は小さい。なので、以下のように K をスムージングすれば INT 8 に量子化しても精度を維持できる。

この変換でアテンションスコア P は変化しない。

\(\tilde{P}\) はブロック単位の量子化、V はチャンネル単位の量子化を適用する。理由は以下の通り。

1. \(\tilde{P}\) にチャンネル単位の量子化と V にブロック単位の量子化とを適用するのは不可能。逆量子化は outer axis のスケールファクターが必要なので
2. \(\tilde{P} = exp(S_i - rowmax(S_i)), S_i \; \text{は}\;QK^\top \;\text{ブロックの行列積の結果}\) なのでそれぞれの行の最大値の \(\tilde{P}\) は１。よって \(\tilde{P}\) ブロックに s= 1/127 を適用すれば精度はトークンごとの量子化と同じ
3. チャンネルごとの量子化は V のチャンネルごとの外れ値に対処可能

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

QK を INT4、\(\tilde{P}, V\) を FP8 で計算する。K だけでなく Q にもスムージングを適用する。

SageAttention2++: A More Efficient Implementation of SageAttention2

PV の計算に、mma.f32.f8.f8.f32 ではなく、mma.f16.f8.f8.f16 を使用して最大４倍（SageAttention2 比）高速化したバージョン。

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を達成している。

GGUF

llama.cpp で実装されている量子化方法。Georgi Gerganov氏（GGMLライブラリの作者）によって、オープンソースプロジェクトとして進化してきたファイルフォーマット。

1. レガシー量子化
• Qx_0 や Qx_1 のような形式で表示
• 各レイヤーの重みを、ブロック（通常256個の重み）に分割し、それぞれのブロックに対して異なるスケール係数を適用して量子化
• _0 はバイアスなし、_1 はバイアスあり
• 基本的には、重み = スケール × 量子化された値 または 重み = 最小値 + スケール × 量子化された値 の形式で、各ブロックの重みが低ビット（例: 4ビット）の整数に変換される
2. K 量子化
• Qx_K_S, Qx_K_M, Qx_K_L のような形式で表示
• llama.cppのPR #1684 で導入された
• レガシー量子化よりもビット割り当ての性能が上がっている
• ブロックごとの定数も量子化されるなど、より効率的なビット割り当てと精度維持とを目指している
3. I 量子化
• IQ2_XXS, IQ3_S などの形式で表示
• PR #4773 で導入された
• QuIP# (Quantization with Incoherence and Preconditioning) 使用

K 量子化

ブロック単位のベクトル量子化

例えば、基本的に 32要素/ブロック（=128bit）のfloat重みを含むブロックを対象に、以下の手順で量子化する

• ブロックごとに min/max または scale/zero_point を求める
• 各値を k-bit 整数にスケーリングして量子化
• 推論時にはブロック単位で復元（dequant）

typedef struct {
    uint8_t qs[16];  // 量子化データ
    float d;         // scale（1ブロックごと）
    float m;         // offset（optional）
} block_q4_K;

実装例：K量子化された行列との乗算

ggml_vec_dot_q4_K のような関数で、量子化行列との内積計算が行う。K-Quant は以下のように量子化されたまま計算するのではなく、「復元しながら計算する」

• 一定のブロック数を並列SIMD演算でデコード
• デコード時に量子化されたint4をfloatに復元（scaleとoffsetを適用）
• ブロック内で再構成された値との内積を行う

S, M, L の違い

M

• ブロック内を 2つのサブブロック（16要素 + 16要素）に分割
• 各サブブロックに 別々の scale 値を持たせる（dual scale）
• 一部の実装では補正ベクトル（微調整量）を使って残差補正を適用

L

• ブロック内を 複数のサブブロック（例：8要素×4）に分割
• 各サブブロックに異なるスケール
• 非線形な正規化（L2正規化やGroupNormに近い）を適用可能
• 残差補正量を別途保存し、復元時に補完
• 高速化最適化が難しくなるが、推論精度は非常に高い。Q6_K と同性能を Q4_K_L で達成することもある

そのほか

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

数学

Mathematics for Machine Learning and Data Science: Optimization with Mathematica Applications

Machine Learning: a Lecture Note

Reinforcement Learning: An Overview

そのほか

拡散モデルの時系列

LLM推論に関する技術メモ

LLM を利用するサービスを提供する場合の技術的注意点について網羅的に解説されている。

マスクテンソルの作成

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

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 # １辺のブロック個数
    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)