RHO-1: Not All Tokens Are What You Need

■ イベント
LLM論文読み会
https://sansan.connpass.com/event/353351/

■ 発表者
技術本部 研究開発部 Data Analysisグループ
大田尾 匠

■ 研究開発部 採用情報
https://media.sansan-engineering.com/randd

■ Sansan Tech Blog
https://buildersbox.corp-sansan.com/

Transcript

1. Sansan株式会社 部署 名前 RHO-1: Not All Tokens Are What You Need LLM論⽂読会 Sansan株式会社 技術本部 研究開発部 ⼤⽥尾 匠
2. ⼤⽥尾 匠 Sansan株式会社 技術本部 研究開発部 Data Analysisグループ 研究員 京都⼤学⼤学院情報学研究科修⼠課程修了。 在学中は、最適輸送を⾃然⾔語処理に応⽤した⼿法の研 究に取り組む。 2024年に新卒としてSansan株式会社に⼊社し、メール署 名取り込みにおける名刺情報抽出の研究開発に従事。 1
3. 論⽂情報 - LLMの事前学習において、学習に有⽤なトークンを選択して学習すること で、少ない学習トークン数で性能が向上する - NeurIPS 2024 の Runners ups for main track - NeurIPS版のタイトル: Not All Tokens Are What You Need for Pretraining
4. - LLMはnext-token predictionで事前学習を⾏い、汎⽤的な⾔語理解能⼒を 獲得する - 特定のタスクに特化させるには、事前学習後のLLMに対してファイン チューニングを⾏う - 今回の論⽂は事前学習にフォーカスした話 背景: LLMは⾔語理解能⼒を事前学習で獲得する 例: ⿃が空を () 正解 ⾶ぶ 予測確率 ⾶ぶ: 0.7 ⾒る: 0.2 ⾷べる: 0.01 … 正解トークンの確率を⾼くする ように学習する
5. - ⼤きなモデルを学習させるためには⼤きなデータセットが必要だが、 中には品質が低い⽂書(記事、Webページなど)も含まれる - 学習に使うデータセットにおいて、⽂書レベルでフィルタリングを⾏う ⼿法が提案されている [Brown et al., 2020, Wenzek et al., 2019] - 不要な⽂書例: HTMLタグが⼤多数を占めるWebページ - ⽂書レベルでフィルタリングを⾏ったとしても、トークンレベルで学習 に悪影響を及ぼす可能性があるのではないか？というのが研究課題 背景: LLMの事前学習における課題
6. - 実際に事前学習を⾏ったときに、トークンごとにロスがどのように 変化するかを確認した - 実験設定 - モデル: Tiny-llama-1B [Zhang et al., 2024] - データ: OpenWebMath (15Bトークン) - ⽅法: 1B学習させるごとに、各トークンのロスを検証⽤データで確認する 予備実験: 事前学習時のトークンごとのロスを調べる
7. - ロスの曲線で、トークンは⼤きく4つに分類できる - L -> L (教えなくてもわかる)や、H -> H (教えてもわからない) は、学習開始時と 終了時でロスが変わっておらず、学習のノイズとなっている可能性がある - 学習に有⽤なトークンだけを使って学習を⾏えば、より品質の⾼い事前学習済み モデルができるのではないか？ 予備実験: 学習にノイズとなるトークンが多く存在する
8. - Selective Language Modeling (SLM ≠ Small Lauguage Model) - ノイジーな学習データから、有⽤なトークンのみを選択してロスを計算する 提案⼿法: 学習に有⽤なトークンのみで学習する
9. - Reference Model - 厳選した⾼品質なデータを使って、望ましいデータ分布を学習させたモデル - データ量の例: reference model⽤: 0.9Bトークン c.f. 事前学習⽤:15Bトークン - ノイズの多い事前学習⽤データの⽬標となるデータ分布を捉える - 学習データが少ないため汎化能⼒は期待できず、あくまで「ガイド」の⽬的 - 事前学習の時は、Reference Modelのロスと差が⼤きい (i.e. 学習余地がある) トークンに絞って学習する 提案⼿法: 望ましいデータ分布を学習するモデルを⽤意
10. - トークンに対して、事前学習モデルとreference modelのロスの差を計算 する - N個のトークンからなる⽂書に対して、excess loss (学習余地)が⼤きい トークンを上位k%選別し、最終的なロスに組み込む 提案⼿法: reference modelを使ってトークンを選別する excess loss 理想的な値 学習中の値 excess loss 学習余地
11. - モデル - Tiny-llama-1B [Zhang et al., 2024]、トークン使⽤率: 60% - Mistral-7B [Jiang et al., 2023]、トークン使⽤率: 70% - データセット - 数学ドメイン - reference model学習⽤: 0.5B (⾼品質) - 事前学習⽤: 14B - ⼀般ドメイン - reference model学習⽤: 1.9B (⾼品質) - 事前学習⽤: 80B 実験: 実験設定
12. - Few-shot CoT 推論 - 提案⼿法 (RHO-1-Math)は、全トークンを使って事前学習を⾏ったモデル (Tiny-llama-CT)よりも平均16.5%の性能向上 - 全15Bトークンのうち、9Bトークン(60%)を⽤いて学習した - 学習トークン数を9B -> 30B に増やすとさらに性能が向上 - 150Bの巨⼤データセットで学習させたDeepSeekMathにも匹敵する性能 実験: 数学ドメインのタスクにおいて、⼤きな改善
13. - 15個のベンチマークにおいて、通常の事前学習を⾏う場合に⽐べて平均 6.8%の性能向上 実験: ⼀般ドメインのタスクにおいても、⼤きな改善
14. - 全てのトークンを使う従来の事前学習と⽐べて、同じ性能に達するまで の時間が5~10倍速い - データを効率よく利⽤して学習できている 実験: 同じ性能を達成するまでの速度が速い
15. - 数学ドメインだと、数学に深く関係した単語が選択されている -> 事前学習⽤の⼤きなデータセットからドメインに特化した学習が できている - ⻘: 学習されたトークン - ⿊: 学習されなかったトークン 実験: 具体的に選択されるトークン
16. - 有⽤なトークンを選んで事前学習を⾏うことで、下流タスクのロスも 下がる - 選択されないトークンでロスが上がっており、実験では確認されていないが、 汎化性能は下がる可能性がある 実験: トークン選別は下流タスクにも好影響
17. - 事前学習において、有⽤なトークンのみを選択して学習する⼿法である Selective Language Modeling (SLM) を提案した - 数学ドメイン・⼀般ドメインにおいて、通常の事前学習と⽐べて、 少ない学習トークン数で性能が⼤きく向上した - 今後の課題 - reference modelのデータ分布に収束させ過ぎると過学習を起こす可能性が あり、reference modelや事前学習両⽅の学習データセットの拡張が必要 - reference modelは学習させるのではなく、強⼒なAPIを使ってもよいかも しれない まとめ
18. - reference modelの学習データが1B程度というのが、追加学習コストが ⾼くなりすぎず良いなと思った - トークンレベルのフィルタリングがこれほどまで効果が⼤きいことに 驚いた - ドメイン特化させたいLLMであれば、ドメインによく出現するトークンを あらかじめ定義しておき、特定のトークンのロスの⽐率を⼤きくすることも 効果があるのかは気になった 感想
19. - Lin et al., “RHO-1: Not All Tokens Are What You Need”, NeurIPS, 2024. - Brown et al., “Language models are few-shot learners”, NeurIPS, 2020 - Wenzek et al., “Extracting high quality monolingual datasets from web crawl data”, LREC, 2020 - Zhang et al., “Tinyllama: An open-source small language model”, arXiv preprint, 2024 - Jiang et al, “Mistral 7B”, arXiv preprint, 2023 参考⽂献
20. Sansan 技術本部 募集ポジション紹介 https://media.sansan-engineering.com/