今週のTOPIC

[blog] GraphRAG Toolkit を使って Amazon Bedrock で GraphRAG を構築する

• GraphRAG Toolkit は AWS が提供するオープンソースの GraphRAG 用ライブラリで、 GraphRAG を使ったシステムの構築を簡易化する機能を多く提供している
• https://github.com/awslabs/graphrag-toolkit

• Amazon Neptune Database をグラフ DB として、 Amazon OpenSearch Serverless をベクトル DB として使用可能

• このツールキットを使用することで、非構造化データからエンティティとそのリレーションを抽出し、ナレッジグラフとして保存、ユーザーの質問に対してこのナレッジグラフをクエリすることで、より関連性の高い情報を提供するアプリケーションを構築できる

• ドキュメント取り込みの手順
1. ドキュメントを細かい単位に分割する（チャンキング）
2. チャンクをLLMに渡して、命題（proposition）を抽出する
3. 命題を LLM に渡して、トピック、文言、事実、エンティティを抽出する
4. 抽出されたエンティティなどからナレッジグラフを構築し、グラフ DB に取り込む
5. 文言を LLM に渡してエンベディングを行い、ベクトル DB に取り込む

---

メインTOPIC

RegMix: Data Mixture as Regression for Language Model Pre-training

3行まとめ by Claude 3.7 Sonnet

• データミクスチャーの回帰による最適化: 言語モデル事前学習のためのデータミクスチャー選択を回帰問題として定式化し、小規模モデルの訓練結果を使って大規模モデルの最適データ混合比を予測する「RegMix」を提案。

• スケーラビリティと効率性: 1Mパラメータの小さなモデルを複数訓練して回帰モデルを構築することで、1Bや7Bのような大規模モデルの最適データミクスチャーを予測可能であり、従来手法（DoReMi）と同等以上の性能を10%の計算リソースで達成。

• 実証的発見: データミクスチャーは下流タスク性能に最大14.6%の影響を与え、高品質と認識されるWikipediaよりもCommonCrawlなどのWebコーパスが下流タスク性能と最も高い相関を示すことを発見。またドメイン間の相互作用は複雑で直感に反することが多く、自動最適化の必要性を裏付けている。

背景

• LLMの事前学習には多様なデータセットが必要

• 事前学習データの混合（data mixture）の仕方が性能に大きく影響する
• 例）Wikipedia、Github、Booksをそれぞれどれくらいの割合で含めるのが最適なのかを考えたい

• 最適なデータミクスチャーを決定する方法は確立されていない
• GPT-3では、人間の直感に基づいてWikipediaを過剰サンプリング (Brown et al., 2020)
• 手動によるデータ選択は拡張性に欠け、最適でない可能性が高い (Albalak et al., 2024)

既存手法

• トークンレベル選択
• 学習中に学習効果の低そうなトークンを動的に取り除く

• サンプルレベル選択
• 最適化アルゴリズムやperplexity、別のLLMを使って学習に使うサンプルを選ぶ

• グループレベル選択
• DoReMi (Xie, 2024): 100B tokens で訓練し、ドメインごとの損失を見て比率を決める
• 利用するモデル・データセットが大きく、計算コストが高い

提案手法：RegMix

• 小さいモデル・少ないデータで最適な比率は、大きいモデル・多くのデータでも最適

1. Train small-scale proxy models
• ディリクレ分布から混合比率を複数パターンサンプリング
• 小さいモデルでそれぞれ学習・推論して評価値を計算

2. Regression
• 混合比率を特徴量、評価値をラベルとして回帰モデル（線形モデルやTreeモデルなど）を学習

3. Simulation
• 回帰モデルを使って未見の混合分布に対する評価値の分布を推定
• 推定された分布から最適な比率を特定

4. Train a large-scale model
• 特定した最適な混合比率で大規模モデルを訓練

実験

• データセット: Pile データセットから利用可能な17ドメインを利用

• 線形回帰モデル: 一般的なL2正則化付き線形回帰

• Treeモデル: LightGBM

• 回帰データセット
• 訓練
• 1Mモデルを1B tokensで学習した時の評価値（Pile-CC Validation Loss）
• 混合比率は512パターン
• 評価（未知の混合比率で学習した時の評価値を当てられるか評価する）
• 1Mモデルを1B tokensで学習した時の評価値 x 未知の混合比率256パターン
• 60Mモデルを1B tokensで学習した時の評価値 x 未知の混合比率256パターン
• 1B モデルを25B tokensで学習した時の評価値 x 未知の混合比率64パターン

• 評価指標
• スピアマン相関係数、MSE（1Mモデル同士の比較のみ）

• 結果
• 予測性能は線形回帰 < LightGBM → 予測ではLightGBM、分析では線形回帰を使う
• プロキシモデル数は訓練トークン数よりも重要

• 1Bパラメータモデルを64種類の異なるミクスチャーで訓練

• 最高スコアと最低スコアの差は、最大14.6ptになった → 比率が性能に大きく影響する

• 定番のWikipediaデータセットはあまり相関が高くなかった

• Webコーパスの Pile-CC が最も相関が高くなった

• Webドメイン（一部抜粋）ごとの比較: 多くのドメインでPile-CCと同様の相関
ドメイン full versionの表

 
• diverse coverage across various topics and domains が効いてそう

• 比較手法
• Human: Pileデータセットを使っていた過去論文 (Gao, 2021)で決められていた比率を利用
• DoReMi: 論文内で示されていた最適な比率を利用
• PPL: perplexity-based filtering method
• ODM: 論文内で示されていた最適な比率を利用
• Pile-CC: （上で有効だと示されたので）Pile-CCだけを使ったデータセット
• RegMIX: 提案手法

• 結果
• 平均的にはRegMIXが最適な結果。次点でPile-CC。
• Estimated FLOPsは既存手法の10%程度。



• Training Tokensを変えても一貫して提案手法が良いValidation Loss (Pile-CC) になった

• ドメインごとに、Validation Lossと線形回帰の係数との相関を比較



• あるドメインが別のドメインに影響を与えることがあり、人間の直感で最適な混合を選ぶのは難しい
• 例えばPhilPapers（哲学分野の文献）が多くのドメインのValidation Lossと強い相関にある

コメント

• 提案手法はわかりやすく、効果も出ていてよい

• Wikipediaがあまり効果的でないのは意外

• 事前学習する時に試してみる価値はあるかも