技術サマリー：SPOKES – 多様な事前学習データ選択の最適化

1. 問題提起

データの多様性は、事前学習コーパスを構築する上で極めて重要な要素であり、固定されたデータ予算の下で冗長性を減らし、ダウンストリームの性能を向上させます。しかし、多様性の最適化は、データポイント間の相互作用に依存する集合レベルの特性であるため、本質的に困難です。サブセット選択の組み合わせ爆発的な性質により、特に事前学習のスケールにおいては、多様性を直接最適化することは計算量的に実行不可能です。

既存のアプローチは通常、粗粒度なクラスタリング（トピック、スキル、または教師なしクラスタによるもの）やセマンティック・デデュプリケーション（意味的な重複排除）といった、プロキシ（代理指標）や近似法に依存しています。これらの手法はクラスタの表現力を確保してはいるものの、低解像度の分割に依存しており、ロングテールや過小評価された知識に含まれる情報を捉えきれないことがよくあります。さらに、現在の手法は多様性指標（G-Vendiスコアなど）を事後的な評価ツールとして扱うことが多く、最適化の目的関数として活用できていません。本研究が取り組む中心的な課題は、スケーラビリティに欠け、あるいは微細な多様性を捉えられないヒューリスティックに頼ることなく、大規模な事前学習コーパスからいかにして多様なサブセットを確実に抽出するかという点です。

2. 手法：SPOKES

著者らは、確率的な多様化フレームワーク内でG-Vendiスコアを用いて多様性を直接最適化する、スケーラブルなフレームワークであるSPOKES（Spokes: Optimizing for Diverse Pretraining Data Selection）を導入します。

2.1 コアとなる目的

本手法は、データ選択を**品質（Quality）と多様性（Diversity）**を共同で最大化するグローバルな最適化問題として定式化します。$n$ 個の例を含むデータセットが与えられたとき、以下の式を最大化するサイズ $k$ のサブセット $S$ を選択することを目的とします：
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
ここで：

• $q_i$ は、例 $i$ の品質スコア（例：FineWeb-Edu分類器によるもの）を表します。
• $K_S$ は、選択されたサブセットに制限された類似性カーネルです。
• $\alpha \in [0, 1]$ は、品質と多様性のトレードオフを制御する調整可能なパラメータです。

2.2 勾配ベースの多様性（G-Vendi）

静的な埋め込みを使用する従来の手法とは異なり、SPOKESは勾配空間で動作します。本手法は、プロキシモデルを用いて、各データポイントに対するモデルパラメータ $\theta$ に関する勾配 $\nabla_\theta \ell(x_i)$ を計算します。G-Vendiスコアは、これらの勾配から構築された類似性行列のスペクトルのエントロピーを調べることで、多様性を測定します。このアプローチは、表現ベクトルの直交性を促進し、更新方向の冗長性を減らし、より独立的で情報量の多い最適化ステップを促進します。

2.3 指数型勾配降下法による最適化

組合せ的なサブセット選択問題を解決するために、SPOKESは緩和された最適化戦略を採用しています：

1. 連続的緩和（Continuous Relaxation）: 離散的な選択の代わりに、すべてのデータポイントに対して連続的な非負重みベクトル $w$ を割り当て、これを単体（simplex, $w \in \Delta_n$）に制約します。
2. カーネル構築: 重みに依存するカーネル $K(w)$ を定義し、そのエントリーを $\sqrt{w_i w_j}$ によってスケーリングします。
3. 最適化: 目的関数は指数型勾配降下法を用いて最適化されます。各イテレーションにおいて、重みは $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$ を通じて更新され、少数の高重み例への崩壊を防ぐために正規化されます。
4. 離散的復元: 収束後、学習された重みベクトルの上位 $k$ 個のエントリーを選択することで、最終的なサブセットを復元します。

2.4 スケーラビリティの強化

このアプローチを事前学習のスケールで実現可能にするため、著者らはいくつかの効率化戦略を実装しています：

• 切り捨てられた勾配（Truncated Gradients）: フルモデルの勾配を計算する代わりに、最後の $n$ 個のTransformerレイヤー（通常は2レイヤー）のみを使用して勾配を近似します。これにより、類似性カーネルの構造を維持しつつ、計算コストを大幅に削減できます。
• ジョンソン＝リンデンシュトラウス（JL）投影: 高次元の勾配ベクトル（例：2億パラメータ）を扱うため、ランダム投影（ラデマッハー行列）を適用します。実験結果によれば、低次元の投影（例：$k=1024$）であっても、高次元のリファレンスに対して95.5%のオーバーラップという高いサブセット安定性を維持できることが示されています。
• 並列化: データセットを独立して最適化される互いに素なサブセットに分割することで、線形なスピードアップを実現できます。

3. 主な貢献

1. SPOKESフレームワーク: ヒューリスティックなクラスタリングを超え、大規模な事前学習データセットのスケールにおいても効率的に動作する、多様性最適化のためのスケーラブルで原理的な手法の導入。
2. 直接的なG-Vendi最適化: G-Vendiスコアが、単なる事後的なメトリクスとしてではなく、多様なサブセットを抽出するための指数型勾配降下法を通じて直接最適化可能であることを実証。
3. 品質と多様性の共同最適化: 品質と多様性を共同で最適化することが、いずれか一方の目的のみを最適化する場合や標準的なベースラインを使用する場合よりも優れた結果をもたらすことを証明。
4. 効率性の向上: 切り捨てられた勾配計算とJL投影による次元削減が、膨大な計算コストをかけることなく、大規模なコーパスへの適用を可能にすることを検証。

4. 実験結果

本手法は、2つの大規模な事前学習コーパス、DCLM（厳格にフィルタリングされたDolminoサブセット）およびFineWeb（フィルタリングが比較的緩やかなもの）を用いて評価されました。モデルは、1,000億トークンを用いてLLaMA-1Bアーキテクチャからゼロから事前学習されました。

4.1 多様性メトリクス

• G-Vendiスコアの増加: 50万サンプルのサブセットにおいて、SPOKESはランダムサンプリングと比較して**+489ポイントのG-Vendiスコア増加**を達成しました。対照的に、セマンティック・デデュプリケーション（SemDeDup）による増加はわずか+7ポイントでした。
• バッチの多様性: 学習を通じて、SPOKESによって選択されたバッチは、ランダムサンプリングと比較して一貫して高い平均多様性スコアを維持しました（DCLMにおいて315から520へ増加）。

4.2 ダウンストリーム性能

評価は、OLMESスイートの10個の英語ベンチマークに対して行われました。

• 多様性のみ: SPOKES（多様性のみ）は、ランダムサンプリングに対して平均ダウンストリーム性能を**+0.4**（DCLM）および**+0.5**（FineWeb）向上させました。
• 共同最適化: 最も強力な結果は、品質と多様性を共同で最適化した場合に得られました。SPOKESは、ランダムサンプリングに対して**+1.5**（DCLM）および**+1.4**（FineWeb）の利得を達成し、セマンティック・デデュプリケーションや品質のみのフィルタリングを含むすべてのベースラインを上回りました。
• 相補的な強み: 結果は、品質と多様性が異なるが相補的な特性を捉えていることを示しています。品質フィルタリングは知識集約的なベンチマーク（例：MMLU）に優れていますが、多様性の最適化は常識的推論（例：HellaSwag）を向上させます。共同最適化はこのギャップを埋め、あらゆるタスクにおいて競争力のある性能を発揮します。

5. 意義と主張

本論文は、多様性の直接的な最適化は、大規模なスケールにおいても実行可能であり、かつ実用的な利益をもたらすと主張しています。粗い近似やヒューリスティックから脱却することで、SPOKESは以下を実証しました：

• 多様性は重要な軸である: 明示的な品質監視がなくても、多様性を最適化するだけで一貫した利得が得られることは、データの分布に関する多様な勾配表現が汎化性能を向上させることを示唆しています。
• 品質メトリクスは不十分である: 標準的な品質メトリクスは、データの有用性を完全には捉えきれません。多様性を共同で最適化することで、選択されたサブセットの平均品質スコアが品質のみのベースラインよりも低い場合であっても、より良いダウンストリーム性能につながる相補的な学習シグナルを提供できます。
• スケーラビリティ: 提案手法は、膨大な計算コストをかけることなく、大規模な事前学習コーパスから高密度で極めて多様なサブセットを抽出することに成功しており、既存のセマンティック・デデュプリケーションやランダムサンプリングに依存する選択パイプラインに代わる実行可能な選択肢を提供します。

著者らは、SPOKESを高品質かつ多様なサブセットを構築するためのツールとして位置づけており、多様性の恩恵が、異なるフィルタリング体制やノイズプロファイルを持つデータセット間で一貫して得られることを強調しています。