Expected Value Alignment for Generative Reward Modeling in Formal Mathematics Verification
Shihao Ji, Haotao Tan, Zihui Song, Mingyu Li
cs.AI
May 31, 2026 · v2
TL;DR
Trains Leibniz, a generative process reward model that scores proposed reasoning steps against Lean 4 tactic states, using an annotated dataset.
Abstract
Large Language Models (LLMs) are increasingly used with formal interactive theorem provers such as Lean 4. Scaling these systems with reinforcement learning or search methods requires process reward models (PRMs) that can evaluate intermediate reasoning steps. Existing reward-model designs expose a practical trade-off. Value-head models provide continuous scores but modify the generative model interface, while generative reward models preserve textual rationales but are poorly matched to continuous floating-point regression because numeric values are split across tokens. We introduce Expected Value Alignment (EVA), a reward-modeling procedure that keeps the surface output discrete while extracting continuous scores from the model's token distribution. The model emits integer scores in a structured JSON format, and EVA computes a continuous score as the expectation over the logits of the corresponding anchor tokens. Training combines the causal language modeling objective with an auxiliary mean squared error loss on these expected values. We instantiate EVA in Leibniz, a reward model for Lean 4 formal verification, and evaluate it against zero-shot and reward-modeling baselines. The evaluation demonstrates that continuous logit-based scoring significantly reduces discretization artifacts while retaining the interpretability of generative critiques.
Problem
Process reward models for LLM-based Lean 4 theorem proving face a trade-off. Value-head models give continuous scores but no rationale. Generative judges give critiques, but they emit scores as text tokens, which fits poorly with continuous regression.
Approach
Expected Value Alignment (EVA) has the model output a chain-of-thought critique followed by integer scores (1-5) in JSON. A continuous score is then computed as the expectation over the softmaxed logits of the anchor tokens at the score positions. Training adds an MSE loss on these expected values to the causal language-modeling loss. The method is instantiated as Leibniz-1.5B (Qwen2.5-1.5B base) on about 1,800 annotated Lean 4 tactic-state/proposal pairs scored for Logic, Alignment and Clarity.
Results
On a 400-example hold-out set, Leibniz achieves the highest Pearson correlation (0.824) and pairwise ranking accuracy (84.6%). Its MSE is close to a value-head reward model, and it keeps generative critiques.
| Model | Logic MSE ↓ | Align MSE ↓ | Pearson r ↑ | Ranking Acc. ↑ |
|---|
| Zero-Shot GPT-4o | 0.842 | 1.154 | 0.612 | 72.4% |
| Standard SFT Generative RM | 0.485 | 0.521 | 0.725 | 78.5% |
| Value-Head RM | 0.312 | 0.345 | 0.782 | 81.2% |
| Leibniz-1.5B | 0.334 | 0.368 | 0.824 | 84.6% |
Performance on the Lean 4 verification hold-out set