🦿
LLMを数学タスクにアラインする手法の系譜 - GPT-3からQwen2.5まで

2024/12/27に公開
 はじめにこの記事では、これまでに提案されてきたLLMを数学タスクにアライン(=適合)する各種手法について概観する。

現時点でオープンウェイトのモデルの中では最も性能が良いQwen2.5に至るまでの提案手法を俯瞰的に把握することを目的としている。
後半では紹介した提案手法の特徴と課題に関して主観を交えながら考察した。
なお、以下で紹介する内容は私がこれまでに読んだことのある論文に基づいて説明している。網羅的なサーベイではないことに注意されたい。

 数学タスクにアラインするための学習手法の分類初見の人が頭の中で整理しやすいように、これから紹介する提案手法を以下の4通りに分類する。どれも候補選択の方法に違いがある。
ベースライン - Few shotプロンプト [Hendrycks, 2021]
ナイーブ手法 - 多数決 [Huang, 2022a] [Huang, 2022b]
ルールベース手法 - rejection sampling [Zelikman, 2022] [Yuan, 2023] [Singh, 2023]
強化学習手法 - reward model [Cobbe, 2021] [Uesato, 2022] [Singh, 2024] [Yang, 2024]
ただし、以下の説明では上記のグループ順にはまとめずに、論文の公開順にこれらの手法の概要を記述する。

 数学タスクにおけるLLMの自己改善手法の系譜この一連の研究の発端はGPT-3の時代まで遡る。[Hendrycks, 2021] は競技レベルの数学のベンチマークであるMATHデータセットを公開し、これらの問題に対しては当時世界最高の性能を誇っていたGPT-3を持ってしても、LLMのパラメータのスケーリングの効果が全く期待できないことを示した。これを受けて当時の研究者達は「数学の問題に対してモデルのスケーリングにより性能改善が見込める」学習手法の探索に着手し始めた。


(図は[Hendrycks, 2021]より引用)
この問題に最初の一石を投じたのはOpenAIで、[Cobbe, 2021]では生成モデルに複数の回答を生成させ、その良し悪しを評価するLLM(verifier)を別途学習することで、モデルサイズに従って性能をスケーリングさせられる可能性を示した。これが最初の論文かどうかはわからないが、その後、2022年から現在までにかけて、数学のように回答が1つに定まるタスクではこの手法とその派生手法が主流となっていく。


(図は[Cobbe, 2021]より引用)
[Huang, 2022a]では、greedy decodingやbeam searchを用いて解法を1つだけ生成するよりも、確率的デコーディングを用いて解法を複数生成させ、最終的な回答を多数決により決定する方が性能が良いことを示した。さらに[Huang, 2022b]では、この手法をさらに発展し、majority voteで多数派になったCoT推論過程のデータを用いてfine-tuningすることにより、さらに性能が良くなることを示した。これらの手法はともに多数決というルールベースで回答の候補を選択しているという点に特徴がある。
(図は[Huang, 2022a]より引用)
[Huang, 2022b]の手法は正解ラベルを必要としない一方で、多数決のようなナイーブな手法ではモデルが利用できる情報に限界がある。これに対し[Zelikman, 2022]では、正解へと導くCoT推論過程のデータのみをフィルタリングしてfine-tuningする手法を提案した。この手法は正解ラベルを候補選択に利用するという点で[Huang, 2022b]の手法と異なる。適用可能なタスクは正解ラベルが付与されたデータに限定されるものの、ラベルから豊富な情報を供給することができる点に特徴がある。提案手法(STaR)はCommonQAタスクにてモデルのパラメータ数が30倍ある当時のGPT-3と同等の性能を示し、この手法の手応えを世間に知らしめた。なお、STaRで用いられた「解法の候補をある基準でフィルタリングする」手法は統計学ではrejection samplingと呼ばれる手法と密接に関連している。以降ではこの手法をrejection sampling fine-tuningと呼ぶことにする。


(図は[Zeikman, 2022]より引用)
その後の手法は、基本的にはSTaRを踏襲したものが主流になっている。[Uesato, 2022]では、回答の候補選択にSTaRのようなルールベースの代わりに、Reward model（RM）を使ってデータから学習するアプローチを提案した。この論文ではReward modelのdesign choiceとして、ORM(Outcome reward model)とPRM(process reward model)の2種類の手法が提案されている。ORMは一連の計算過程の結果(outcome)に対して報酬を与えるモデルなのに対し、PRMは個々の計算過程(process)に対して逐一報酬を与えるモデルである。後者の方がより細かい教師信号をモデルに供給できる一方で、人間がより詳細なラベルを与える必要があり、ラベリングのコストが高い。この論文で評価に用いたデータセットはGSM8Kのみであるので、データセットに特有の結果である可能性があるが、STaRのようなルールベース手法よりもRMを学習する手法の方が個々の推論過程の誤りに強いモデルを学習できる可能性を示した。
(図は[Uesato, 2022]より引用)
[Yuan, 2023]では基本的にはSTaRと同様、正解へと導く推論過程のみをフィルタリングする手法を踏襲するが、より実践的なアレンジを加えている。また、オープンウェイトモデル（Llama/Llama2）をベースにファインチューニングした結果を報告している。この論文の内容で特筆すべきは、解法の候補選択に単一のモデルの生成結果のみでなく、複数のモデルの生成結果を混ぜて用いている点が挙げられル。これは解法候補の多様性を上げるために導入したもので、ドキュメント検索タスクにおけるretriever-rerankerモデルで用いられる発想に近いと言える。解法データの混合なしのケースでは33Bあたりで性能改善が止まっていたのに対し、混合データを用いたfine-tuning結果では70Bまでの性能のスケールに成功している。一方で、65B以上のモデルではRejection samplingを用いない、通常のfine tuningと比較した性能改善は認められなかった。
(図は[Yuan, 2023]より引用)
[Singh, 2023]はrejection sampling手法を統計学のEMアルゴリズムの観点から説明し、理論的な整理を行った。また、Math/AppというGMS8Kよりも難しいタスクに対してもrejection sampling手法が有効であることを示した。
報酬モデルの学習については性能改善以外に、計算量を節約する手法も提案されている。本来の強化学習では1) policy model, 2) value function model, 3) reward modelの3つが必要であるが、2024年2月に公開されたDeepSeekMath[Shao, 2024]ではvalue function modelを省き、policy modelとreward modelのみで強化学習を行う手法が提案されている。
(図は[Singh, 2024]より引用)
[Yang, 2024]はオープンウェイトのLLMの中では現状最も数学の性能が強い Qwen2.5-Math の学習手法に関する技術文書で、ここではpythonコードとして実行可能なのみを採用するTIR(Tool Integrated Reasoning)手法を用いているほか、[Yuan, 2023]と同様、モデルサイズの小さいモデルや学習途中のモデルが生成した解答候補をサンプリングすることで、候補の多様性を維持しつつ、候補に正解と不正解が満遍なく含まれるような工夫を採用している。
(図は[Yang, 2024]より引用)

 まとめと考察本投稿ではGPT-3からQwen2.5に至るまでのLLMを数学タスクにアラインする各種手法についての系譜を確認した。
これらの手法はどれも「LLMを使って複数解法の候補をサンプリングし、その中から最適な1つを選択する」という共通点がある。以下ではこの手法を便宜的にsample&select手法と呼ぶことにする。
なぜこの方法がうまくいくのかについては議論の余地があるが、以下では主観を交えて考察する。

 なぜsample&selectがうまくいくのか？最も一般的な仮説としては「LLMが問題の回答率を最大化する」ようなファインチューニングは学習データに対して過剰適合しやすい、というのがある。

[Cobbe, 2021]によると、ファインチューニングのepoch数を増やしすぎると1shotの問題の正答率(Pass@1)は徐々に上がっていくものの、可能な解法の多様性(Pass@N)は徐々に下がっていることを報告している。
(図は[Cobbe, 2021]より引用)
これに対し、sample&selectアプローチでは「生成モデルにおいて多様な解法を生成する余地を残しつつ、報酬モデルによって最適な解法を選択する」ことによって解法候補の多様性を失わずに問題の正答率を上げることができている可能性がある。

 sample&selectの課題次に、この手法の課題について述べる。この手法の大きな欠点は正答率を上げるのに大量の候補サンプリングが必要という点である。候補数をNとすると推論時間が単純にN倍になるので、実用的なモデルでは致命的な欠点である。
この欠点は本質的に性能とコストのトレードオフで回避不能である可能性もあるが、個人的には人間の場合はランダムに思いついた解法を片っ端から試すようなアプローチは通常取らずに、有望そうな解法から順に思いついて試しているはずなので、まだ改善の余地はありそうに思う。

ルールベースの方法としては[Snell, 2024]で提案されている、問題の難しさをモデルに予測させ、それに応じて推論コストをコントロールする手法が挙げられる（例えば、sampling&select手法の場合は難しい問題ほどサンプル数を増やすなど）。
いずれにせよ、ここで紹介した手法はどれも候補となる解法を虱潰しに試すという点であまり「知的」とは言えないかもしれない。特に、良質な解法がサンプリングできるが確率というガチャに頼っているところに致命的な問題がある。

目標地点がAGIならば、現状よりも「知的」かつ効率的なアプローチを探索していく必要がありそうだ。

 Reference[Hendrycks, 2021] "Measuring mathematical problem solving with the math dataset", Mar 2021, Hendrycks et. al.
[Cobbe, 2021] "Training Verifiers to Solve Math Word Problems", Oct 2021, Cobbe et. al.
[Huang, 2022a] "Self-Consistency Improves Chain of Thought Reasoning in Language Models", Mar 2022a, Huang et. al.
[Huang, 2022b] "Large Language Models Can Self-Improve", Oct 2022b, Huang et. al.
[Zelikman, 2022] "STaR: Bootstrapping Reasoning With Reasoning", Mar 2022, Zelikman et. al.
[Uesato, 2022] "Solving math word problems with process- and outcome-based feedback", Nov 2022, Uesato et. al.
[Yuan, 2023] "Scaling Relationship on Learning Mathematical Reasoning with Large Language Models", Aug 2023, Yuan et. al.
[Gulcehre, 2023] "Reinforced Self-Training (ReST) for Language Modeling", Aug 2023, Gulcehre et. al.
[Singh, 2023] "Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models", Dec 2023, Singh et. al.
[Shao, 2024] "DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models", Feb 2024, Shao et. al.
[Snell, 2024] "Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters", Aug 2024, Snell et. al.
[Yang, 2024] "Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement", Sep 2024, Yang et. al.