- Brian KitanoはTinyShakespeareで縮小版Llamaを自作し、論文実装は小さなモデルから始めて部品を1つずつ差し替え、そのたびに学習・評価するのが安全だと整理している
- データ分割、バッチ生成、損失評価、生成関数といった検証用補助関数を先に用意し、単純なモデルでコンパイルと学習の可否を確認したあとにLlamaの構成要素を追加していく
- RMSNorm、RoPE、SwiGLUを順に組み込みながら、テンソルのshape、数式の性質、アテンションマップで各レイヤーが期待どおり動くかを確認する
- RoPEアテンションでcausal maskを外すと検証損失は0.16まで下がったが、生成品質は悪化し、未来トークンを見てしまう情報漏洩が原因だった
- 最終的な縮小版Llamaは4ブロック、約237万パラメータで検証損失を約1.0まで下げており、gradientの流れと学習率スケジュールも併せて点検する必要がある
小さく始めて反復的に確信を積み上げる
- 論文実装は小さなモデルから始め、構成要素を1つずつ変え、そのたびに学習と評価を繰り返すやり方が重要
- まずモデルを定量的に確認するための補助関数を準備する
- データ分割
- 学習ループ
- 損失の可視化
- 検証損失の評価
- 論文の構成要素を一度にすべて移植するのではなく、すでに実装経験のある単純で高速なモデルで生成結果を見る定性評価関数も用意する
- テンソルレイヤーは
.shape、assert、plt.imshowで確認し、最初から行列積の最適化に入るのではなく、手計算で期待結果を検算してからtorch関数で効率化する
- バッチサイズ、シーケンス長、埋め込み次元を変えながらテストすべきであり、1つのサイズでしか合わないコードは推論時に壊れる可能性がある
データセットと基本設定
- 実装対象はMeta AIのLlamaを大幅に縮小したバージョンで、学習データはTinyShakespeare
- Llamaは1.4Tトークンで学習されているが、ここでは約111万文字規模のTinyShakespeareを使う
- 元のLlamaはSentencePieceのbyte-pair encodingトークナイザーを使うが、この実装では単純な文字単位トークナイザーを使用する
- vocabulary sizeは65
- データセットが小さいため、メモリ保存方式は別途最適化していない
MASTER_CONFIG辞書でvocab_size、batch_size、context_window、d_modelといったモデル設定を管理する
- 定数やマジックナンバーを減らし、コードを読みやすくするのが目的
get_batches関数はデータをtrain 80%、val 10%、test 10%に分け、ランダムな開始点から入力xと1文字後ろのラベルyを生成する
基本モデルでコンパイルと学習を確認
- 最初のモデルは埋め込みと単純なfeed-forwardネットワークで構成された
SimpleBrokenModel
nn.EmbeddingLinearReLULinear
- 論文実装でモデルが「動く」とは、次の2条件を両方満たすことを意味する
- コンパイル: テンソルのshapeがレイヤー間で合っている
- 学習: 損失が実際に下がる
evaluate_loss関数はtrainとval splitで10回バッチをサンプリングし、平均損失を計算するSimpleBrokenModelは1000 epochs学習後も検証損失が3.94程度で、初期cross-entropy 4.17からほとんど下がらなかった- 原因は
F.cross_entropyにすでにsoftmax済みの値を渡していたこと
- PyTorchの
F.cross_entropyは正規化されていないlogitsを直接受け取る
- softmaxを取り除いた
SimpleModelは検証損失を2.51程度まで下げた
- その後
generate関数を追加してモデルが生成した文字を直接確認し、基本モデルは不完全ながらも検証損失が下がる状態になった
Llama構成要素 1: RMSNorm
- Llamaは元のTransformerと比べて3つの主要なアーキテクチャ変更を使っている
- RMSNorm pre-normalization
- Rotary embeddings
- SwiGLU activation function
- 元のTransformerはBatchNormalizationを使うが、Llamaはベクトルをcenteringせず、varianceでスケーリングするRMSNormを使う
- 元のTransformerがattention layer出力にnormalizationを適用するpost-normalization方式なのに対し、Llamaは入力に先に適用するpre-normalization方式を採る
- 実装した
RMSNormは入力shapeを(batch, seq_len, d_model)と仮定する
- RMSNormの結果は、レイヤーノルムがレイヤー要素数の平方根になる性質でテストする
assert- row-wise comparison
torch.allclose
- 基本モデルにRMSNormを追加した
SimpleModel_RMSは検証損失を2.5015程度へわずかに改善した
Llama構成要素 2: RoPEとcausal mask
- RoPEはTransformer向けの位置エンコーディング方式で、トークン位置をembeddingの回転で表現する
get_rotary_matrixはcontext windowとembedding dimensionに対して、位置ごとの回転行列を生成する- RoPE実装は次の性質でテストする
- 位置
m、nで回転した2つのベクトルの内積が、相対位置n-mの回転と一致すること
RoPEAttentionHeadはw_q、w_k、w_vを作り、queryとkeyにRoPE回転を適用したあとF.scaled_dot_product_attentionを使う- 学習時と推論時のテンソルshapeの違いに注意が必要
- 学習時は
(config['batch_size'], config['context_window'], config['d_model'])のように設定どおりであることが多い
- 推論時は
(1, 1, config['d_model'])のような単一サンプルを処理することがある
forward内ではモデル設定値ではなく、入力から得たshapeを基準にインデックスすべき
- causal maskなしでRoPE multi-head attentionを追加したモデルは検証損失が0.1623まで急低下したが、生成結果は
OOOO...、IIII...のように良くなかった
- attention mapを確認すると、すべての位置がすべての位置を参照しており、次トークン予測で未来トークンを見てしまう情報漏洩が起きていた
F.scaled_dot_product_attentionにis_causal=Trueを適用したRoPEMaskedAttentionHeadへ置き換えると、未来に相当するupper triangular attentionはほぼ0になった- causal mask適用後の検証損失は2.0815となり、より長く学習させると1.8985まで下がった
Llama構成要素 3: SwiGLUとブロック積み上げ
- LlamaはReLU非線形性をSwiGLU activation functionに置き換える
- 実装した
SwiGLUはSwish-gated linear unitで、2つのlinear変換と学習可能なbetaパラメータを使う
- feed-forward部分にSwiGLUを入れたRopeModelはパラメータ数が592,706で、検証損失は1.8963程度
- その後
LlamaBlockを作り、次の構成を1つのブロックにまとめる
- RMSNorm pre-normalization
- masked RoPE multi-head attention
- residual connection
- RMSNorm pre-normalization
- SwiGLU feed-forward
- residual connection
- 最終的な
Llamaモデルはn_layers=4に設定し、OrderedDictベースのnn.Sequentialで4つのLlamaBlockを積み重ねる
- 最終モデルのパラメータ数は2,370,246で、学習結果は次のとおり
- 初回の4-layer学習後の検証損失は1.5532
- 10,000 epochsまで追加学習すると検証損失は1.1479
- さらに学習すると検証損失は0.9997
- test splitの1バッチ損失は1.2358
生成結果とデバッグ時の点検
- 最終モデルはShakespeare風の名前、改行、単語の断片を生成するが、実際の文の品質には限界がある
- cross-entropy損失はトークン選択の観点で直感化できる
- 初期損失4.17はvocabulary size 65においてランダム選択に近い
- 損失1.08は、ランダムに約2.9個のトークンから選んでいるのと同程度と解釈できる
- gradientの流れは
show_grads関数で確認する
- 各パラメータで絶対値の小さいgradientの割合を計算する
- ほとんどのパラメータgradientが0に近すぎなければ、流れは良好な状態
- 元のLlamaはCosine Annealing learning scheduleを使うが、この実装では実験結果がより悪かった
- Cosine Annealing実験では、非常に低いtoleranceでもattention biasがほとんど信号を受け取れず、理由ははっきりしないため、実装ではまず単純に始めるほうが安全
1件のコメント
Hacker Newsの意見
SwiGLUの実装にバグがあるように見える：参照論文では、feed-forward networkのbetaは学習可能な値ではなく定数で、
FFnSwiGLU = Swish1...としているhttps://arxiv.org/pdf/2002.05202.pdf の式6に基づく
公式llama実装でも定数betaは取り除かれている：https://github.com/facebookresearch/llama/blob/main/llama/mo...
ブログのログにある
"feedforward.1.beta', 0.0"の行を見ると、学習中にbetaが0へ退化しているが、本来は定数1であるべきネットワークが意図したかどうかにかかわらず変更に適応することも多く、学習後には複数のアーキテクチャ変種が似たように動作することもあるため、原本と必ず一致していなければならないのか曖昧な場合もある
こうしたミスを見つける方法としては、参照実装と出力値を正確に一致させる方式がある。HuggingFaceのtiny-randomモデルのように、ランダム重みであっても出力が正確に同じであるべきで、違えばバグの兆候になる
ただしこの方式は推論中に生じるバグにはよく効く一方、データ処理・optimizer・学習中にだけ発生する問題を捕まえるのはより難しい
個人的には自己回帰的でODEのような性質のためだと思うが、確信できるほどではない
仕事ぶりは素晴らしいが、初期の
SimpleBrokenModelとSimpleModelにはかなり多くの無駄な演算がある。embedding 65 -> 128、linear 128 -> 128、ReLU、linear 128 -> 65という順序だが、前の2層の間に非線形性がなく、どちらも線形なので、2番目の線形層は実質的に役に立っていないこのモデルは結局、古典的な単一隠れ層MLPと同じで、FLOPS基準では
128*128=16k演算を、全体の128*128+65*128=24kのうち無駄にしていることになる埋め込み層はトークンインデックスを埋め込みベクトルに変換する特殊な構造なので、取り除くことはできなさそう
全体として基本原則をよく示している。特に「
.shapeを宗教的に使え。assertとplt.imshowは友達だ」という言葉がよく、shapeの事前・事後条件は常にassertすべきbearやtypeguardがデコレータでこうした検査をサポートしているのかも気になるただし「小さく単純で速いモデルを選び、定性的に評価するヘルパーを作れ」という部分は、定量評価を意味しているのではないかと思う。そうすれば、より高度な手法と比較できる数値的なベースラインができる
論文の構成要素を一つずつ実装せよという助言も、もっと正確であるべきだ。論文は通常、複数の変更を一度に試したうえでablation studyによって各要素の寄与を示すため、コアとなるアーキテクチャ変更から始め、ablationで影響が大きい順に、依存関係を守りながら原子的な変更ごとに評価するほうがよいと思う
bearやtypeguardの代わりに、https://peps.python.org/pep-0646/ のおかげで、一部は Pythonの型アノテーションへ直接押し込める例えば
ndarray[float, Dim1, *Shape]のような形で軸ごとのshapeを型に表現し、axis値に応じて戻り値のshapeをオーバーロードできるbear/typeguardを通じて行列shapeの基本的なランタイム検査をサポートしているそれでもPythonがJuliaほど優れるのは難しそうだ。Juliaの型システムなら、行列サイズが合っていることをはるかに簡単に保証できる
SwiGLUをReLUの代わりに使う原則が何なのか気になる。著者たちが可能な非線形関数をとにかく全部試しただけなのか、もっと深い理由があるのか分からない
bearblogがDDoSを受けている最中なので、リポジトリを残しておく：https://github.com/bkitano/llama-from-scratch
AIを学ぶ立場から、記事に出てきた用語を簡単に整理してみた。トークンはテキスト片を表す整数識別子で、LLMでは限られた語彙サイズの中で、よく使われる文字片をまとめて使う
損失関数は予測と正解の差を測る値で、低いほどよい。PyTorchはテンソルとニューラルネットワークを扱うライブラリで、テンソルはスカラー・ベクトル・行列を含む多次元の数値配列である
ニューラルネットワークは重みとバイアスを持つニューロンの接続構造で、線形層はすべての入力と出力が接続された単純な構造である。ReLUは
Math.max(0, x)のような活性化関数で、線形層だけを積み重ねると結局1つの線形関数と同じになるため、非線形性を入れて学習可能性を高める勾配は学習中にモデルをより正確にするために計算される数値の変化量で、バッチ正規化は流れる数値を調整して学習を助ける方法である。位置エンコーディングはトークン同士の相対位置をベクトルで伝える
Pythonの
@演算子は__matmul__の別名で、行列積に使われる。エポックはデータセット全体を一度学習することで、バッチはパラメータ更新の前に一度に投入するデータの数であるアテンションはLLMを動かす中核で、入力トークンを並列処理して中間テンソルを作り、その後、出力トークンを生成するために使う
例えば
writing、written、writerに共通して含まれるwritが1つのトークンになり得て、writerはwritとerにトークン化され得る埋め込みは、こうしたトークンを固有の数値表現に変換する段階である
モデルの既存実装とチェックポイントがあるなら、自分の実装が正しいか確認する最も効果的な方法は、そのチェックポイントを読み込んで出力値を比較することだ
出力が合わなければ、たいていは細部の実装を間違えているので、各層を体系的に追いながら実際の差分を見つけられる。その過程で既存実装のおかしな点を発見することもある
これはモデル自体についての話であり、学習は別の軸だ。それでもハイパーパラメータをある程度似せておけば、モデル実装が正確なときは概して問題なくなる
論文の読み方も当該論文の内容もどちらもよく、KarpathyのMakemoreシリーズもおすすめする
要約された助言はとてもよく、テンソルのshapeをassertせよという助言は、どんな一般的な線形代数ライブラリにも当てはまると思う。複雑な線形代数コードを書くときは、小さなステップで進め、防御的にコーディングすることが非常に重要だ
主流言語で線形代数をプログラミングするのは、コンパイル時のshape検査がないためひどいものだ。テンソルのshapeは型の一部であるべきで、
3x4と3x4を転置なしで掛けようとしたら、コンパイル自体が通らないべきだ長い計算を回した後で、次元不一致の演算で失敗するのは本当に最悪だ
PyTorchのテンソルも、デバイスが静的に型付けされるべきだと思う。現状ではCPUメモリ上のテンソルとGPUメモリ上のテンソルを掛けようとすると、ランタイムエラーになる