LLaMAのアーキテクチャを論文から読み解く
概要
LLaMAは、7Bから65Bの4種のモデルを公開データのみで学習し、GPT-3(175B)やPaLM-540Bに匹敵する性能を達成したモデルである。核心的な主張は「パラメータ数を増やすよりも、小さいモデルをより多くのトークンで長く学習したほうが推論コストの観点で有利」ということ。
これにより、一般向けGPUやMacBook等のローカル環境でLLMを動かす道が開け、現在のオープンソースLLM生態系の原点となった(※初代LLaMAは研究用途限定の非商用ライセンスで公開。商用利用が条件付きで解禁されたのはLLaMA 2から)。
プロフィール
いつできたの?
- LLaMA 1: 2023年2月
- LLaMA 2: 2023年7月
誰が作ったの?
- 開発元:Meta AI Research
- 開発国:アメリカ
LLaMAの核心的主張
Chinchillaスケーリング則とLLaMAの視点
GPT-3に代表されるように、2020年ごろのLLM開発は「計算予算が増えたらモデルを大きくすべき」という考え方が主流だった(Kaplan et al., 2020)。
Chinchillaスケーリング則とは
2022年にDeepMindが出した論文「Training Compute-Optimal Large Language Models」(Hoffmann et al.)の主張。一定の計算予算に対して、モデルサイズとデータ量は同じ比率でスケールすべきだとした。具体的には、パラメータ数の約20倍のトークン数で学習するのが学習コスト最適であるとした。
実証として、70BパラメータのChinchillaを1.4Tトークンで学習させると、4倍の280BパラメータであるGopherを300Bトークンで学習した場合より高性能になった。大きさではなく「十分なデータで学習したか」が重要だった。
推論コスト最適化という新たな視点
Chinchillaは「学習コストの最小化」に対する最適解を求めた。しかしLLaMAは「推論コストの最小化」という別の問いを立てた。
モデルは一度学習すれば終わりだが、推論は何百万回も繰り返される。だから学習に多少余分にコストがかかっても、小さいモデルを長く学習して同じ性能を出す方がトータルでは安い。
具体的には、Chinchillaの法則に従えば7Bモデルの最適学習量は約140Bトークンだが、LLaMAは7Bモデルを1Tトークン(約7倍)で学習した。結果として13Bモデルが175BのGPT-3を主要ベンチマークで上回る性能を達成し、推論時にはGPT-3の1/10以下のコストで動作する。
アーキテクチャ詳解
全体構造:Decoder-Only Transformer
元祖Transformer(Vaswani et al., 2017)はEncoder-Decoder構成だった。Encoderが入力を理解し、Decoderが出力を生成する、翻訳タスク用の設計である。
しかしGPT系が示したように、「次のトークンを予測する」という言語モデルのタスクにはDecoder部分だけで十分であることがわかった。Encoderを取り除くことでアーキテクチャがシンプルになり、スケーリングもしやすくなる。LLaMAもこの系譜にある。
Transformer Blockの処理フロー
各ブロックは以下の順序で処理を行う:
- 入力にRMSNormを適用(Pre-normalization)
- Masked Multi-Head Self-Attentionを計算(Q, KにRoPEを適用)
- Attentionの出力と入力の残差接続
- 残差接続の結果にRMSNormを適用
- SwiGLU Feed Forward Networkを計算
- FFNの出力と中間結果の残差接続
この1ブロックを32回繰り返し(7Bの場合)、最後にFinal RMSNormとLinear output layerで語彙サイズ(32,000)の確率分布を出力する。
Causal Mask
Decoder型LLMでは、推論時に未来のトークンを参照できないようにする必要がある。Attentionスコア計算時に、現在位置より後ろのトークンへのスコアを-∞に設定することで、位置iのトークンは位置0〜iまでの情報のみを参照する。これがcausal(因果的)maskと呼ばれるものである。
Pre-normalizationの意義
元祖TransformerはPost-norm(サブレイヤーの後に正規化)だったが、LLaMAはPre-norm(サブレイヤーの前に正規化)を採用している。残差接続のパスに正規化が挟まらないため、勾配がスムーズに流れ、深いモデルの学習が安定する。GPT-3以降、事実上すべてのLLMがこの構成を採用している。
RMSNorm
なぜ正規化が必要なのか
層を深く重ねていくと各層の出力の分布がズレていき、勾配の爆発・消失が発生する。正規化は各層でこのズレをリセットする役割を持つ。
LayerNormの仕組み
Transformerで標準的に使われてきたLayerNorm(Ba et al., 2016)は4つの処理を行う: 平均μの計算、平均の引き算、分散σで割る正規化、γとβによるアフィン変換。
数式:
LayerNorm(x) = γ · (x - μ) / √(σ² + ε) + β
RMSNormへの改良
RMSNorm(Zhang & Sennrich, 2019)はLayerNormから平均に関する演算をすべて取り除いたもの。RMS(二乗平均平方根)で割り、γでスケールするだけ。
数式:
RMSNorm(x) = γ · x / √((1/n) Σ xᵢ²)
LayerNormとの差分は3つ:
- 平均μの計算が不要
- 平均の引き算が不要
- シフトパラメータβが不要(学習パラメータが半分)
Zhang & Sennrich (2019)の実験により、正規化で重要なのは「スケールを揃える」ことであり、「平均を0にする」ことではないと示された。
なぜLLaMAはRMSNormを選んだか
- 計算効率: 演算ステップが少なく高速
- 性能: LayerNormとほぼ同等の結果
- 配置: GPT-3で有効性が示されたPre-normalization(サブレイヤーの前に配置)と組み合わせ、学習安定性と計算効率を両立
SwiGLU活性化関数
標準FFNからの進化
元祖TransformerのFFNは「Linear → ReLU → Linear」の2層構造で、入力を一度大きい次元(4d)に射影してReLUで非線形化し、元の次元に戻すだけのシンプルな構成だった。
SwiGLU(Shazeer, 2020)はここに「ゲート機構」を導入した。入力のすべての要素を同じように通すのではなく、どの情報を通してどの情報を抑えるかを学習で決める。
数式:
SwiGLU(x) = (SiLU(xW₁) ⊗ xW₃) W₂
ここで SiLU(z) = z · σ(z) (σはシグモイド関数)
3つのLinear層がそれぞれ異なる役割を持つ:
- W₁(gate): ゲートの生成。SiLUで活性化される(4,096 → 11,008)
- W₃(up): 値の変換。活性化なしでそのまま通る(4,096 → 11,008)
- W₂(down): ゲーティング後の結果を元の次元に戻す(11,008 → 4,096)
入力が2つのパスに分岐し、片方がゲート、片方が値となり、element-wise multiplyで合流する。
隠れ層次元が 2/3 × 4d である理由
標準FFNはLinear 2本(d→4d→d)で合計パラメータ数は8d²。SwiGLUはLinear 3本あるため、同じ隠れ層次元(4d)だと12d²になってしまう。
パラメータ数を揃えるために、3dh = 8d²を解くと h = 8d/3 となり、隠れ層次元を8d/3に縮小する。d=4,096の場合、8×4,096/3 ≈ 10,922を128の倍数に丸めて11,008となる。
SiLUの利点
ReLUは負の入力を完全に0にするため、一度死んだニューロンは復活しない。SiLU(z · σ(z))は負の値も小さく通すため勾配が完全に0にならず、学習が安定する。
RoPE(Rotary Position Embedding)
従来手法との違い
- 絶対位置(GPT-3等): トークンごとに固定のベクトルを加算する。実装は簡単だが、学習時より長い文章が入ると性能が劣化する。
- 相対位置(T5等): トークン間の距離をアテンション計算に組み込む。長文に強いが、計算が複雑で処理が重い。
- RoPE(LLaMA): 絶対位置として埋め込むのに、計算すると自然に相対位置に化ける。両者のいいとこ取りをしたハイブリッド手法。
仕組み:回転のメカニズム
トークンの埋め込みベクトルを、トークンの位置(インデックス)に応じた角度だけ2次元平面上で回転させる(数学的には複素数の掛け算と等価)。
位置mのトークンに対して、次元iのペアには角度 m × θᵢ(θᵢ = 1 / 10000^(2i/d))を適用する。次元ごとに異なる周波数を持つことで、さまざまなスケールの位置関係を表現できる。
この回転させたQueryとKeyで内積(アテンション計算)を行うと、三角関数の加法定理によって式が整理され、2つのトークンの相対位置(差分)だけに依存する形に分解できる。
採用された3つの利点
- 長文への耐性: 位置情報が角度で表現されているため、後から角度の周波数を調整(RoPE Scaling)することで、学習時より長い文脈への拡張が比較的容易にできる。ただし実用上は少量のファインチューニングを組み合わせることが多い。
- 直感的な位置把握: トークン同士がどれくらい離れているかをモデルが自然に理解できるため、文脈の依存関係を正確に捉えられる。
- 計算効率の高さ: 理論上は行列の掛け算だが、実際の実装では符号反転と要素ごとの掛け算で実現できるため、絶対位置埋め込みや相対位置埋め込みと比べて計算コストが低い。
学習データと学習手法
公開データへのこだわり
LLaMA以前の巨大LLM(GPT-3など)は、一部に非公開のプライベートデータ(独自の書籍データや対話ログなど)を含んで学習されていた。
これに対し、LLaMAは「誰もが再現・検証できること」を重視し、公開されているデータのみで構成されている。合計約1.4Tトークンで、内訳はCommonCrawl(67%)、C4(15%)、GitHub(4.5%)、Wikipedia(4.5%)、Books(4.5%)、ArXiv(2.5%)、StackExchange(2%)。トークナイザはSentencePieceによるBPEで、数値は1桁ずつ分割する。
高速化と安定化の工夫
1.0T〜1.4Tという膨大なトークンを小さなパラメータサイズに詰め込むには、莫大な計算時間がかかる。Metaはこれを解決するため、いくつかのテクニックで学習を高速化した。
アテンション計算の効率化
FlashAttentionに代表される効率的なアテンション計算手法を採用した(実装はxformersライブラリ)。通常のアテンション計算ではGPUのHBM(高帯域幅メモリ)上に巨大なアテンション行列をすべて展開する必要があるが、この手法ではGPUのSRAM(高速だが小容量のオンチップメモリ)を活用し、タイル単位で分割して計算することでHBMへのアクセス回数を大幅に削減する。
勾配チェックポインティングの最適化
通常の逆伝播では、すべての層の中間アクティベーション(計算結果)をメモリに保持しておく必要があり、深いモデルではメモリを大量に消費する。勾配チェックポインティングはあえて中間アクティベーションを保存せずに捨てておき、逆伝播時に必要になったタイミングで再計算する手法。計算量はやや増えるが、メモリ消費を大幅に削減できる。
オプティマイザの構成
アダプティブな学習率制御を行うAdamW(β₁=0.9, β₂=0.95)を採用。コサイン形式の学習率スケジュールを適用し、学習の終盤に向けて滑らかに学習率を落とすことで、重みが安定して収束するように制御した。
結果と意義
推論効率と性能の両立
13Bという軽量サイズでありながら、当時の代表的モデルだったGPT-3(175B)を主要ベンチマークで上回った。パラメータの「大きさ」よりも「十分な学習量」こそがLLMの真のポテンシャルを引き出すことを実証した。
LLaMA-65BはChinchilla-70BやPaLM-540Bと同等の性能を示した一方、MMLUベンチマークではやや弱い結果となった。これは学習データ中の書籍・学術論文の割合が少ないためであり、学習データの構成がドメイン知識の深さに直結することも示している。
ローカルLLM文化の誕生
モデルが軽量(7B/13B)になったことで、一般のコンシューマ向けGPUやMacBook等のローカル環境でLLMを動かす道が開けた。これにより、世界中の開発者が独自のファインチューニングを競い合う、現在のオープンソースLLM生態系の原点となった。
モデルサイズごとのスペック
| モデル | dim | heads | layers | 学習トークン |
|---|---|---|---|---|
| 7B | 4,096 | 32 | 32 | 1.0T |
| 13B | 5,120 | 40 | 40 | 1.0T |
| 33B | 6,656 | 52 | 60 | 1.4T |
| 65B | 8,192 | 64 | 80 | 1.4T |
参考文献
- Touvron et al. (2023). “LLaMA: Open and Efficient Foundation Language Models.” arXiv:2302.13971
- Hoffmann et al. (2022). “Training Compute-Optimal Large Language Models.” arXiv:2203.15556
- Vaswani et al. (2017). “Attention Is All You Need.” arXiv:1706.03762
- Zhang & Sennrich (2019). “Root Mean Square Layer Normalization.” arXiv:1910.07467
- Shazeer (2020). “GLU Variants Improve Transformer.” arXiv:2002.05202
- Su et al. (2021). “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv:2104.09864
- LLM Architecture Gallery by Sebastian Raschka: https://sebastianraschka.com/llm-architecture-gallery/