Zロス（ゼットロス）

Zロス（Z-loss）とは

Z-lossは、LLMの学習において出力ロジット（またはMoEのルーターロジット）の絶対値が過大になることを防ぐための補助損失（auxiliary loss）です。Googleが2022年にPaLMモデルで導入し、大規模モデルの学習安定化技術として広く採用されています。

数式

Z-loss = α × (1/B) × Σ_i log²(Σ_j exp(z_ij))

- z_ij: i番目サンプルのj番目ロジット値
- B: バッチサイズ
- α: 係数（典型値: 1e-4〜1e-2）

log(Σ exp(z))はlog-sum-expで、これを最小化することでロジットの最大値付近を抑制します。

なぜロジット爆発が起きるのか

Transformerの学習において、ソフトマックスの前のロジット値が学習が進むにつれて大きくなる傾向があります：

• 勾配更新の累積でパラメータのスケールが増大
• 長い学習ステップでAdamオプティマイザの分母が追いつかない
• 数値精度（bfloat16）の限界に達して演算が不安定化

Z-lossはこの「ロジットドリフト」を正則化項として抑制します。

MoEモデルにおけるZ-loss

Mixture-of-Experts（MoE）モデルでは、ルーター（どのエキスパートにトークンをルーティングするか決める部分）のロジットが崩壊しやすいため、Z-lossが特に重要です。

Switch Transformer、ST-MoE、DeepSeek-MoE等の主要MoEモデルで採用されており、ルーター崩壊（全トークンが同一エキスパートに集中する問題）の防止に効果的です。

実装例

def compute_z_loss(logits, alpha=1e-4):
    # logits: [batch, seq_len, vocab_size]
    log_z = torch.logsumexp(logits, dim=-1)  # [batch, seq_len]
    z_loss = torch.mean(log_z ** 2)
    return alpha * z_loss

# 総損失 = CE loss + Z-loss
total_loss = ce_loss + compute_z_loss(logits)

QK-Normとの違い

• QK-Norm: Attention Q/K ベクトルを正規化（magnitude制御）
• Z-loss: 出力/ルーターロジットを正則化（ペナルティ）
• logit soft-capping: Attentionスコア/出力ロジットをtanhで制限

3手法は独立した問題を解決しており、組み合わせて使用することが多いです。

係数αの設定指針

• 小さすぎる（1e-5未満）: 効果が薄い
• 適切（1e-4〜1e-3）: 安定化効果あり、損失への影響小
• 大きすぎる（1e-2超）: モデル性能の低下

PaLMではα=1e-4が報告されており、この値がデフォルトとして広く用いられています。