T-Fixup とは、Transformer の残差接続にくるまれた各ブロックのパラメータの初期値を層数 $N$ に応じて小さくすることで、LayerNorm と warmup を不要にする初期化手法です。2020年の ICML 2020 に採択された論文 Improving Transformer Optimization Through Better Initialization で提案されました。

以下に文献概要とトイコードを記します。

参考文献

Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving Transformer Optimization Through Better Initialization. In Proceedings of the 37th International Conference on Machine Learning (ICML 2020), 2020.

文献概要

本記事で「ブロック」とは、残差接続にくるまれた部分 (マルチヘッドアテンション層や FFN) を指します。原論文では "residual branch" と呼ばれますが、以下ではブロックと呼びます。

以下の不安定性の分析、および T-Fixup が更新の大きさを抑えられるという主張は、原論文の Theorem 3.1 (SGD 更新の大きさの上界) に基づきます。私はその証明を追っていません。

オリジナルの Transformer (Post-LN) の訓練には、LayerNorm と学習率の warmup が経験的に必要で、どちらか一方でも外すと学習が発散しやすいことが知られていました。しかし、なぜ必要なのかは自明ではありませんでした。
Huang らは、この不安定性の原因を、Adam の更新が層をまたいで増幅されることにあると分析しました。
- 訓練初期は Adam の 2 次モーメントの推定がまだ不安定で、更新幅が大きくなりがちです。この大きな更新が、積み重なった残差ブロックを通じて増幅され、1 ステップの更新による出力・損失の変化量が深さとともに増大しうることを示しました。
- warmup は訓練初期の学習率を下げてこの大きな更新を抑え、LayerNorm は活性を正規化して抑えている、という解釈です。
そこで Huang らは、各ブロックのパラメータの初期値を小さくして、1 ステップの更新による出力変化を深さ $N$ によらず $\Theta(1)$ に抑えることができれば LayerNorm も warmup も不要になると考え、T-Fixup 初期化を提案しました。具体的に以下です ($N$ は層数)。
- 入力埋め込み以外の全パラメータを Xavier 初期化する。
- 入力埋め込みを $\mathcal{N}(0, d^{-1/2})$ のガウス分布で初期化する ($d$ は埋め込み次元)。
- エンコーダ: 各 MHA ブロックの $V$ 変換と出力変換、各 FFN ブロックの重み、および入力埋め込みを $0.67 , N^{-1/4}$ 倍する。
- デコーダ: 各 MHA ブロックの $V$ 変換と出力変換、各 FFN ブロックの重み、および入力埋め込みを $(9N)^{-1/4}$ 倍する。
ポイントは、マルチヘッドアテンション層のうち $Q$ 変換・$K$ 変換はスケールしないことです。$Q, K$ は softmax で正規化されるためブロックの出力の大きさに直接効かず、出力の大きさに効く $V$ 変換・出力変換と FFN の重みだけを縮小します。
- 指数が $-1/4$ である直感は、各ブロックが 2 つの重み行列 ($V$ 変換と出力変換、あるいは FFN の 2 層) を通ることにあります。各行列を $N^{-1/4}$ 倍すると 1 つのブロックの寄与は $N^{-1/2}$ 程度になり、$N$ 個のブロックを足しても全体で $\Theta(1)$ に収まります。
検証実験は機械翻訳ベンチマークで行われました。
- LayerNorm も warmup もなしで学習でき、精度もベースラインと同等以上でした。
- さらに、エンコーダ・デコーダともに 200 層 (合わせて 1000 超の MHA/FFN ブロック) の超深層 Transformer も、特別な工夫なく訓練できました。

トイコード

以下は、LayerNorm のないブロックからなるエンコーダ層を $N$ 層積んだだけの簡易モデルで、Xavier 初期化と T-Fixup 初期化 (エンコーダ側の $0.67 N^{-1/4}$ 倍) とで、順伝播した出力の大きさが深さとともにどう変わるかを確かめるトイコードです。

各ブロックは、マルチヘッドアテンション層を線形化した残差 ($x \leftarrow x + W_o W_v x$) と FFN 残差 ($x \leftarrow x + W_2 , \mathrm{relu}(W_1 x)$) からなります。

import torch


def xavier(*shape):
    return torch.nn.init.xavier_uniform_(torch.empty(*shape))


def build_layers(N, d, tfixup):
    torch.manual_seed(0)
    scale = 0.67 * N ** -0.25 if tfixup else 1.0
    layers = []
    for _ in range(N):
        Wv = xavier(d, d)
        Wo = xavier(d, d)
        W1 = xavier(4 * d, d)
        W2 = xavier(d, 4 * d)
        # T-Fixup: V/出力変換と FFN の重みだけを縮小 (Q/K に相当する部分はなし)
        layers.append((Wv * scale, Wo * scale, W1 * scale, W2 * scale))
    return layers


def forward(x, layers):
    for Wv, Wo, W1, W2 in layers:
        x = x + x @ Wv.T @ Wo.T  # MHA 残差 (softmax を平均で線形近似)
        x = x + torch.relu(x @ W1.T) @ W2.T  # FFN 残差
    return x


def main():
    d = 32
    torch.manual_seed(1)
    x0 = torch.randn(64, d)  # 64 サンプル、次元 d
    in_norm = x0.norm(dim=1).mean()

    print('出力ノルム / 入力ノルム')
    print(f'{"N":>5} | {"Xavier":>14} | {"T-Fixup":>14}')
    for N in [6, 12, 200]:
        out_x = forward(x0, build_layers(N, d, tfixup=False))
        out_t = forward(x0, build_layers(N, d, tfixup=True))
        rx = (out_x.norm(dim=1).mean() / in_norm).item()
        rt = (out_t.norm(dim=1).mean() / in_norm).item()
        print(f'{N:>5} | {rx:>14.4e} | {rt:>14.4e}')


if __name__ == '__main__':
    main()

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

出力ノルム / 入力ノルム
    N |         Xavier |        T-Fixup
    6 |     1.5594e+01 |     1.1502e+00
   12 |     2.9281e+02 |     1.1698e+00
  200 |            nan |     1.1337e+00

Xavier 初期化では層数 $N$ が増えると出力ノルムが指数的に増大し、$N=200$ では発散します。一方 T-Fixup 初期化では $N$ が増えても出力ノルムがほぼ一定に保たれます。