Technical Summary: TLA-Prover

Problem Statement

TLA+ is a formal specification language used for verifying distributed systems and safety-critical protocols, paired with the TLC model checker for exhaustive state-space exploration. While Large Language Models (LLMs) show promise in code generation, they struggle significantly with TLA+ due to the language's sparse representation in pre-training corpora and its idiosyncratic syntax (e.g., temporal operators, set comprehensions).

Existing baselines, such as the FormaLLM study, indicate that untuned open-weight LLMs achieve only an 8.6% success rate in passing the TLC model checker across 25 models. A critical challenge in training LLMs for formal verification is reward hacking. In the context of TLA+ and TLC, a model can optimize for the binary "pass" signal by generating trivial, tautological invariants (e.g., TypeOK == TRUE). These specifications pass TLC on every reachable state but convey no meaningful properties about the protocol, rendering the verification useless. Standard reinforcement learning approaches that rely solely on TLC pass/fail as a reward signal are vulnerable to this failure mode.

Methodology

The authors present TLA-Prover, a 20-billion-parameter model (based on gpt-oss-20b) designed for TLA+ specification synthesis. The approach combines Supervised Fine-Tuning (SFT) with a novel training pipeline that integrates repair-based Group Relative Policy Optimization (GRPO) and a four-tier validation hierarchy.

1. Four-Tier Validation Hierarchy

To prevent reward hacking and ensure semantic validity, every model output is graded by a deterministic external pipeline into one of four tiers:

• Bronze: The output fails the SANY parser (syntactically invalid).
• Silver: The output parses but fails to load or run in TLC (e.g., timeouts, type errors).
• Gold: The output passes both SANY and TLC; the model checker exhausts the state space, and all invariants hold.
• Diamond: The output passes Gold and the invariant is mutation-sensitive.
  • Mechanism: An automated AST-level mutation is applied to the invariant predicate (e.g., negating a conjunct, flipping relational operators, or swapping constants). The mutated specification is re-run through TLC.
  • Criterion: If the mutation causes TLC to find a counterexample, the original invariant is deemed non-trivial (Diamond). If TLC still passes (e.g., because the invariant was TRUE), the output fails the Diamond check. This step explicitly blocks tautological invariants.

2. Training Pipeline

The training process consists of two primary stages, with a DPO ablation for comparison:

• Stage 1: Supervised Fine-Tuning (SFT)

  • Corpus: A curated dataset of 1,053 specifications, strictly filtered to the Diamond tier. The data is sourced from the TLA+ Foundation distribution and the tlaplus/examples repository, augmented by an Ollama generator and filtered by an LLM judge and the four-tier pipeline.
  • Architecture: A LoRA adapter (rank 8, alpha 16) is applied to the attention projections of the 20B parameter base model.
  • Goal: To teach the model the structural vocabulary and syntax of TLA+.

• Stage 2: Repair-Based Group Relative Policy Optimization (GRPO)

  • Motivation: Standard GRPO fails when all samples in a group fail the verifier (sparse rewards), resulting in zero advantage and no gradient flow.
  • Mechanism: The model is presented with a "broken" specification (harvested from previous failed attempts) and tasked with repairing it.
  • Reward Signal: A continuous improvement score based on tier transitions (Bronze→Silver: 0.1, Silver→Gold: 0.3, Gold→Diamond: 0.6). This provides dense gradients even when the final output is not perfect, as long as the tier improves.
  • Process: The model iteratively fixes its own outputs using TLC error messages as feedback.

• Ablation Study (DPO): A Direct Preference Optimization variant was trained from the same SFT checkpoint using static preference pairs (Diamond > Gold > Silver > Bronze). This served as a baseline to compare against the dynamic, repair-based GRPO approach.

Key Results

The model was evaluated on a held-out 30-problem benchmark using single-shot greedy decoding.

• Performance: TLA-Prover (final repair-GRPO checkpoint) achieved a 30% pass rate (9/30) at both the Gold and Diamond tiers.
• Baseline Comparison: This represents a 3.5x improvement over the untuned FormaLLM baseline of 8.6%.
• Ablation Comparison: The DPO variant achieved a 20% Diamond pass rate, outperformed by the repair-GRPO approach (30%).
• Anti-Reward-Hacking Validation: Crucially, the Gold and Diamond pass rates coincided at every checkpoint in the production pipeline (0 difference). This indicates that the model did not learn to generate tautological invariants that pass TLC but fail the mutation test. The only exception was a single instance in the DPO ablation, which the Diamond tier successfully caught.
• Inference Strategies:
  • Best-of-8 sampling increased the Diamond pass rate to 37% (11/30).
  • The Greedy ∪ Best-of-8 union suggests a ceiling of 43% (13/30), indicating the model possesses the capability to solve more problems given multiple attempts.
  • Self-correction (3-shot TLC feedback) yielded no improvement, suggesting the failures are due to knowledge gaps rather than random sampling errors.

Significance and Contributions

The paper claims three primary contributions:

1. TLA-Prover with a Four-Tier Hierarchy: The introduction of a validation hierarchy where the top tier (Diamond) uses mutation testing to rule out tautological invariants. This addresses the specific reward-hacking risk inherent in using TLC as a reward signal, ensuring that "passing" implies a non-trivial specification.
2. Repair-Based GRPO Pipeline: A two-stage training pipeline combining SFT on a Diamond-tier corpus with repair-based GRPO. This approach leverages online feedback from the verifier to provide dense reward signals, outperforming static DPO ablations.
3. Empirical Anti-Reward-Hacking Analysis: The observation that Gold and Diamond pass rates coincide across all checkpoints provides empirical evidence that the pipeline successfully prevents the "always-true invariant" failure mode.

The authors position this work as the first TLA+-targeted training study to substantially exceed the FormaLLM baselines. They emphasize that while the model significantly improves specification synthesis, it currently specializes in invariant-style specifications with small state spaces and struggles with complex protocols requiring case analysis or quantifier manipulation. The work underscores that verifier-guided training without anti-reward-hacking validation (like the Diamond tier) is insufficient for producing meaningful formal specifications.