いのいち勉強日記

Turingで自動運転の開発をしています。京大でPhDをとりました。Kaggle Grandmasterです。

LLaMA: Open and Efficient Foundation Language Models

Metaから発表された大規模言語モデルに関する論文です。どういうデータやモデルが使われているかを簡単にメモしました。

Codes: GitHub - facebookresearch/llama: Inference code for LLaMA models
Paper: LLaMA: Open and Efficient Foundation Language Models - Meta Research
Request form: Request Form
Request formから学習済みモデルを申請することができます。

概要

  • 誰でもアクセスできるオープンデータセットだけを用いて学習
  • 13Bのモデルが175BのGPT-3に相当する精度
  • 65Bのモデルが70BのChinchillaや540BのPaLMに相当する精度

推論のコストも考えたらモデルサイズを大きくするのは実用的ではないので、小さめモデルでより多くのデータを使って学習させています。このあたりの話はDeepMindから出た論文で議論されているようです。
arxiv.org

データ

使用したデータセット

English CommonCrawl [67%]

2017~2020年の5つのCommonCrawl データをCCNetを使って処理してます。CommonCrawlは毎月クロールしてデータをためている非営利団体のデータセット。
commoncrawl.org

C4 [15%]

こちらもパブリックに使えるデータセットです。クオリティフィルタリングの方法が違うらしい。

Github [4.5%]

Apach, BSD, MITライセンスのもののみ使用。ヒューリスティックな方法でフィルタリング。

Wikipedia [4.5%]

2022年の6~8月の期間のデータを使用。20カ国の言語をカバー。

Gutenberg and Book3 [4.5%]

2つのパブリックに使えるブックコーパスを使用。90%内容がかぶってるものはカット。

ArXive [2.5%]

Latexのファイルを科学のデータとして追加。最初のセクションと書誌、コメントの.texファイル、インラインの定義などを削除。

Stack Exchange [2%]

クオリティの高いQAサイトのデータ。コンピュータ・サイエンスから化学まで網羅しています。

Tokenizer

Bytepair encoding algorithmを使用し、1.4兆のトークンを生成しています。Wikipediaとブックコーパスのデータは2エポック、それ以外は1エポックだけ学習しました。

モデルのアーキテクチャ

基本的にはTransformerです。オリジナルのTransformerからその後の研究でわかった改善点をいくつか加えています。

Pre-normalization [GPT-3]

TransformerのサブレイヤーのアウトプットではなくインプットでRMSNorm normalizationしています。RMSNormはRoot mean square normalizationです。

SwiGLU [PaLM]

活性化関数にSwiGLUを使用しています。

Rotary Embedddings [GPTNeo]

相対位置埋め込みを回転行列で示したものです。

Optimizer

AdamW(β1 = 0.9、β2 = 0.95)を使っています。200ステップWarm upして、その後はCosine schedularで最終1/10の学習率になるようにしています。Weight decayは0.1、Gradient clippingは1です。

実装の工夫

xformerの使用

facebook researchが作ったPytorchベースのTransformer用のライブラリ。カスタマイズが簡単で、メモリ効率がよく計算速度も高速です。

github.com

カスタムなバックワード

checkpointingを使うけど、逆伝播のとき計算量が多い活性化関数は再計算しないようにしています。そのためPytorchの自動微分を使わずにカスタムで実装したようです。
さらにできる限り活性化関数の計算とネットワーク越しのGPUの通信を重ねるようにしたらしいです(よくわからんけど効率良くなる?)。

その他

計算に必要なCO2

計算に必要なエネルギーをCO2換算で出していて非常に重要で興味深い点だと思いました。
合計で80GBのA100を2048枚を5ヶ月使って、2,638MWhのエネルギーを使い、CO2換算で1,015トンだそうです。
それぞれのモデル、他のモデルとの比較は本文 Table 15より引用した図を参照してください。

論文のTable 15より引用

所感

大規模な言語モデルがオープンで公開されるのはインパクトがすごいですね。データもコードもすべてオープンということで、紙面上からはわからない実用上の細かいTipsもたくさんありそうです。またほとんどのデータは1エポックしか学習されてないので、精度としてはまだまだ上がる余地があるのかもしれないです。学習済みモデルも申請したら使用可能ということで、活用面でもいろいろと出てきそうです。大規模言語モデルの話題はまだまだつきそうにないので、今後の研究にも期待です。