BigCode の背後にある大規模ニア重複除去

Hugging Face は、MinHash と Locality Sensitive Hashing (LSH) を用いた大規模ニア重複除去パイプラインを実装し、BigCode プロジェクトのトレーニングデータ品質を向上させました。このプロセスはベンチマーク汚染を減らし、プライバシーリスクを緩和し、モデルがより小さなデータセットで同等またはそれ以上の性能を達成できるようにすることで、トレーニング効率を高めます。

LLM トレーニングにおける重複除去の重要性

Large Language Model (LLM) のトレーニングにおけるデータ重複は、モデルがトレーニングデータをそのまま出力したり、プライバシー攻撃に対して脆弱になったりするなど、いくつかの重大な問題を引き起こします。効果的な重複除去は主に次の 3 つの利点を提供します。

• トレーニング効率: モデルはより少ないトレーニングステップで同等または優れた性能を達成できます。
• 評価の完全性: 重複を除去することでデータ漏洩やベンチマーク汚染を防止し、性能向上が本当にモデルの改善によるものか、テストデータをトレーニングで見たことによるものかを保証します。
• アクセシビリティ: データセットの物理サイズを削減することで、保存、転送、共同作業が容易になります。

技術的実装: MinHash と LSH

BigCode は、スケールでニア重複文書を特定・除去するために 3 段階のワークフローを利用しています。

1. シングリングとフィンガープリント

プロセスは、テキストを n-gram（シングル）にトークナイズすることから始まります。たとえば、単語レベルの tri-gram を文書の表現に使用します。各シングルは複数回ハッシュおよび置換され、文書内のすべてのシングルに対して各置換の最小ハッシュ値を取ることで「MinHash」フィンガープリントが作成されます。この操作の時間計算量は $\mathcal{O}(NM)$（$N$ は文書数、$M$ は文書長）であり、並列化により線形スケーラビリティを実現します。

2. Locality Sensitive Hashing (LSH)

すべての文書ペアを比較するという計算コストの高いタスクを回避するため、LSH は MinHash フィンガープリント配列をバンドに分割します。特定のバンド内で同じハッシュ値を共有する文書は同じバケットにグループ化され、重複除去の候補ペアとしてフラグが付けられます。

3. 重複除去とクラスタリング

候補ペアが特定されると、BigCode はグラフベースの手法を用いて重複文書を連結成分にクラスタリングします。パイプラインの初期バージョンでは偽陽性を除去するために Jaccard 類似度を二重チェックしていましたが、"The Stack" データセットでの実験により、すべての LSH 候補を真陽性として扱う方が下流モデルの性能が最も高くなることが判明しました。

Spark を用いたパイプラインのスケーリング

テラバイト規模のデータを処理するため、Hugging Face はローカルの Python フレームワークから Apache Spark へ移行しました。これにより分散 groupBy 操作や連結成分検出アルゴリズムの実装が可能となり、GCP DataProc を使用して 1.4 TB のデータを 4 時間未満で、約 $15/時間 のコストで重複除去できました。

モデル性能への影響

ニア重複除去はコードモデルの品質に大きな影響を与えます。主な知見は以下の通りです。

• データセットサイズと性能: ニア重複除去により、より小さなデータセット（例: 3 TB 対 6 TB）でモデルがより良い性能を発揮できます。
• 積極的な重複除去: 類似度閾値を下げ、シングルサイズを増やす（例: unigram から 5-gram へ）ことで、偽陽性率を減らし、さらなる性能向上が得られます。
• リコール: 類似度閾値を下げることで高類似ペアのリコールが向上し、冗長データをより多く除去できます。

制限事項と今後の方向性

ニア重複除去は基礎的なステップですが、毒性、バイアス、PII などのデータ品質フィルタリングの必要性を置き換えるものではありません。さらに、ベンチマーク汚染は依然として課題であり、たとえば MBPP ベンチマークは GitHub 上で一般的に見られる Leetcode 問題と高い類似性を持っています。

今後の研究課題としては、コードに対するサブストリング重複除去、単一文書内の繰り返し段落検出、そして多様性と冗長性のバランスを取るためのモデル埋め込みを用いたセマンティック重複除去の探索が挙げられます。