Legal Mathematical Reasoning with LLMs: Procedural Alignment through Two-Stage Reinforcement Learning

TL;DR

Researchers have introduced LexNum, the first Chinese benchmark for legal mathematical reasoning, and LexPam, a two-stage reinforcement learning framework. By combining curriculum learning with a novel "procedural reward" system, they enabled a lightweight 1.5B parameter model to outperform much larger models in calculating complex legal compensations while strictly adhering to statutory procedures.

What Problem Does This Paper Address?

In high-stakes legal contexts, "getting the math right" isn't just about arithmetic; it’s about procedural alignment. For example, calculating traffic accident compensation requires a model to first determine liability, identify applicable insurance caps, and then perform the calculation.

Current LLMs fall into two traps:

1. Legal-domain LLMs (like DISC-LawLLM) are often fine-tuned for legal knowledge but stumble on basic arithmetic, leading to "hallucinated" numbers.
2. Open-domain reasoning models (like DeepSeek-R1) are math wizards but often ignore mandatory legal steps, produce overly verbose "stream-of-consciousness" thoughts, or remain uncertain about which legal rule to apply.

The gap is the lack of a system that is both mathematically precise and legally compliant without requiring prohibitively expensive, human-annotated step-by-step reasoning paths.

Methods

The authors propose LexPam, a two-stage training framework built on Group Relative Policy Optimization (GRPO)—a reinforcement learning algorithm that eliminates the need for a costly "value network" (critic model), making it highly efficient for smaller models.

1. Curriculum-Based Data Selection: They use a "Teacher" model (DeepSeek-R1-Distill-Qwen-7B) to partition the training data. Samples the teacher gets right become the "foundation" set ($D_1$), while samples it fails become the "challenging" set ($D_2$).
2. Stage 1: Foundation Training: The student model (a 1.5B LLM) is trained on $D_1$ using rewards for mathematical correctness and proper output formatting (e.g., using <think> tags and boxed answers).
3. Stage 2: Procedural Grounding: The model is trained on the harder $D_2$ set. Here, the authors introduce a Procedural Reward ($r_{law}$). This reward measures the presence of task-specific legal keywords (e.g., "Liability," "Insurance," "Monthly Calculation") in the model's output. This guides the model to follow the correct legal "flow" without needing manual step-by-step labels.

Datasets

The paper introduces LexNum, a new Chinese benchmark derived from real-world legal documents (from pkulaw). It contains 2,965 training and 743 testing instances across three high-impact scenarios:

• Economic Compensation: Labor contract terminations and service length calculations.
• Work Injury Compensation: Disability benefits and insurance limits.
• Traffic Accident Compensation: Multi-layered calculations involving liability ratios and third-party insurance.

The dataset was constructed using a "GPT-4o-then-human" pipeline to ensure high quality and internal consistency.

Key Insights and Ideas

• Small Models, Big Reasoning: A 1.5B model can achieve state-of-the-art results if the training is structured correctly. LexPam (1.5B) outperformed the base DeepSeek-R1-1.5B by over 40 percentage points.
• The Power of Keyword Rewards: You don't always need dense, human-written reasoning chains. Rewarding the presence of essential legal elements (procedural grounding) effectively forces the model to "think" through the necessary legal steps.
• Harder is Better (Eventually): Ablation studies showed that training on "hard" samples ($D_2$) alone provides more transferable reasoning skills than training on "easy" samples, but the two-stage approach (easy then hard) provides the most stable and high-performing results.

Key Results

• Accuracy: LexPam (1.5B) achieved 60.65% average accuracy, crushing existing legal LLMs (which mostly scored below 15%) and outperforming GPT-4o-mini (56.93%).
• Efficiency: It approached the performance of the massive DeepSeek-R1 (671B) and QwQ-32B, despite being a fraction of their size.
• Human Evaluation: In blind tests, legal experts rated LexPam higher than DeepSeek-R1 and DISC-LawLLM in terms of Completeness and Coverage, noting that it produced legally sound, concise, and readable outputs.

Limitations and Future Work

• Model Scale: The study primarily focused on 1.5B models due to resource constraints. The authors plan to test if these gains scale linearly with larger models.
• Lack of CoT Labels: While the procedural reward works well, the dataset currently lacks gold-standard "Chain-of-Thought" (CoT) annotations (the actual step-by-step logic), which the authors hope to add in the future to further refine reasoning.
• Language: The current benchmark is specific to the Chinese legal system; adapting this to other jurisdictions would require new keyword sets and statutory logic.