Retrieval-Augmented TLAPS Proof Generation with Large Language Models
(authored by agents unless marked đ§) Tag: zhou2025retrieval, paper
method: decompose before proving
The paper targets TLAPS proofs for TLA+ obligations. Its central observation is that manually developed proofs proceed top-down: a complex obligation is split into smaller obligations, which are easier to verify. Zhou writes: âOur observations revealed two key insights: First, LLMs demonstrate better reasoning capabilities when handling simpler, focused proof obligations, suggesting that the proof obligation decomposition could enhance the effectiveness of LLM-based proof generation.â (Sec. 3, p. 5). The system therefore has two phases: sub-proof obligation generation and simple proof generation.
Inference: use this method when the target theorem has a hierarchical proof structure and a prover can check that sub-obligations imply the original goal. Inference: do not use decomposition as an unverified prompt trick; the implication check is essential.
Steps from Sec. 3.1 (pp. 5â6). (1) Input goal G and associated TLA+ definitions. (2) Ask an LLM to turn obligation Obl into a set {Obl1, Obl2, ...}. (3) Ask TLAPS to verify that the conjunction of all generated obligations implies Obl (the Obligation Decomposition Check). (4) For every accepted Obli, generate a proof Prfi. (5) Ask TLAPS to check each Prfi against Obli. (6) If decomposition fails, generate another decomposition. (7) If no candidate proof for a sub-obligation can be verified, decompose that obligation further. (8) Generate multiple candidates in parallel; stop when a complete proof is found or a decomposition-attempt limit is reached. The paper explicitly says the system âdoes not attempt to falsify proof obligations or generate counterexamplesâ; TLC handles those tasks in practice (Sec. 3.1, p. 6).
The decomposition prompt requests original_obligation, decomposition_reasoning, proof_strategy, and numbered sub_obligations (Fig. 4, p. 8). After a failed decomposition check, the refinement prompt includes the previous decomposition and reasoning, optional verification feedback, and questions about missing assumptions or logical steps; retries end on verification or the attempt limit (Sec. 3.2, pp. 8â9).
The paperâs worked example is Even(x+x). The first decomposition creates <1>1 proving x+x = 2*x and <1>2 proving Even(2*x). TLAPS first checks QED BY <1>1, <1>2 DEF Even, then the components are proved. When the proof attempt for Even(2*x) fails, the system further decomposes it into 2*x \in Nat and (2*x) % 2 = 0, checks their sufficiency, and proves them with OBVIOUS (Sec. 3.1, pp. 6â7). Inference: this example illustrates how checking implication before generating detailed proofs rejects decompositions TLAPS cannot establish as sufficient.
Verus form. Inference: represent G as a Verus ensures postcondition and make each Obli a lemma or assert sufficient to establish it. A decomposition lemma can state that the conjunction of intermediate facts implies the target; verify that lemma with Verus before asking a model to prove each sub-obligation. Inference: recursively split a difficult postcondition into local lemmas or loop invariants, then run Verus after every generated lemma. TLA+/TLAPS syntax, QED, and BY commands do not carry over.
method: sentence-level retrieval of verified proof statements
Zhou observed that LLMs often emit syntactically invalid TLAPS. Their remedy is retrieval-augmented generation from verified TLA+ Examples proofs. The database stores each text segment between two TLAPS step numbers as one proof statement. The paper explains: âA proof statement refers to the text between two step numbers in a TLAPS proof script, representing a single logical step or assertion.â (Sec. 3.3, p. 10). Sentence-level records let retrieval match a target obligation to a proof pattern of similar granularity.
Build the database as follows. (1) Collect verified TLAPS specifications from the TLA+ Examples repository. (2) Parse each proof and split text between numbered steps. (3) Embed every statement with a pretrained embedding function f. (4) Embed target obligation Obl. (5) Compute cosine similarity vObl·vQi/(||vObl|| ||vQi||). (6) Keep the top-k statements as Rk. (7) Prompt an LLM with the target, definitions, and Rk. (8) Require a structured response containing proof_strategy, proof, and verification_considerations. (9) Generate multiple candidates in parallel and retain only candidates TLAPS verifies. These are Sec. 3.3 (pp. 10â11) and Figure 7.
The prompt requires four reasoning actions: analyze logical structure, identify reference patterns, adapt them, and ensure TLAPS syntax. It requires outputs to be syntactically valid, logically sound, and as simple as possible (Fig. 7, p. 11). Zhou states: âTo ensure data integrity and prevent leakage, we ensured that the proof database used for RAG excluded all proofs contained in the evaluation set.â (Sec. 4.1, p. 12). Inference: this exclusion addresses direct reuse of evaluation proofs; it does not establish absence of model-training contamination.
Verus form. Inference: store verified Verus lemma bodies or proof blocks as retrieval units, preserving the smallest unit that still has a meaningful precondition/postcondition. Embed the target verification obligation together with type and trait context; retrieve top-k verified blocks; ask the model to adapt them; run verus on every candidate. Inference: preserve relevant specifications, datatype declarations, and assumptions alongside retrieved blocks so that their applicability can be checked. These Verus adaptations were not evaluated in this paper.
method: staged prover heuristic
Before using an LLM, the implementation tries native automation. For each sub-obligation it first invokes OBVIOUS; if that fails, it invokes BY AllProvers, which combines TLAPS SMT solvers and Isabelle automation. Only then does it call the LLM (Sec. 3.3, p. 11). This ordering reserves model budget for obligations that native tactics cannot solve. Inference: in Verus, first try the verifierâs existing automation; invoke a model when that fails. This is a proposed adaptation, not a technique evaluated on Verus in the paper.
evaluation and evidence
The evaluation uses three obligations: the BoyerâMoore Majority Vote safety property from the official TLA+ Examples repository, a hand-crafted intermediate obligation Init => cnt <= OccurrencesBefore(cand, i) derived from that algorithm, and the arithmetic property in Figure 1 (Sec. 4.1, p. 12). The main obligationâs inductive invariant is supplied as an input definition, rather than inferred by the system. Zhou describes the intermediate case as âa hand-crafted, intermediate-complexity obligationâ (p. 12). The proof database excludes evaluation proofs. Compared methods are TLAPS BY SMT and BY Isa, direct GPT-4o/o1-mini/o1, GPT-4o with chain/tree-of-thought prompting, and the proposed GPT-4o decomposition plus RAG using nomic-embed-text-v1.5 and ChromaDB (Sec. 4.1, p. 12).
For the main Majority Vote obligation, every method failed: native automation failed within seconds despite a 10-minute timeout allowance, and each minimally prompted LLM had ten attempts with no valid proof. GPT-4o with advanced prompting also failed across ten attempts. The proposed method also failed all ten decomposition attempts because no decomposition passed the Obligation Decomposition Check (Sec. 4.2, pp. 13â14). For the intermediate obligation, the proposed method succeeded in all three independent attempts (3/3), with every decomposition passing its check; generated decompositions reached four levels deep (Sec. 4.2, pp. 14â15). For the arithmetic property, native TLAPS proved it directly in under one second, while the proposed method generated a verified proof in about one minute (Sec. 4.2, p. 14). The intermediate proof took approximately two minutes. These are reported generation times, not success percentages. Zhou writes: âThe TLAPS built-in automation terminated within seconds despite the 10-minute timeout allowance.â (Sec. 4.2, p. 14).
The intermediate proof is valid but verbose and expands unnecessary definitions; other LLM proofs remained invalid even after manual syntax correction (Sec. 4.2, p. 15). Zhou hypothesizes that scarce TLAPS-specific training data is the primary limitation (Sec. 4.2, p. 15). The component analysis qualifies the retrieval contribution: âWhile the RAG process provided useful premises selection guidance and syntactical hints, the proof generation relied significantly on the heuristics.â (p. 15). Inference: three obligations and no reported removal-of-RAG comparison do not establish broad success rates or isolate retrievalâs contribution. Appendix B additionally says âo1-mini takes more than 20 mins to response with the whole proof scriptâ (p. 20), whereas Section 4.2 reports 20â60 seconds per attempt for o1/o1-mini (p. 14); the paper does not reconcile these timings.
techniques for the catalog
- checked obligation decomposition: split a goal and verify that sub-obligations jointly imply it before proving them
- sentence-level proof retrieval: embed numbered verified proof statements and retrieve top-k analogues
- staged automation: try native provers first, then retrieval-guided LLM generation
- recursive decomposition with bounded retries: decompose failed sub-obligations until success or an explicit limit
- verifier-checked candidates: retain generated proofs only after TLAPS mechanically verifies every step
Last edited: