【文献メモ】σReparam: アテンションの崩壊を防ぐ再パラメータ化 (行列ノルム分離) (ICML 2023)

σReparam とは、Transformer のすべての線形層の重み行列を、「行列の大きさ」と、「行列の向き」に分離して訓練する再パラメータ化です。具体的には、各線形層の重み $W$ を、そのスペクトルノルム $\sigma(W)$ で割り、学習可能なスカラ $\gamma$ を掛けたものに置き換えます。これにより、重みのスペクトルノルムが訓練中に際限なく育つのを抑え、アテンションのエントロピーの崩壊を防ぎます。この手法は、Transformer の訓練の不安定さがアテンションエントロピーの崩壊と共に起こる、という分析に基づいています。2023 年の ICML 2023 に採択された論文 Stabilizing Transformer Training by Preventing Attention Entropy Collapse で提案されました。

以下に文献概要と、簡単に数値を確認するコードを記します。

参考文献

Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Josh Susskind. Stabilizing Transformer Training by Preventing Attention Entropy Collapse. In Proceedings of the 40th International Conference on Machine Learning (ICML 2023), 2023.

1. arXiv: Stabilizing Transformer Training by Preventing Attention Entropy Collapse
2. PMLR: Stabilizing Transformer Training by Preventing Attention Entropy Collapse
3. GitHub: apple/ml-sigma-reparam
4. Apple ML Research: Stabilizing Transformer Training by Preventing Attention Entropy Collapse
5. arXiv: Spectral Normalization for Generative Adversarial Networks
6. arXiv: Entropy-based Attention Regularization Frees Unintended Bias Mitigation from Lists
7. ACL Anthology: Entropy- and Distance-Regularized Attention Improves Low-Resource Neural Machine Translation

文献概要

そもそもアテンションエントロピーとは

• アテンションは、クエリ $Q = XW_Q$、キー $K = XW_K$ からロジット $QK^\top/\sqrt{d}$ (アテンションロジット) を作り、これを行ごとに softmax して重み行列 $A$ を得ます ($A_{ij}$ はクエリ $i$ がキー $j$ に置く重み)。
  • $A = \mathrm{softmax}!\left(\dfrac{QK^\top}{\sqrt{d}}\right) = \mathrm{softmax}!\left(\dfrac{XW_Q W_K^\top X^\top}{\sqrt{d}}\right)$
• アテンションエントロピーとは、この $A$ の各行のエントロピー (とその行平均) です ($T$ は系列長)。
  • $\mathrm{Ent}(A_i) = -\sum_{j=1}^{T} A_{ij} \log A_{ij}, \quad \mathrm{Ent}(A) = \dfrac{1}{T}\sum_{i=1}^{T}\mathrm{Ent}(A_i)$
  • エントロピーはアテンションの散らばりの指標です。アテンションが一様なら最大 $\log T$、1 点に集中すると 0 に近づきます。

アテンションエントロピーの崩壊

• Zhai らは、このアテンションエントロピーを訓練中ずっと追跡しました。すると、アーキテクチャやタスクを問わず「アテンションエントロピーが極端に低い (注意が過度に集中している) 状態」と「訓練の不安定さ (損失の振動や発散)」が同時に現れる、という共通のパターンが見られました [1]。この低エントロピー状態を論文は entropy collapse と呼びます。

エントロピーの下界とロジットのスペクトルノルム

• Zhai らは、アテンションエントロピーのタイトな下界を導きました。その下界は、アテンションロジットのスペクトルノルム $\bar\sigma$ に対して指数のオーダー ($e^{-\bar\sigma}$) で 0 に近づきます [1]。
  • ここで $\bar\sigma$ は、重み側の $\sigma(W_K W_Q^\top)$ と入力側の $\sigma(XX^\top)$ を合わせたロジットの大きさの尺度です ($\sigma(\cdot)$ はスペクトルノルム)。
• つまり、重みのスペクトルノルムが大きくなってロジットのスケールが増すほど、エントロピーは崩壊しうる、ということです。ロジットが大きいほど softmax が鋭くなり注意が一点に集中する、という直感に対応します。
• ここから「重みのスペクトルノルムの増大を抑えれば entropy collapse を防げるはず」という σReparam の動機が出てきます。

σReparam

• σReparam は、Transformer のすべての線形層の重み $W$ について、次のように再パラメータ化します。
  • $\widehat{W} = \dfrac{\gamma}{\sigma(W)}, W$
    • $\sigma(W)$ は $W$ のスペクトルノルムで、GAN の学習安定化で用いられたスペクトル正規化 [5] と同じく、べき乗法で推定します (順伝播ごとに 1 反復だけ回して訓練全体で償却)。
    • $\gamma$ は層ごとに 1 つだけ持つ学習可能なスカラです。$\sigma(W)$ で割ることで $\widehat{W}$ の作用素ノルムは常に $\gamma$ に固定され、大きさの自由度を $\gamma$ 1 つに集約します。
• こうすると、生の重み $W$ のノルムが訓練中に育っても、実際に使われる $\widehat{W}$ のスペクトルノルムは $\gamma$ のまま暴走しません。結果としてロジットのスケールが抑えられ、アテンションエントロピーの崩壊が防がれます。
• 追加の正規化層や warmup とは独立で、既存の線形層を置き換えるだけで使えるのが利点です。

検証実験

• 画像分類・画像の自己教師あり学習・機械翻訳・音声認識・言語モデリングの 5 種類のタスクで検証され、σReparam がハイパーパラメータの選び方に対して安定性・頑健性をもたらすことが報告されています [1]。
• 特に、ViT を warmup・weight decay・LayerNorm・適応的最適化 (Adam など) のいずれも使わずに、標準的な学習法と同等の精度まで訓練できたこと、機械翻訳で深いアーキテクチャを発散させずに学習できたこと、音声認識で warmup なしに学習できたことが示されています [1]。

補足: 他の研究におけるアテンションエントロピー

• Zhai らの本研究は、訓練の不安定さを避けるために、アテンションエントロピーが 0 に崩壊するのを防ぐ立場でした [1]。他方、アテンションエントロピーを上げるべきか下げるべきかは、研究の目的によって分かれます。
  • Attanasio らは、アテンションが少数のトークンに集中しすぎてエントロピーが低くなると、特定語に過適合してバイアスが生じるとして、エントロピーを上げる正則化を提案しました [6]。
  • 逆に、低資源の機械翻訳を扱う研究では、アテンションが広く散ってエントロピーが高いと対応が曖昧になるとして、エントロピーを下げてアテンションを鋭くする正則化が提案されています [7]。

補足: σReparam を利用するには

• σReparam そのものは、RAdam の torch.optim.RAdam のような形では PyTorch に組み込まれていません (専用のレイヤーやパッケージは標準にはありません)。
• ただし核となるスペクトル正規化 [5] は組み込みで、torch.nn.utils.parametrizations.spectral_norm が使えます (べき乗法で計算。旧 torch.nn.utils.spectral_norm は非推奨)。σReparam はこれに学習可能なスカラ $\gamma$ を 1 つ足すだけなので、自作は容易です。
• 既製の実装としては、公式の apple/ml-sigma-reparam [3] (PyTorch・JAX) と、非公式の PyTorch 実装 ywchan2005/sigma-reparam-pytorch があります。

補足: ロジットのスケールとエントロピーの関係を見るだけのコード

以下は論文の手法を実装したものではありません。訓練は一切せず、重みを定数倍しただけの順伝播を 1 回計算して、ロジットのスケールとアテンションエントロピーの向きの関係を眺めるだけのコードです。示せるのは「ロジットを大きくすれば softmax が尖ってエントロピーが下がる」「$\widehat{W} = W/\sigma(W)$ はスケール不変なので σReparam 側は $s$ を変えても何も動かない」という、ほぼ定義どおりの当たり前の対応だけで、訓練中の崩壊やその防止そのものを再現するものではありません。

• $T = 16$ トークン・次元 $d = 32$ のランダムな入力 $X$ を固定します。
• クエリ・キーの重み $W_Q, W_K$ を用意し、訓練で重みノルムが育つ状況の代わりに、単に係数 $s$ を掛けて重みを大きくするだけです ($s$ を上げるほど重みが大きい)。
• 各 $s$ について、plain (生の重みをそのまま使う) と σReparam (各重みを $\gamma/\sigma(W)$ 倍、$\gamma = 1$) の 2 通りで、ロジットの作用素ノルム $\sigma(W_Q W_K^\top)$ と、行平均のアテンションエントロピー $\mathrm{Ent}(A)$ を見ます。

import math
import torch


def attn_entropy(logits):
    # 行ごとに softmax したあと、各行のエントロピーを求め、行平均を返す
    A = torch.softmax(logits, dim=-1)
    ent = -(A * A.clamp_min(1e-12).log()).sum(-1)
    return ent.mean().item()


def spectral_norm(W, iters=100):
    # べき乗法でスペクトルノルムを近似する
    u = torch.randn(W.shape[0])
    for _ in range(iters):
        v = W.t() @ u
        v = v / v.norm()
        u = W @ v
        u = u / u.norm()
    return (u @ (W @ v)).item()


def main():
    torch.manual_seed(0)
    T, d = 16, 32
    X = torch.randn(T, d)
    Wq = torch.randn(d, d) / d ** 0.5
    Wk = torch.randn(d, d) / d ** 0.5

    def logits(Wq, Wk):
        Q, K = X @ Wq, X @ Wk
        return Q @ K.t() / d ** 0.5

    print(f'T={T},  log(T)={math.log(T):.4f}  (アテンションが一様なときのエントロピー)')
    print(f'{"s":>4} | {"plain: sigma":>12} {"Ent":>7} | {"sReparam: sigma":>15} {"Ent":>7}')
    for s in [1.0, 1.5, 2.0, 3.0, 4.0, 6.0]:
        # 訓練の不安定化で重みノルムが育つ状況を s 倍で再現する
        Wq_s, Wk_s = s * Wq, s * Wk
        # plain: 生の重みをそのまま使う
        sig_p = spectral_norm(Wq_s @ Wk_s.t())
        ent_p = attn_entropy(logits(Wq_s, Wk_s))
        # sReparam: 各線形層を (gamma / sigma(W)) W に置換する (gamma = 1)
        Wq_r = Wq_s / spectral_norm(Wq_s)
        Wk_r = Wk_s / spectral_norm(Wk_s)
        sig_r = spectral_norm(Wq_r @ Wk_r.t())
        ent_r = attn_entropy(logits(Wq_r, Wk_r))
        print(f'{s:>4} | {sig_p:>12.3f} {ent_p:>7.4f} | {sig_r:>15.3f} {ent_r:>7.4f}')


if __name__ == '__main__':
    main()

実行すると以下のようになります。

T=16,  log(T)=2.7726  (アテンションが一様なときのエントロピー)
   s | plain: sigma     Ent | sReparam: sigma     Ent
 1.0 |        2.416  2.3512 |           0.684  2.7349
 1.5 |        5.436  1.4887 |           0.684  2.7349
 2.0 |        9.663  0.8743 |           0.684  2.7349
 3.0 |       21.743  0.4034 |           0.684  2.7349
 4.0 |       38.653  0.2137 |           0.684  2.7349
 6.0 |       86.970  0.0859 |           0.684  2.7349

読み取れることは以下です。

• plain : $s$ を上げると、重みのスペクトルノルム $\sigma(W_Q W_K^\top)$ が $s^2$ に比例して 2.4 から 87 へ増え、アテンションエントロピーは 2.35 から 0.09 へ急落します。一様なら $\log T \approx 2.77$ なので、アテンションエントロピー崩壊が起きています。
• σReparam : $\widehat{W} = W/\sigma(W)$ はスペクトルノルムが 1 に固定され、$s$ を掛けても変わりません。そのため再パラメータ化後のスペクトルノルム $\sigma(\widehat{W}_Q \widehat{W}_K^\top)$ は 0.684、エントロピーは 2.73 と、どちらも $s$ に依らず一定です。生の重みが育っても、エントロピーは一様な $\log T \approx 2.77$ の近くに保たれ、崩壊しません。
• これは「アテンションエントロピーはロジットのスペクトルノルムに支配される」という下界の帰結を順伝播だけで見たものです。実際の訓練では、重みのスペクトルノルムが育つのを σReparam が抑えることで、この崩壊を防ぎます。