形状付きトランスフォーマー：無限の深さと幅の限界における注意モデル

要旨

深層学習理論において、表現の共分散行列はネットワークの学習可能性を検証するための代理指標として機能する。Transformerの成功に着想を得て、我々は無限深度・無限幅の比例極限における、スキップ接続を備えた修正版Softmaxベースの注意機構モデルの共分散行列を研究する。初期化時において、その極限分布は深度対幅比によってインデックス付けされた確率微分方程式（SDE）によって記述可能であることを示す。明確な確率的極限を達成するため、Transformerの注意機構を修正し、Softmax出力を単位行列を中心に配置し、Softmaxのロジットを幅に依存する温度パラメータでスケーリングする。対応するSDEを通じてネットワークの安定性を検証し、ドリフトと拡散のスケールが残差接続の助けを借りてエレガントに制御可能であることを示す。安定したSDEの存在は、非常に大きな深度と幅においても共分散構造が良好に振る舞うことを意味し、深層注意モデルにおける悪名高いランク縮退問題を防ぐ。最後に、シミュレーションを通じて、SDEが対応する有限サイズモデルを驚くほど良く記述することを示す。我々はこれらのアーキテクチャ修正を「shaped Transformer」と名付ける。

English

In deep learning theory, the covariance matrix of the representations serves as a proxy to examine the network's trainability. Motivated by the success of Transformers, we study the covariance matrix of a modified Softmax-based attention model with skip connections in the proportional limit of infinite-depth-and-width. We show that at initialization the limiting distribution can be described by a stochastic differential equation (SDE) indexed by the depth-to-width ratio. To achieve a well-defined stochastic limit, the Transformer's attention mechanism is modified by centering the Softmax output at identity, and scaling the Softmax logits by a width-dependent temperature parameter. We examine the stability of the network through the corresponding SDE, showing how the scale of both the drift and diffusion can be elegantly controlled with the aid of residual connections. The existence of a stable SDE implies that the covariance structure is well-behaved, even for very large depth and width, thus preventing the notorious issues of rank degeneracy in deep attention models. Finally, we show, through simulations, that the SDE provides a surprisingly good description of the corresponding finite-size model. We coin the name shaped Transformer for these architectural modifications.

形状付きトランスフォーマー：無限の深さと幅の限界における注意モデル

The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit

要旨

Support