Improving Lean4 Autoformalization via Cycle Consistency Fine-tuning

1. Problem Statement

Autoformalization is the task of translating natural language (NL) mathematical statements into formal proof languages like Lean4. While Large Language Models (LLMs) have shown promise in reasoning, they frequently generate Lean4 code that is syntactically valid (compiles) but semantically incorrect, failing to capture the intended mathematical meaning.

The core challenge addressed in this work is improving the semantic fidelity of NL-to-Lean4 translation without relying on expensive external tools (like a full Lean4 compiler or symbolic checker) during the training loop. The author investigates whether Cycle Consistency—the ability to translate NL $\to$ Lean4 $\to$ NL and recover the original meaning—can serve as an effective training signal to enhance model performance.

2. Methodology

2.1 Model Architecture and Setup

• Base Model: Qwen3.5-2B (a 2-billion parameter LLM).
• Fine-tuning Strategy: Low-Rank Adaptation (LoRA) is applied exclusively to the feed-forward network (FFN) sublayers ($r=16, \alpha=32$), resulting in ~9.4M trainable parameters.
• Dual-Model System:
  1. $f_\theta$ (nl2lean): The trainable model converting NL to Lean4.
  2. $g_\phi$ (lean2nl): A back-translator converting Lean4 back to NL. This model is trained via Supervised Fine-Tuning (SFT) and then frozen during the Reinforcement Learning (RL) phase.

2.2 Training Pipeline

The training occurs in two distinct phases:

Phase 1: Supervised Fine-Tuning (SFT)

• Data: FineLeanCorpus (FLC), a dataset of 509k NL/Lean4 pairs. The data is decontaminated against PutnamBench (a held-out benchmark of 660 competition problems) using TF-IDF similarity.
• Variants Tested:
  1. Curriculum SFT: Training sequentially on data sorted by difficulty (1 to 10).
  2. No-Curriculum SFT: Training on the same shuffled dataset.
• Goal: To establish a strong baseline and test if difficulty ordering aids learning.

Phase 2: Reinforcement Learning (RL) via GRPO

• Algorithm: Group Relative Policy Optimization (GRPO). This avoids the need for a separate critic model by normalizing rewards within groups of sampled completions.
• Reward Signal (Cycle Consistency):
  For a generated Lean4 statement $\hat{l}$ from prompt $s$, the reward $r$ is defined as:
  $$r(\hat{l}, s) = \cos(e(s), e(g_\phi(\hat{l})))$$
  Where $e(\cdot)$ is the sentence embedding from all-MiniLM-L6-v2. This measures the cosine similarity between the original NL statement and the back-translated NL.
• Objective: Maximize the cycle consistency reward while maintaining a KL penalty to stay close to the SFT reference policy.

3. Key Contributions

1. Cycle Consistency as a Training Signal: The paper demonstrates that using embedding-based cycle consistency as a reward in GRPO significantly improves the semantic quality of autoformalization, outperforming standard SFT.
2. Ablation of Curriculum Learning: The study rigorously tests whether training on data sorted by difficulty (1 $\to$ 10) provides benefits over shuffled training. The results show no measurable benefit for the 2B model scale, suggesting that curriculum ordering does not compensate for limited model capacity in this specific task.
3. Efficient RL Pipeline: The work presents a parameter-efficient pipeline (LoRA + GRPO) that improves performance on a single GPU without requiring the latency of a live Lean4 compiler during training.
4. Comprehensive Evaluation: The models are evaluated on both in-distribution (FLC held-out) and out-of-distribution (PutnamBench) datasets, providing a robust assessment of generalization.

4. Experimental Results

4.1 Performance Metrics

The primary metric is Mean Cycle Consistency ($\bar{C}$). Cross-entropy (CE) loss is used as a secondary metric to ensure the model's basic language modeling capabilities are not degraded.

• RL vs. SFT: The RL model achieves a 0.156 improvement in mean cycle consistency on the FLC held-out set and 0.139 on PutnamBench compared to the SFT baseline.
• Statistical Significance: The improvements are highly significant ($p < 10^{-56}$).
• Cost: The improvement comes with a negligible increase in cross-entropy loss (only +0.011 nats), indicating the model retains its language modeling ability.
• Curriculum Finding: The SFT Curriculum and SFT No-Curriculum variants are statistically indistinguishable, confirming that difficulty ordering offers no advantage at this scale.

4.2 Analysis of Results

• Difficulty Scaling: The RL improvement is most pronounced on easier problems (Difficulty 1) and decreases monotonically as difficulty increases. This is attributed to the "noisier" reward signal on hard problems, where the frozen back-translator struggles to interpret complex Lean4 output.
• Domain Performance: RL outperforms SFT across all mathematical domains. Discrete and algebraic domains (Number Theory, Algebra) show the highest gains, while continuous domains (Calculus, Analysis) show lower absolute scores due to the complexity of $\epsilon$-$\delta$ definitions and verbose Lean4 structures.
• Distribution Shift: On FLC, RL creates a "high-score tail" (many scores > 0.8), suggesting it may exploit structural regularities in the training data. On PutnamBench (more diverse), the improvement is a uniform shift, suggesting more genuine generalization.

5. Significance and Limitations

Significance:

• Semantic Alignment: The work proves that semantic consistency (via cycle translation) is a viable and powerful objective for training formalization models, potentially reducing reliance on brittle compilation feedback loops during training.
• Scalability vs. RL: While scaling the model to 9B parameters improves performance, the 2B model trained with RL (GRPO) outperforms the 9B model trained only with SFT on in-distribution data, highlighting the efficiency of RL signals.
• Practicality: The approach is computationally efficient, requiring only sentence embeddings rather than a full Lean4 compiler during the RL reward calculation.

Limitations:

• Reward Hacking: The cycle consistency metric has blind spots. The back-translator can sometimes "echo" the original statement verbatim if the generated Lean4 is structurally similar to natural language, yielding a perfect reward even if the formalization is mathematically vacuous.
• Proxy Metric: Cycle consistency is a proxy for semantic correctness. A model could produce syntactically consistent but mathematically false statements that still score high.
• Back-Translator Bias: Using a single frozen back-translator for both reward and evaluation introduces a systematic bias; errors in the back-translator affect all models equally.

Conclusion

The paper establishes that Cycle Consistency Fine-tuning via GRPO is a highly effective strategy for improving Lean4 autoformalization. It significantly outperforms standard SFT and curriculum learning, particularly on in-distribution data, while maintaining computational efficiency. The findings suggest that for small-to-medium models, optimizing for semantic round-trip fidelity is a more effective training signal than difficulty-based curriculum ordering. Future work should focus on integrating hard compiler feedback to mitigate reward hacking and further validating on diverse, out-of-distribution benchmarks.