Spokes: Optimizing for Diverse Pretraining Data Selection
Abstractの概要
本論文は、勾配空間におけるG-Vendi多様性スコアを直接最適化することで、多様な事前学習サブセットを選択するスケーラブルな手法「Spokes」を提案しています。本手法は、固定サイズのサブセット選択をサンプルごとの確率論的な重み付け問題に緩和し、指数勾配降下法を用いて解決すると同時に、多様性とサンプルレベルの品質との間の明示的なトレードオフを可能にします。大規模なデータで実用的にアプローチを適用するため、著者らはモデルの最終層からの切り捨て勾配とJohnson-Lindenstraussランダム射影を使用し、表現コストを削減しています。DCLMとFineWebでの実験では、多様性のみのバリアントと品質重み付けバリアントの両方を、ランダムサンプリング、意味的重複排除、および品質フィルタリングと比較して評価しています。
新規性
主な新規性は、勾配空間のG-Vendiを事後的な多様性指標やクラスタリングベースの近似に頼るのではなく、事前学習データ選択の直接的な最適化目的として扱う点にあります。また、この最適化目的を層の切り捨て、低次元ランダム射影、および指数勾配最適化といったスケーラブルな近似手法と組み合わせることで、大規模なWebスケールのコーパス上で動作可能にしています。
成果
Spokesはランダムサンプリングよりもはるかに多様なサブセットを生成し、バッチレベルでG-Vendiスコアが315から520に増加し、50万サンプルのサブセットではスコアが489向上したと報告されています。下流の事前学習において、多様性のみを考慮したSpokesは、DCLM（52.5から52.9）とFineWeb（45.1から45.6）の両方でランダムサンプリングによる平均ベンチマークスコアを上回りました。品質と多様性の同時最適化が最も優れた性能を発揮し、DCLMで54.0、FineWebで46.5を達成して、同一サンプリング予算下での意味的重複排除や品質のみのフィルタリングを上回りました。
論文の注目点
- Spokesは、G-Vendiと指数勾配降下法に基づく緩和された重み付き目的関数を使用し、勾配空間におけるサブセットの多様性を直接最適化する。
- 安定したサブセット選択を維持しつつ、最終層の勾配の切り捨てやJohnson-Lindenstrauss射影などの実用的な近似手法を通じて、大規模データに対応するよう設計されている。
- 経験的に、多様性のみの選択でもランダムサンプリングに対し一貫した改善が見られ、多様性と品質を組み合わせることでDCLMおよびFineWeb全体で最高の事前学習性能が得られる。