ロジットソフトキャッピング（ロジットソフトキャッピング）

ロジットソフトキャッピング（Logit Soft-Capping）とは

Logit soft-cappingは、TransformerのAttentionスコアや最終出力ロジットの値域を「ソフトに」制限する学習安定化技術です。GoogleがGemma 2シリーズ（2024年）で採用して注目され、大規模LLMの長期学習における数値安定性の向上に寄与しています。

動作原理

キャッピングは以下の変換式で適用されます：

soft_cap(x, cap) = cap × tanh(x / cap)

- x: 元のロジット値
- cap: 上限値（Gemma 2ではAttentionに50、出力に30）
- tanh(x/cap): x/cap を [-1, 1] の範囲に圧縮
- 乗算 cap で [-cap, +cap] の範囲に拡張

ハード（クリッピング）との違い: ハードクリッピング clip(x, -cap, cap) は閾値付近で勾配が0になる（勾配消失）のに対し、soft-cappingはtanhのなめらかな変換のため勾配が常に存在します。

Gemma 2での実装

Gemma 2は2段階のsoft-cappingを採用：

1. Attentionレベル（attn_logit_softcapping = 50.0）:

   • QK^T / sqrt(d_k) の後に適用
   • アテンションエントロピー崩壊を防止

2. 出力ロジットレベル（final_logit_softcapping = 30.0）:

   • 最終線形層の後、softmax前に適用
   • 語彙上の確率分布が極端にならないよう制御

# Attentionでのsoft-capping実装例
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
if attn_logit_softcapping:
    scores = scores / attn_logit_softcapping
    scores = torch.tanh(scores)
    scores = scores * attn_logit_softcapping

QK-Normとの組み合わせ

Gemma 2はQK-NormとLogit soft-cappingを同時採用しています：

• QK-Norm: Q/Kを正規化 → Attention計算前
• Soft-capping: スコアの範囲制限 → Attention計算後

二重の安定化により、27Bパラメータ規模でも安定した学習を実現しています。

cap値の選定根拠

• Attentionロジット cap=50: 通常のAttentionスコア範囲（sqrt(d_k) ≈ 8〜16）より十分大きく、大半の学習期間で不活性。発散しそうな異常値のみをキャッチ。
• 出力ロジット cap=30: 語彙サイズ（数万〜10万）に対応したスケールで設定。

パフォーマンスへの影響

• 品質への影響: 正常範囲のロジットはほぼ変更されないため、モデル性能への影響は軽微
• 速度への影響: tanh計算の追加コストは数パーセント程度
• FlashAttentionとの互換性: 標準FlashAttentionカーネルはsoft-cappingに非対応のため、カスタムカーネルが必要（Gemma 2推論速度への制約）

採用状況と今後

現時点でGemma 2特有の技術として知られていますが、学習安定化の観点から他のモデルへの応用が研究されています。FlashAttention v3以降での公式サポートにより、採用モデルが増える見込みです。