Safe: Enhancing Mathematical Reasoning in Large Language Models via Retrospective Step-aware Formal Verification

生成日:

Safe: Enhancing Mathematical Reasoning in Large Language Models via Retrospective Step-aware Formal Verification

要約

論文の面白いところ

この論文の要点は、LLM の推論を「もっとよく採点する」だけでなく、採点の根拠を機械的に確認できる形に寄せた点にある。通常の PRM は、ある途中式が正しそうかを自然言語の文脈から点数化するが、その判断は外から検査しにくい。Safe は、途中式の正しさを Lean 4 の定理として書き直し、自動定理証明器に証明を試みさせる。証明が通れば、その推論ステップには少なくとも形式的な裏付けがある。証明が通らない場合も、誤りとは限らず、自動形式化や証明探索の失敗である可能性が残る。そこで著者らは、単純な正誤ではなく、「検証不要」「形式化失敗」「形式化成功かつ証明成功」「形式化成功だが証明失敗」という 4 状態を使う。さらに、その状態列を小さな Long Short-Term Memory(LSTM)で集約する。自然言語の見込みを測る PRM と、過去のステップを形式的に見る検証器を併用する設計は、地味だが筋がよい。

問題設定

Chain-of-Thought(CoT)による段階的推論は、数学問題で LLM の性能を上げる手段として広く使われている。しかし、途中の一手が誤っていても、最終答だけを見る評価ではその原因を見落としやすい。多数の候補解を生成し、報酬モデルで最もよい候補を選ぶ Best-of-N も有効だが、報酬モデル自体は必ずしも証明可能な根拠を返さない。数学では、主張の正しさを示す標準的な方法は証明である。したがって、自然言語で書かれた各推論ステップを、形式言語の主張として表し、機械で証明できるかを調べることには意味がある。ただし、数学の解答全体を一度に形式化して証明するのは難しく、計算量も大きい。著者らはこの負担を避けるため、解答を小さなステップへ分け、それぞれを個別に検査する問題として扱う。

提案手法

Safe は、まず LLM にゼロショット CoT で複数の解答候補を生成させる。各候補は、規則または LLM によって個別の推論ステップへ分割される。次に、検査対象のステップについて、前のステップや問題文を前提にした Lean 4 の定理を LLM が生成する。単なる言い換えやまとめのように検証が不要なステップは、そのまま検証不要と判定される。Lean 4 で表しにくい内容や、形式化に失敗した内容は、形式化失敗として扱われる。形式化に成功した場合は、DeepSeek-Prover-V1.5 や COPRA などの自動定理証明器が証明を試みる。得られた状態列は、語彙数 4 の小さな LSTM に入力され、推論候補全体の回顧的な点数になる。最後に、この点数を既存 PRM の予測的な点数と組み合わせ、候補の中から最終解を選ぶ。

結果

実験は MATH-500、GSM8K、CollegeMath の 3 データセットで行われた。推論を生成するモデルとして、Llama-3-8B-Instruct、Llama-3.1-8B-Instruct、GPT-4o、DeepSeek-Math-7B-Instruct が用いられた。比較対象には、ゼロショット CoT、自己整合性による多数決、Skywork や ArmoRM などの結果報酬モデル、Math-Shepherd や RLHFlow などの PRM が含まれる。Safe は、多くの組み合わせでこれらのベースラインを上回った。たとえば MATH-500 では、Llama 3.1 で 60.0、Llama 3.0 で 36.3、GPT-4o で 80.4、DeepSeek-Math で 52.4 の BoN@5 精度を得ている。GSM8K では改善幅が小さく、著者らは問題が比較的易しく、誤った候補が少ないため学習信号が乏しいと説明している。新たに作られた FormalStep は 30,809 件の Lean 4 形式主張からなり、数論、代数、組合せ、幾何などを含む。推論時の計算費用は PRM より大きく、単一ステップの検証には平均で複数回の自動形式化と定理証明の試行が必要である。

具体例

たとえば、「整数 n が 7 で割って 2 余るとき、(n + 2)(n + 4)(n + 6) を 7 で割った余りを求めよ」という問題を考える。LLM は、n ≡ 2 (mod 7) だから n + 2 ≡ 4、n + 4 ≡ 6、n + 6 ≡ 1 となり、積は 4 × 6 × 1 に等しい、と途中で述べるかもしれない。Safe はこの一文をそのまま信用せず、「n が 2 mod 7 なら、(n + 2)(n + 4)(n + 6) は 4 × 6 × 1 と合同である」という Lean 4 の定理に変換する。自動定理証明器は、合同式の性質を使って各項の余りを導き、積の余りが保たれることを証明しようとする。証明が通れば、この途中ステップは形式的な証拠をもつものとして扱われる。その後、4 × 6 × 1 = 24 で、24 は 7 で割ると 3 余る、という別の数値計算も検査対象になりうる。間違えやすい点は、途中で n + 6 の余りを 8 あるいは 0 のように扱ってしまうこと、または合同な値への置換を積全体に適用できる理由を省いてしまうことである。Safe はこのような小さな飛躍を、自然言語の印象ではなく、証明可能な主張として確認する。