Data Selection via Optimal Control for Language Models

■ イベント
Kyoto AI Meetup 分科会 #1 NLP論文読み会~査読コメントまで見てみよう~
https://sansan.connpass.com/event/395169/

■登壇概要
タイトル：Data Selection via Optimal Control for Language Models
登壇者：技術本部 研究開発部 DataAnalysisグループ 山内敏嗣

■ 技術本部 採用情報
https://media.sansan-engineering.com

Transcript

1. Kyoto AI Meetup 分科会 #1 NLP論⽂読み会~査読コメントまで⾒てみよう~ Data Selection via Optimal Control for Language Models Sansan株式会社 技術本部 研究開発部 DataAnalysisグループ ⼭内敏嗣
2. ⼭内 敏嗣 Sansan 株式会社 技術本部 研究開発部 Data Analysisグループ 研究員 写真が⼊ります 京都⼤学⼤学院⼯学研究科電気⼯学専攻修⼠課程修了。 在学中は制御理論研究に取り組む。⾃動⾞メーカーでの品質保証 業務を経た後、制御システムベンダーでベイズ統計や機械学習を 活⽤した多変量時系列データ向けのアルゴリズム研究開発に従事。 現在は契約書データ化のためのVLMの研究開発に取り組む。
3. 論⽂の概要 - 下流タスク性能向上を意識したLLMの事前学習のためのデータ選択を最適 制御問題として解く理論を構築し、豊富な実験でその有効性を⽰した。 - ICLR2025のOralでAcceptされている。
4. 背景と課題 • 背景 • • 下流タスクを考慮したデータ選択の既存⼿法 • • 事前学習向けのデータ選択⼿法は、様々な⼿法が提案されているが、重複除去、下流データと 特徴が近いデータを選ぶ、事前学習の中で動的に選ぶ⼿法などが提案されている。 既存⼿法の課題 • • • ⾔語モデル（LM）の発展に伴い、下流タスクの性能向上や学習を加速することを⽬的とした、事前 学習におけるデータ選択が重要になっている。 ヒューリスティックな⼿法が多く、下流タスクも考慮した最適性が理論的に⽰されていない。 事前学習軌道（ダイナミクス）を考慮せず特定の時点の評価のみ⾏うため、学習初期の影響を無視 することになり、最終的により良い結果に辿り着けなくなってしまう。 本研究 • “データ選択”により事前学習軌道を制御し、下流タスク性能を最⼤化する理論を構築した。 3
5. 問題設定 • LLMのパラメータ𝜃をデータ𝑥! ∈ 𝐷で𝑇ステップ事前学習させることを考える。 • 事前学習プロセスにわたる累積下流損失 最⼩化問題 min ∑&#$% 𝐽 𝜃# " s.t 𝜃#'% = 𝜃# − 𝜂∇𝐿 𝜃# , 𝛾 𝐿 𝜃, 𝛾 = ∑! 𝛾! 2 𝑙 𝑥! , 𝜃 下流損失 勾配降下 事前学習損失 学習曲線下の⾯積(AUC)を最⼩化 • 𝛾! によるscore付けがデータの選択に相当する。 本研究では、この最適化問題を解けるscorerを作ることを⽬指す。 • このようにダイナミクスが制約となる最適化問題は、最適制御問題として知られており、 その分野で理論的に⽰された解法を⽤いることができる。 4
6. 最適制御とは • • • 時間発展する状態に、コストが最⼩になるよう⼊⼒（制御）を加え続ける最適化。 ポイントは、ある1点ではなく 軌道全体 を最適化すること。 最適制御はポントリャーギンの最⼤原理（PMP）という⼿法で⼊⼒を定めることができる。 例①：⾃動⾞の経路追従 例②：エアコンの温度制御 状態：位置・姿勢 制御⼊⼒：ハンドル・アクセル コスト：経路逸脱＋急操作 状態：室温 制御⼊⼒：冷暖房出⼒ コスト：快適さ＋電気代 5
7. データ選択を最適制御として⾒る 本研究（LM事前学習） 最適制御 イメージ図 𝑢! 制御対象 𝑥!"# = 𝑓 𝑥! , 𝑢! 𝛾 事前学習 𝜃!"# = 𝜃! − 𝜂∇𝐿 𝜃! , 𝛾 ブロック線図 制御器 𝑥! コスト∑%!$# 𝐿 𝑥! , 𝑢! + 𝜑 𝑥% を 最⼩化するよう設計されている データscorer 𝜃! 下流損失AUC∑%!$# 𝐽 𝜃! を最⼩化 するするよう設計されている 6
8. PMP による必要条件 定理 2.1（データ選択の PMP 条件） 𝛾 ∗, 𝜃#∗ が以下の最適解である。 l • ⟹ 各時刻 0 ≤ 𝑡 < 𝑇 に対し随伴ベクトル 𝜆∗# ∈ 𝑅) が存在して、以下が成り⽴つ。 • ∗ (4) 𝜃!"# = 𝜃!∗ − 𝜂∇𝐿 𝜃!∗ , 𝛾 (5) 𝜆∗! = 𝜆∗!"# + ∇𝐽 𝜃!∗ − 𝜂∇' 𝐿 𝜃!∗ , 𝛾 ∗ 𝜆∗!"# , 𝜆∗% = ∇𝐽 𝜃%∗ • ∗ ∗ (6) 𝛾 ∗ = arg max ∑)$# 𝛾) ∑%,# !$+ 𝜆!"# ∇𝑙 𝑥) , 𝜃! • * % ( 定理の解釈 • • • • l 式(4)は、パラメータ𝜃の事前学習による勾配更新。 式(5)は下流タスクも考慮した理想的な勾配⽅向𝜆を得る式。 式(6)は𝜆 に勾配が近いデータに⾼スコア𝛾を割り当てる式（右図）。 実応⽤では計算都合で多くの近似を⼊れて、定理2.1の式(4)-(6) を繰り返し計算する。 7
9. データ選択のPMP 条件の実応⽤ • PMP-based Data Selection (PDS)： • • • • • ① Proxy LM(~160M)を⽤いてPMPを適⽤して、Proxy Datasetにscore付けする。 ② Data Scorer(~125M)𝛾 𝑥) = 𝑤 ∗ % ℎ8 𝑥) , 𝜙 ∗ + 𝑏 ∗ を学習する。 ③ Data Scorerで⼤規模データのscoreを予測し、Gumbel-Top-Kサンプリングを⾏う。 本命LM(160M〜1.7B)を学習する。 ①〜③はオフライン（本命LMの事前学習とは独⽴）であるため⼀度Data Scorerを得れば、 任意サイズの本命LMに再利⽤できる。 8
10. 実験設定 • データ DにはCommonCrawl（Redpajama）を⽤いて、⽬標 J(θ)にはLIMAの損失を⽤いた。 • 過学習回避のため、評価はLIMA以外で実施した。 • モデルはMistral構成 160M / 470M / 1B / 1.7B をスクラッチ学習した。 • ベースラインには以下4つを⽤いた。 • Conventional：D から⼀様にサンプル（＝Redpajama そのまま）。 • RHO-Loss：下流データで学習したモデルとの損失差が⼤きいデータを優先。 DSIR：下流データ(LIMA)との n-gram 特徴の重なりが⼤きい⽂書を、重要度リサンプリングで選ぶ。 IF-Score：影響関数を⽤い、下流ロス J(θ) への影響が⼤きいデータを選ぶ。 • • 9
11. 結果①：学習加速と下流タスク性能向上 • 下流タスク性能向上：PDSはほとんどのデータセットでベースラインを上回った。160Mか ら1.7BLMの規模が拡⼤してもパフォーマンスの向上が⼀貫している。 • 学習加速：1.7B で同等精度に達するまで 2.0× 演算回数が短縮された。 10
12. 結果②：⼤規模へ外挿しても有効 • 事前学習損失(DCLM loss)の計測：モデルサイズ(160M~1.7B)でPDSは⼀貫して⾼性能。 • ⼤規模へ外挿：以下のスケーリング則（Scaling Law; Hoffmann et al., 2022）を⽤いてテ スト損失を外挿し、⼤規模LMの事前学習においても、PDSによる改善が持続することを ⽰した。 𝐿 𝑁, 𝐷 = 𝐸 + * + + )! ," 11
13. 査読結果 • 査読スコア • • • • 評価されている点 • • • • • 10 (strong accept, should be highlighted at the conference)：1名 8 (accept, good paper)：3名 6 (marginally above the acceptance threshold)：1名 最適制御をデータ選択の設定に適⽤するアプローチに新規性がある。 scorerを学習するという発想が巧妙である。 豊富なablationがなされている。 複雑な⼿法を分かりやすく説明できている（Presentation 4点が3名）。 査読では主に3つの論点があったが、rebuttalでratingを 6に引き上げた査読者もいた。 12
14. 査読論点①：1サンプル最適問題 • 指摘内容：式(6)のデータに最適なスコア𝛾を割り当てる式について、 one-hot（つまり、 1サンプルで学習すること）が理論的に最適に⾒える。 • 著者の回答 • 確かに「単⼀の学習サンプル」は式(6)の解の1つではあるが、最⼤値が同点となるデータ集合も存 在しその場合は複数サンプルに⾮ゼロ𝛾 が割り当てられる。 • 定理2.1は必要条件であり、one-hotな解が式(6)の最適解だからといって、本問題の最適解ではない。 • 必要条件では最適とは⾔えないが、最適制御分野での実証結果ではPMP条件を解くことで良い解が 得られている。深層学習における勾配降下法 が⼤域最適を保証しないことと同様である。 [再掲]定理 2.1（データ選択の PMP 条件） l 𝛾 ∗, 𝜃#∗ が以下の最適解である。 l ⟹ 以下が成り⽴つ。 • ∗ (4) 𝜃!"# = 𝜃!∗ − 𝜂∇𝐿 𝜃!∗ , 𝛾 (5) 𝜆∗! = 𝜆∗!"# + ∇𝐽 𝜃!∗ − 𝜂∇' 𝐿 𝜃!∗ , 𝛾 ∗ 𝜆∗!"# , 𝜆∗% = ∇𝐽 𝜃%∗ • ∗ ∗ (6) 𝛾 ∗ = arg max ∑)$# 𝛾) ∑%,# !$+ 𝜆!"# ∇𝑙 𝑥) , 𝜃! • * % ( 13
15. 査読論点②：理論⇄実装ギャップ • 指摘内容：PDSには近似が多く、理論(定理2.1)との対応が不透明である。 • 著者の回答 • PMP-Solverの近似後も性能(下流損失の低減効果)の⼤部分が維持されることは確認済み。 • Scorerによる近似誤差についても、PMP-Solverに よる真のスコアとの⽐較評価を⾏ったところ、 0.55という相関を⽰した。Scorerを⽤いて上位40% のデータを閾値として選択した場合、検証データの 81.6%が正しく分類された。 14
16. 査読論点③：スケーリング則の外挿は信⽤できるか • • 指摘内容：2B以下・4サイズの⼩規模実証から 175B〜405B の性能を外挿している点につい て、3⼈の査読者から以下の指摘があった。 • 複雑なモデルを適合させるにはデータポイントが不⾜している。 • ⼤規模な検証が⽋如している。 • 適合度が報告されておらず外挿の説得⼒に⽋ける。 著者の回答 • スケーリング則の数式（以下再掲）の5つのパラメータに対し、データ数80個で⼗分である。 𝐿 𝑁, 𝐷 = 𝐸 + - / . *" ! + • 先⾏研究（GPT-4／MiniCPM／DeepSeek）でも少数点や⼩規模なモデルの精度から外挿している。 • スケーリング曲線の適合度を計算し、 R²≈0.99（モデル⽅向・データ⽅向の両⽅）を提⽰した。 15
17. まとめ • データ選択を最適制御として定式化し、PMP で最適選択の必要条件を導出し た。 • 実⽤⾯ではproxyモデルやproxyデータを使ってオフラインで選定できる⼿法 を提案し、下流タスクでの有効性やモデルサイズを⼤きくしも⼀貫されるこ とを⽰した。 • 査読コメントでは、理論的かつ定量的に回答することで採択の後押しと なっている。 16
18. Sansan 技術本部 採⽤情報 https://media.sansan-engineering.com/