正規化なしのTransformers技術

• Transformerで必須のように使われてきた Layer Norm/RMSNorm の代わりに Dynamic Tanh（DyT）を入れても、従来の正規化モデルと同等またはそれ以上の性能を出せることを示している
• DyTは DyT(x) = tanh(αx) 形式の 要素ごとの演算 で、Layer NormalizationがTransformer内でしばしばtanhに似たS字型の入出力マッピングを作るという観察から出発している
• 実装はPyTorchで数行レベルであり、学習可能な alpha、weight、bias により tanh(alpha * x) の出力にスケールとバイアスを適用する
• 評価は ViT, ConvNeXt, MAE, DINO, DiT, LLaMA, wav2vec 2.0, HyenaDNA, Caduceus まで、ビジョン・言語・音声・DNAシーケンスモデリング全般を網羅している
• 別途 ハイパーパラメータチューニングなしでも 複数の設定で正規化ベースの対応モデルと同等またはより良い結果が得られ、正規化層が必須だという前提を見直させる

---

Dynamic Tanhが変えるポイント

• DyT はTransformerブロックのLayer NormまたはRMSNormを置き換えるシンプルな層である
• 中核となる演算は要素ごとに適用される DyT(x) = tanh(αx) である
• 正規化層を除いたTransformerでも、従来の正規化Transformerと性能が同等かそれ以上になる可能性を示している
• アイデアの出発点は、Transformer内のLayer Normalizationがよく作る入出力関係が scaled tanh 関数に似ているという観察である

実装方法

• DyTモジュールはPyTorchで短く実装できる

class DyT(nn.Module):
    def __init__(self, num_features, alpha_init_value=0.5):
        super().__init__()
        self.alpha = nn.Parameter(torch.ones(1) * alpha_init_value)
        self.weight = nn.Parameter(torch.ones(num_features))
        self.bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x):
        x = torch.tanh(self.alpha * x)
        return x * self.weight + self.bias

• alpha は学習可能なパラメータで、初期値は 0.5 に設定される
• weight と bias も学習可能なパラメータで、tanh(alpha * x) の出力に適用される

Layer Normalizationから得られた観察

• Transformerの Layer Normalization（LN） は、scaled tanh関数に近い入出力マッピングを生成する
• 初期層ではこのマッピングは概ね線形に近い
• 深い層に進むほど、tanh関数の特徴である S字型の曲線 がより明確に現れる
• 観察対象にはVision Transformer（ViT）、音声Transformerモデルであるwav2vec 2.0、Diffusion Transformer（DiT）の選択されたLN層が含まれる

評価範囲と結果

• DyTは複数のアーキテクチャとタスクで評価されている
  • 教師あり学習のビジョン: ViT, ConvNeXt
  • 自己教師あり学習のビジョン: MAE, DINO
  • 拡散モデル: DiT
  • 大規模言語モデル: LLaMA
  • 自己教師あり学習の音声: wav2vec 2.0
  • DNAシーケンスモデリング: HyenaDNA, Caduceus
• すべての場合で、DyTを適用したTransformerは正規化ベースの対応モデルと同等またはより良い性能を示した
• 評価範囲は認識から生成、教師あり学習から自己教師あり学習、コンピュータビジョンから言語モデルまで幅広い

参考資料

• Download Paper: 研究の全詳細を含む論文
• View on GitHub: 実装の詳細を確認できるリポジトリ
• View Summary: 研究結果の簡単な要約
• Transformers without Normalization は CVPR 2025 論文として採択されている

[GN⁺]

Hacker News のコメント

• 本当なら、かなり良い漸進的な改善です。モデル性能を大きく引き上げるものではなさそうですが、現在の最先端 LLM の多くが使っている RMSNorm より計算コストが低いため、学習がより速く、安価になる可能性があります
  • ただし Transformer 全体の計算量の中で RMSNorm が占める割合はかなり小さいです。通常、縮約演算は前後の演算と融合できます
  • 自分の Transformer 学習ベンチマークに先ほど適用してみましたが、結果は非常に期待外れでした。RMSNorm を使った場合よりもはるかに遅く収束しました
    alpha の調整はほとんど効果がなく、かなりのハイパーパラメータ調整や、より精密な初期化が必要かもしれません。PyTorch のデフォルト初期化と直交初期化の両方を試しましたが、差はありませんでした
    あるいは、自分が使っているスカラー最適化器との相性が悪いのかもしれません。Adam より収束を速めるカスタムのスカラー最適化器を使っていますが、DyT 層では Adam と同程度にしか見えませんでした
    もしくは、数十億トークン後になってようやく追いつく方式なのかもしれませんが、そこまで長くテストする予算はありません
• float8 のような低精度形式を使う場合、通常は正規化の前に活性値を BF16 に上げる必要があります。そのため、より低い精度へ進むほど、正規化層が占める計算割合は大きくなります
  こうした層を置き換えられるなら、計算コストをかなり大きく下げる助けになります
• 詳細は読んでみる必要がありますが、正規化の除去は大きな意味を持つ可能性があります。新しいアーキテクチャを試すとき、ネットワークが適切に正規化されるよう調整するのはいつも面倒です
  tanh にも別の影響はあるはずです。正規化は時に条件付けの問題を解決しているからです。それでも、選択肢が増えるのは歓迎です
• では、勾配消失はもう問題ではないのでしょうか？
  • 層を適切に初期化すれば、深いネットワークでも勾配の大きさが消失したり爆発したりしないように保てます。たとえば各層の出力の平均を 0、標準偏差を 1 に合わせると、勾配も妥当な範囲に収まります
    Kaiming He らによる元の ResNet 論文と、その後続論文をおすすめします
    RNN に対する現代的なアプローチとしては、DeepMind の https://arxiv.org/abs/2303.06349 を読む価値があります
    核心は、最大固有値、つまりスペクトル半径が 1 付近であるべきだということです。線形変換を繰り返し適用しても、活性値が大きくなったり小さくなったりしないという意味です
  • ResNet が勾配消失はほぼ解決したと見てよいでしょう。勾配爆発は通常、適切なパラメータ初期化と正規化で扱います。この論文は正規化の代替を提案している形です
  • 良い質問です。それは tanh を活性化関数として使っていた時代、そして残差接続と正規化層が登場する前の問題でした。別の活性化関数と残差接続がある状態で、tanh を正規化のように使うのは問題なさそうです
  • 図にあるように、Transformer は残差を学習します。y = x + f(x) という形です
• 私だけそう見えるのか分かりませんが、論文のグラフは LNinput と LNoutput を比較する際に、tanh(a*x) の後にも重みとバイアスを付けているように見えます
  類似性を見るなら、LayerNorm の出力から重みとバイアスを取り除いた結果と比較すべきではないかと思います
  最終結果が良ければ問題ないのでしょうが、実際に置き換えている部分だけを切り出して見ると、何が起きているのかをより理解しやすい気がします
  • 実装を見ると、tanh を計算した後に重みとバイアスを適用しているようです
• 実用上はどんな意味があるのでしょうか？
  • 要旨によると、DyT を入れると、正規化なしの Transformer はほとんどハイパーパラメータ調整なしでも、正規化された対応モデルと同等かそれ以上の性能を出せる可能性があります