アテンション+MLP、繰り返し
2つのサブレイヤー
Transformerブロックは正確に2つのサブレイヤーを含み、それぞれが形状[batch, seq_len, d_model]のトークンシーケンスに対して動作します:
1. マルチヘッド注意サブレイヤー。 トークンが互いを見ます。Activity 68 でこれを詳しく扱いました。出力形状は入力形状と一致します。
2. フィードフォワード MLP サブレイヤー。 各トークンが独立して2層パーセプトロンで変換されます。トークン間の情報フローはありません。出力形状は入力形状と一致します。
両方のサブレイヤーは [batch, seq_len, d_model] 形状を保持します。この保持により、ブロックを積み重ねることが可能:レイヤー N の出力がレイヤー N+1 の入力に形状の変更なしで供給されます。
各サブレイヤーがもたらすもの
注意は情報を位置間で移動させます:位置 17 のトークンが位置 1 から 16 の情報を引き出せます。MLP は情報を各位置内で変換します:トークンの表現が学習された非線形関数で再形成されますが、隣接トークンは見えません。
ブロックの積み重ねはこれら2つの操作を交互に行います。レイヤー1の注意機構は位置を混合します。レイヤー1のMLPは位置ごとの表現を再形成します。レイヤー2の注意機構は再び混合しますが、今度は再形成された表現に対してです。この交互は深さとともに表現力を成長させます。
ANDREAのスタック
| Variant | n_layer | n_head | d_model | mlp_dim |
|---|---|---|---|---|
| ANDREA-12M | 6 | 12 | 384 | 1536 |
| ANDREA-120M | 12 | 12 | 768 | 3072 |
| ANDREA-480M | 16 | 24 | 1536 | 6144 |
ファミリー全体で mlp_dim = 4 × d_model に注意してください。この比率はほぼすべての現代のトランスフォーマーで維持されています。セクション3でその理由を詳しく説明します。
サブレイヤーの命名
スキップ接続が重要な理由
残差パターン
各サブレイヤーはその計算を残差接続で包みます。出力は入力を加算します:
x = x + Attention(LayerNorm(x)) # attention sublayer
x = x + MLP(LayerNorm(x)) # MLP サブレイヤー
各サブレイヤー内では、Attention(...) または MLP(...) の関数がデルタを生成します。ブロックは入力を置き換えません。学習された修正を加えます。
これが重要な理由
残差接続が深層アーキテクチャで支配的な3つの理由:
1. 勾配の流れ。 バックプロパゲーション中、追加により勾配は出力から入力への直接経路を得て、サブラヤーを迂回します。残差なしの12層スタックでは、埋め込みに到達する前に勾配信号が失われますが、残差により数百層を通しても勾配の大きさが使用可能に保たれます。
2. Identity initialization. 初期化時、サブラヤーの重みは小さな出力を生成します。残差接続により、層Nは最初ほぼ変更せずに通過します。トレーニングは動作する開始点から徐々にデルタを学習します。
3. Composition学習。 各ブロックは置き換えではなく洗練を学習します。層1は位置情報を追加するかもしれません。層2は統語構造を追加。層7は意味的関係を追加。残差ストリームが蓄積します。
層正規化
各サブラヤーの前で、LayerNorm は各トークンの表現を平均0、分散1に再スケーリングし、その後学習された特徴ごとのゲインとバイアスを適用します:
y = gamma * (x - mean(x)) / sqrt(var(x) + epsilon) + beta
平均と分散は d_model 次元にわたって計算され、各トークンごとに別々に行われます。2つの学習されたベクトル(gamma、beta、それぞれ形状 [d_model])が表現力のあるスケールを回復します。正規化は活性化を数値的に安定した範囲に保ちます。それがないと、訓練中の小さな不安定さが雪だるま式に増大し、NaN勾配になります。
Pre-Norm vs Post-Norm
微妙だが重要な選択
レイヤーノーマライズを残差サブラayerに組み込む2つの方法:
Post-norm (元の2017年論文):
x = LayerNorm(x + Attention(x))
レイヤーノーマライズは残差加算の後に配置されます。残差ストリーム自体が各層で正規化されます。
Pre-norm (現代標準、ANDREAで使用):
x = x + Attention(LayerNorm(x))
レイヤー正規化はサブレイヤーの前にあり、残差ブランチ内に位置します。残差ストリームは正規化されず、サブレイヤーへの入力のみが再スケーリングされます。
Pre-Normが勝った理由
Post-normはLRウォームアップと慎重なハイパーパラメータ調整なしではうまく訓練されません。初期層で勾配が爆発するのは、各レイヤー正規化が残差ストリームの蓄積された状態を乱すためです。元の2017年の論文は広範な調整を伴うpost-normを使用しましたが、その後の研究(GPT-2、LLaMA、ANDREA)はpre-normを標準化しました。
Pre-norm は安定して訓練されます。残差ストリームはすべての層でクリーンに蓄積されます;サブレイヤーの入力のみが数値安定性のために正規化されます。現代のトランスフォーマーはデフォルトで pre-norm を使用し、& ANDREA もその選択を継承します。
最終ブロック方程式
残差、pre-norm 位置のレイヤー正規化、& 両サブレイヤーを組み合わせると、ANDREA の完全なブロックが得られます:
```python
def block_forward(x):
```
x = x + Attention(LayerNorm(x)) # attention サブレイヤー
x = x + MLP(LayerNorm(x)) # MLP サブレイヤー
return x
2つのサブレイヤー、2つの残差加算、2つのレイヤー正規化(注: 各サブレイヤーは独自のレイヤー正規化を持つ; ANDREA-120M は12ブロックで24のレイヤー正規化+出力の最終1つで合計25)。12回繰り返す。これが ANDREA-120M の幹となる。
なぜ Pre-Norm が訓練を安定させるのか
2つの線形層、1つの活性化関数
3つの操作
MLP サブレイヤーは、層の間に非線形活性化関数を持つ2層パーセプトロンです:
def mlp_forward(x):
h = x · W_1 + b_1 # 拡張: d_model → mlp_dim
h = GELU(h) # 非線形活性化
y = h · W_2 + b_2 # 縮小: mlp_dim → d_model
return y
3つの操作。2つの線形、1つの非線形。最初の線形は幅を拡張し、2番目は縮小する。
4倍の拡張比率
現代のトランスフォーマーは mlp_dim = 4 × d_model を設定する。ANDREA-120M:
d_model = 768
mlp_dim = 4 × 768 = 3072
W_1 形状 = [768, 3072] # ~2.36M params
W_2 形状 = [3072, 768] # ~2.36M params
ブロックあたりの MLP パラメータ = 4.72M (バイアスを無視)
すべての注意サブレイヤーのペアの間に2つのMLPが配置されます(ブロックあたり1つ)。12ブロック × 4.72M ≈ 56.6M MLPパラメータがANDREA-120M全体で、すべてのパラメータの約半分です。
なぜ4倍か
4倍の比率は経験的に現れました。小さな比率はモデル容量を減少させます。大きな比率はパラメータあたりの収束が減少します。数十年にわたるアーキテクチャ検索において、4倍が維持されてきました。それはGPT、BERT、T5、LLaMA、& ANDREAに現れています。
最近の研究(PaLM、Chinchilla)では、ゲーティング機構(SwiGLU)が8/3倍の拡張を使用でき、同等の容量を少ないコストで実現できることがわかりました。ANDREAはシンプルさのためにクラシックなGELU + 4倍を維持します。
GELU:滑らかな活性化関数
GELUが計算するもの
GELU(Gaussian Error Linear Unit)は、現代のトランスフォーマーのMLP層間の標準的な活性化関数です。その公式:
GELU(x) = x · Φ(x)
Φ(x) は標準正規分布の累積分布関数です:標準ガウス乱数が x 以下である確率。数値的に計算されます:
Φ(x) ≈ 0.5 × (1 + tanh(sqrt(2/π) × (x + 0.044715 × x³)))
領域ごとの振る舞い
- 大きな正の x に対して: Φ(x) ≈ 1 なので、GELU(x) ≈ x。ReLU に似ています。
- 大きな負の x に対して: Φ(x) ≈ 0 なので、GELU(x) ≈ 0。ReLU に似ています。
- x = 0 付近: Φ(x) ≈ 0.5 なので、GELU(0) = 0 ちょうど。原点を通る滑らかな遷移。
ReLU と異なり、GELU は Φ(x) で重み付けされた一部の負の入力を通します。小さな負の入力は、まだ小さな負の出力を寄与しますが、入力全体ほどではありません。
GELUがReLUを上回った理由
実証的に、同じパラメータ数でGELUで訓練されたトランスフォーマーは、ReLUで訓練されたものよりも低い損失に到達します。ゼロ付近の滑らかさが重要です:ReLUのゼロでの急激なカットオフは勾配の不連続性を生じさせますが、GELUの滑らかな曲線はバックプロパゲーションのためのよりクリーンな勾配を提供します。
ANDREAの訓練エンジン microgpt_cuda.cu は、手書きのGELU CUDAカーネルを同梱しています。このカーネルは上記の tanh 近似を使用します;現代のGPUは tanh を単一命令オペレーションとして含んでいます。
MLPパラメータの計算
12個のブロックが ANDREA-120M を構成
ブロックからモデルへ
ANDREA-120M の完全なフォワードパス:
def model_forward(token_ids):
x = token_embed(token_ids) + position_embed(positions)
for block_idx in range(n_layer): # 12 blocks
x = block_forward(x) # attention + MLP w/ residuals
x = LayerNorm(x) # final norm
logits = x · token_embed.T # tied weights for output projection
return logits
6つのライン。主な処理はblock_forwardの中にあり、12回呼び出される。埋め込みがパイプラインを開始し、結合された非埋め込み(入力ルックアップに使用される同じ行列を転置して出力投影に使用)がそれを終了する。
深さとしての構成
各ブロックは残差ストリームを読み、δを計算し、それを追加する。12回の通過後、ストリームにはすべてのブロックからの累積寄与が含まれる。内部的に、層は専門化する傾向がある:
- 初期層 (1-3): 統語的パターン、位置構造
- 中間層 (4-8): 単語の関係、句の境界
- 後期層 (9-12): 意味的内容、事実の想起
この特殊化は、訓練の圧力から生じ、建築上の選択によるものではありません。同じ均一なブロック設計が、言語で訓練されると特殊化された層を生み出します。
ブロック全体のパラメータ数
| コンポーネント | ブロックあたり | 12ブロック合計 |
|---|---|---|
| Attention projections (4×W) | 2.36M | 28.3M |
| MLP 重み (W_1 + W_2) | 4.72M | 56.6M |
| レイヤー正規化 (gamma, beta) | ~3K (無視可能) | ~37K |
| ブロックごとの合計 | ~7.1M | ~85M |
トランクに85Mパラメータ。トークン埋め込みに~13Mを追加(8449語彙 × 768 d_model × 2、入力/出力共有)および最終レイヤー正規化を加えると、ANDREA-120Mのパラメータ数はおよそ120Mとなる。ブロック設計が3分の2を占め、残りが埋め込み。