2. LLaMAのモデルアーキテクチャ

Transformerデコーダの標準構成とLLaMAの設計方針

LLaMA の内部アーキテクチャは、GPT-3 などと同様の「デコーダ専用（Decoder-only）Transformer」を基本としています。しかし、オリジナルの Transformer（Attention Is All You Need, 2017）の設計をそのまま使うのではなく、その後の深層学習研究で提案された優れた改良手法を貪欲に取り入れました。

具体的には、学習の安定性を高めるための「RMSNorm」、非線形表現力を高める活性化関数「SwiGLU」、長文脈での位置関係を的確に捉える「RoPE（回転位置埋め込み）」などを組み合わせています。これらの改良により、モデルのパラメータ数が同じであっても、学習が破綻しにくく、高い予測精度を発揮できるようになりました。

事前正規化（Pre-normalization）とRMSNorm

オリジナルの Transformer では、各層の処理（アテンションや順伝播ネットワーク）の「後」に正規化（LayerNorm）を適用する Post-normalization が使われていました。しかし、Post-norm は層が深くなるにつれて勾配が爆発または消失しやすく、学習が不安定になる欠点がありました。

LLaMA では、各サブレイヤーの「前」に入力を正規化する Pre-normalization を採用しました。さらに、標準的な LayerNorm よりも計算が軽量な「RMSNorm（Root Mean Square Normalization）」を導入しました。 LayerNorm が平均と分散の両方を計算するのに対し、RMSNorm は二乗平均平方根（RMS）のみを用いて正規化します。平均の計算を省略することで計算速度が約 10% から 50% 向上しつつ、学習の安定性は同等を維持します。

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 入力の二乗平均の平方根を計算して割る
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return (x / rms) * self.weight

活性化関数SwiGLUの採用

標準的な Transformer の順伝播ネットワーク（FFN）では、ReLU や GELU といった活性化関数が使われてきました。 LLaMA では、PaLM などでも高い効果が実証されていた「SwiGLU（Swish Gated Linear Unit）」を採用しました。

SwiGLU は、2 つの線形変換の出力のうち一方に Swish（SiLU）活性化関数をかけ、もう一方との要素ごとの積（ゲート処理）を取る構造です。パラメータ数が通常の FFN より約 1.5 倍増えるため、LLaMA では隠れ層の次元数を元の $\frac{8}{3}d$ から 256 の倍数に調整することで、計算量とパラメータのバランスを最適化しています。

import torch
import torch.nn as nn
import torch.nn.functional as F

class SwiGLU(nn.Module):
    def __init__(self, in_features: int, hidden_features: int):
        super().__init__()
        self.w1 = nn.Linear(in_features, hidden_features, bias=False)
        self.w2 = nn.Linear(hidden_features, in_features, bias=False)
        self.w3 = nn.Linear(in_features, hidden_features, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SiLU(w1(x)) と w3(x) の要素ごとの積を w2 で元の次元へ戻す
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

回転位置埋め込み（RoPE: Rotary Position Embedding）

単語の順序をモデルに認識させるため、従来の Transformer は絶対位置埋め込み（各位置に固有のベクトルを足す方式）を使用していました。しかし、絶対位置埋め込みは未知の長さの文章への汎化が苦手でした。

LLaMA は、2021 年に Jianlin Su 氏が提案した「RoPE（Rotary Position Embedding）」を採用しました。 RoPE は、トークンの埋め込みベクトルを 2 次元ごとの組に分け、その位置インデックスに応じた角度で回転行列を掛ける手法です。 2 つのトークンの内積を計算すると、2 つの単語の「相対的な距離」だけが自然に残る数学的性質を持ちます。これにより、位置関係の捉え方が極めて柔軟になり、後のコンテキスト拡張（2048 トークンから数万トークンへの延長）でも重要な役割を果たすことになります。

推論を高速化するKVキャッシュの原理

LLM が文章を生成するとき、単語（トークン）を 1 つずつ順番に出力する自己回帰（Autoregressive）の仕組みを取ります。愚直に計算すると、新しいトークンを 1 つ出すたびに、それまでに生成されたすべての過去トークンの Key と Value の行列を毎回ゼロから再計算することになり、膨大な無駄が発生します。