👏
AIMOコンペ 上位解法まとめ
概要
先日KaggleのAI Mathematical Olympiad - Progress Prize 1というコンペにチームで参加し、1161チーム中53位でした。Early Sharing Prizeという賞の影響で序盤に強力なノートブックが公開されたことで、スコアがなかなか改善せず苦労しましたが、LLMを扱う上で学びの多いコンペになりました。特にvLLMの知見は今後のLLM系コンペでも役に立ちそうです。
以下にコンペ概要とEarly Sharing Prize解法、上位解法を整理します。
コンペについて
- 数学の問題（高校数学中級レベル）をLLMに回答させて、その正解数を競う。
※train.csvの最初の2問
| problem | answer |
|---|---|
| Let （和訳） |
52 |
| Each of the three-digits numbers （和訳） |
250 |
- 問題数は50問（PublicとPrivateそれぞれで50問用意されている）。正解数が同じ場合は、提出タイミングが早い方が上位となる。
- Notebookを提出する。実行時間の制限は9時間。
- モデルは2024年2月23日より前に公開されたオープンソースの物のみ使用可能。
- 提供される学習データ（train.csv）は数学の問題10問のみ。外部データ使用可能。
- 最終順位に対する賞のみではなく以下の賞も設定されている。
-
Overall Progress Prize Winner
Public、Privateいずれの正答数も47/50を超えた1位のチームには最低$794,624の賞金が与えられる。該当チームが無ければ次のコンテストに繰り越される。※今回は該当チーム無し
-
Early Sharing Prize
2024 年 4 月 22 日までに、Public LBで20/50以上 のスコアを獲得したノートブックを公開したチームに$10,000の賞金が与えられる。
-
Early Sharing Prize解法
- モデルはDeepSeekMathを使用。
※DeepSeekMathとは:
DeepSeek-Coder-v1.5 7B をベースに、Common Crawl から収集された 1,200 億個の数学トークンなどでファインチューニングされたもの。数学の問題に関するベンチマークでも上位に位置している。https://paperswithcode.com/sota/math-word-problem-solving-on-math
- モデルに直接回答させるのではなく、その問題を解くPythonのプログラム（Sympyという代数計算ライブラリを用いたもの）を出力させ、それをsubprocessで実行して回答を得るという方法が採用されていた。
※プロンプト
code = """Below is a math problem you are to solve (positive numerical answer):
\"{}\"
To accomplish this, first determine a sympy-based approach for solving the problem by listing each step to take and what functions need to be called in each step. Be clear so even an idiot can follow your instructions, and remember, your final answer should be positive integer, not an algebraic expression!
Write the entire script covering all the steps (use comments and document it well) and print the result. After solving the problem, output the final numerical answer within\\boxed{}.
Approach:"""
- CoT（step by stepで考える）とSelf-Consistency（単一モデルに回答の多様性を確保しながら複数の回答をさせ、多数決をとる）を採用していた。各問題に対する回答数は最大19回。
- ファインチューニングやRAGは使われていない。
上位解法（1位～4位）
概要
- 共通点は以下の通り
- モデルはDeepSeekMathがベース。
- Self-Consistencyに類する手法を採用
- 処理の高速化（1位～3位はvLLMを使用）
- 1位と2位は外部データを用いてファインチューニングを行っていた。
1位解法（スコア：29/50）
-
MuMath-Code論文をベースに、2段階のファインチューニングをフルパラメータで行った。
- 1回目は自然言語による数学の問題と解答の大規模で多様なデータセット。オンラインのPDFや数学のディスカッションフォーラムなどから収集した。各解答は Chain of Thought (CoT) を使用してテンプレート化されている。
- 2回目は、ステージ 1 のモデルを、問題＋一連の根拠＋Python プログラム＋その出力で構成されたデータセットでファインチューニングした。MicrosoftのToRA論文を参考に、約60000問の問題に対しGPT-4でデータセットを生成、使用した。
- 各問題に対し48回の回答を行い、最多数の回答を最終的な答えとした。その際に、コードエラーが出たらモデルにフィードバックしコードを修正させ、再度回答させた（max3回）。出力にはvLLMを用いている。
※1位はHuggingFaceとMistralのエンジニアチーム。H100×8でファインチューニングしたとのことで、スコアのみならず計算資源も圧倒的だった。
2位解法（スコア：22/50）
-
この論文を参照に、Policy ModelとReward Modelの2つをファインチューニングで作成。
- Policyモデルは問題をプログラム的に回答する。データセットにはKaggleで共有されたAMC、AIMEなどの問題を用い、GPT4でソリューションを生成させ、使用した。
- Rewardモデルは、問題とPolicyモデルが出力したPythonのコードを受け取り、評価・重みづけする。これもデータセットにKaggleで共有されたMATHやAMCなどの問題を用い、各問題に対する正答したパターンと誤答したパターンが1:1になるようにデータセットを構成した。（※Rewardモデル用のデータセットをどのように作成したかは説明の内容が良く分かりませんでした。）
- 各問題に対し、vLLMを用いてPolicyモデルに42回回答させ、それぞれの回答でRewardモデルを用い評価（0~1のスコア）を行う。各回答の評価値の幾何平均と出現回数の積を取り、それが最も高いものを最終的な回答として採用する。
3位解法（スコア：21/50）
- Early Sharing Prizeのコードをベースに、vLLMで処理の高速化を行い、各問題に対する回答を120～160個生成。kvキャッシュはfp16を採用した。
- LLMには問題の最終的な回答を「\boxed{}」というフォーマットで出力するよう指示しているが、たびたび無視される。その場合、 それまでの出力に「"The final answer is \boxed{”」という文字列を追加して再度LLMに投げ、所定のフォーマットで出力するよう強制した。
- 一つの問題の各出力で得られたPythonコードをバッチで並列で処理し、処理時間を短縮した。
- LLMの誤答のパターンとして、以下の2つが確認できた。複数の回答候補から最終的な回答を得る際に、以下に該当するものの重みを小さくした。
- 数値が10未満の場合 （コード内のエラーが原因で発生することがよくある）。
- 問題文の一部に含まれている数字。
4位解法（スコア：21/50）
- Early Sharing Prizeのコードをベースに、モデルを 2 つの T4 GPU に並列で展開して実行し、高速化。各問題の回答を30~40個生成。
- 各問題に費やされた時間を追跡し、後続の問題の試行回数を動的に調整して、指定された時間内にできるだけ多くの問題が解決されるようにした
- モデルが解決できないと思われる問題に対する試行回数を減らし、残りの問題に対してより多くの試行回数を設定した。例として、モデルは図形の問題を正答することがほとんどないため、問題文の中に図形の問題分によく現れる「[asy]」の文字列を確認したら試行回数を3回までにした。
- 回答候補生成のロジックを調整（モデルは間違った回答をすると0～5 の小さな整数を出力する傾向があるため、これらの回答の重みを減らすなど）
補足
- 我々の解法（53位）
- vLLMの解説記事（チームメートの方が執筆）
Discussion