GraFT: Grammar-Forced Translation of Natural Language to Temporal Logic

Reproducibility package for "Grammar-Forced Translation of Natural Language to Temporal Logic using LLMs" (ICML 2025).

• Paper: https://openreview.net/forum?id=p411a7WHox
• Preprint: https://arxiv.org/abs/2512.16814

This repository contains the training code, datasets, and experiment scripts needed to reproduce the paper's main result and its convergence study, plus an ablation isolating where GraFT's advantage comes from.

---

What GraFT does

Translating natural language into temporal logic (TL) is a sequence-to-sequence problem where the output has a hard syntactic contract: an ill-formed formula is not a worse answer, it is an unusable one. GraFT enforces that contract during training rather than hoping the model infers it from examples.

GrammarAwareTrainer runs a pushdown automaton (TemporalLogicGrammar) over the gold token sequence under teacher forcing. At each position it computes the set of grammatically legal next tokens and masks every other logit to -inf before the cross-entropy loss. The model therefore only ever distributes probability among valid continuations, and is never penalized for mass placed on invalid ones.

The control (control_trainer.py) is the same architecture and schedule with ordinary cross-entropy over the full vocabulary.

Important: because GraFT's training loss never sees invalid tokens, a GraFT-trained model must be decoded with the matching grammar mask (TemporalLogitsProcessor_Inference.py). Decoded unconstrained it collapses to 2.6% exact match — it emits fluent English paraphrases instead of TL. All GraFT numbers below use constrained decoding. See Results.

---

Repository layout

Path	Purpose
GraFT_trainer.py	Grammar-forced fine-tuning (the proposed method)
control_trainer.py	Standard cross-entropy fine-tuning (baseline)
GrammarAwareTrainer.py	Seq2SeqTrainer subclass that masks logits before the loss
TemporalLogicGrammar.py	Pushdown automaton over T5 token IDs defining valid TL
TemporalLogitsProcessor_Inference.py	The same grammar as a generation-time LogitsProcessor
prepare_model.py	Downloads and caches the t5-base checkpoint
experiments/	Evaluation, data-quantity sweep, ablation, plotting
augmented_*_cleaned.jsonl	The three source datasets
run_logs/	Loss histories, accuracies, and figures from our runs

Trained weights are large and not tracked in git (see .gitignore); the steps below regenerate them.

---

Data

Three datasets of paired natural-language / temporal-logic specifications, combined into one training pool:

File	Samples
augmented_CW_total_cleaned.jsonl	3,382
augmented_GLTL_total_cleaned.jsonl	11,153
augmented_navi_total_cleaned.jsonl	7,474
Total	22,009

Each JSONL record carries four fields. Training uses the lifted pair — masked_sentence → mask_ltl — in which concrete atomic propositions are replaced by prop_N placeholders:

{
  "sentence":        ["go", "to", "the", "blue", "room", "..."],
  "ltl":             ["finally", "(", "blue_room", "and", "finally", "green_room", ")"],
  "masked_sentence": ["go", "to", "the", "[prop_1]", "..."],
  "mask_ltl":        ["finally", "(", "prop_1", "and", "finally", "prop_2", ")"]
}

Inputs are prefixed with Translate the following sentence into Temporal Logic: . The pool is shuffled with seed=42 and split 90/10 into train (19,808) and held-out eval (2,201). Every script in this repo rebuilds that split identically, so results are directly comparable across runs.

---

Setup

Python 3.13 and a CUDA GPU (~8 GB is enough; we used an RTX 4070 Ti SUPER 16 GB).

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

Set USE_TF=0 for every command below. If TensorFlow with Keras 3 is present in the environment, from transformers import Trainer fails with a tf-keras error. Exporting USE_TF=0 once per shell is the simplest fix:

export USE_TF=0

Get the T5 checkpoint

Both trainers start from t5-base. Download and cache it once:

python prepare_model.py

This writes saved_models/t5-base/ (~850 MB, gitignored). Every later script loads the base model from there, so this only needs doing once.

---

Fine-tuning

Each run trains for 3 epochs (7,428 steps at batch size 8, lr 2e-5) and writes periodic checkpoints plus a final model.

# Grammar-forced model → models/GraFT_final
python GraFT_trainer.py

# Control baseline → control_models/control_final
python control_trainer.py

Approximate wall-clock on an RTX 4070 Ti SUPER: ~28 min for GraFT and ~23 min for the control. GraFT is roughly 20% slower per step because the grammar automaton advances in a Python loop over every batch position.

Checkpoints land in models/temporal_logits_base_2000/ and control_models/t5_control_2000/ every 500 steps (last 3 kept). If a run dies, resume from the newest one rather than restarting:

trainer.train(resume_from_checkpoint="./models/temporal_logits_base_2000/checkpoint-5500")

---

Running the experiments

All experiment scripts live in experiments/ and write their results to run_logs/convergence/.

1. Main comparison

Scores both final models on the held-out split — GraFT with the grammar mask, the control without.

python experiments/evaluate.py          # ~4 min

2. Convergence / data-quantity sweep

Trains a fresh GraFT and control model at each requested quantity and records both loss histories and final accuracy. Quantities are per source dataset, so 500 means 1,500 samples before the 90/10 split.

python experiments/convergence_sweep.py 500 2000          # the paper's two settings
python experiments/convergence_sweep.py 50 100 150 200 250 300 350 400 450

Small quantities are quick (~1–4 min per model); 2000 takes ~8 min for the pair.

3. Grammar-mask ablation

Decodes each control model both ways — unconstrained and grammar-constrained — to separate grammar-aware training from grammar-constrained decoding. Control models are retrained per quantity (training is seeded, so they match the sweep's runs); full reuses the saved control_models/control_final.

python experiments/ablation.py                  # every quantity, ~25 min
python experiments/ablation.py 100 500 full     # or pick a few

4. Figures

python experiments/plot_results.py

Writes run_logs/convergence/convergence_comparison.png and data_efficiency.png, skipping any quantity with no results on disk.

---

Results

Our runs, single seed, exact-match accuracy on the held-out split.

Full training data

Model	Decoding	Exact match
GraFT	grammar-constrained	2191/2201 = 99.55%
Control	unconstrained	2192/2201 = 99.59%
GraFT	unconstrained	57/2201 = 2.59%

At full data the task is saturated and the two methods are indistinguishable. The third row is the decoding caveat noted above, not a training failure.

Data efficiency

GraFT's advantage is concentrated in the low-data regime, peaking at N=100 and decaying to zero by N=2000 — consistent with the paper's reported 0.9%–42.10% improvement range.

N per dataset	GraFT	Control	Gap
50	6.7%	6.7%	0.0
100	43.3%	3.3%	+40.0
150	53.3%	33.3%	+20.0
200	55.0%	45.0%	+10.0
250	61.3%	48.0%	+13.3
300	55.6%	48.9%	+6.7
350	71.4%	60.0%	+11.4
400	84.2%	79.2%	+5.0
450	80.0%	76.3%	+3.7
500	87.3%	82.7%	+4.7
2,000	99.3%	99.0%	+0.3
full	99.55%	99.59%	−0.0

Eval sets at small N are correspondingly small (15 samples at N=50, 135 at N=450), so individual low-N points carry meaningful single-seed variance — the trend is robust, the individual values are not precise.

Where the advantage comes from

Decoding the control models with the grammar mask changes nothing at any data quantity — the accuracies are identical to unconstrained decoding, down to the hit count:

N per dataset	Control unconstrained	Control + grammar mask
50	6.7%	6.7%
100	3.3%	3.3%
150–500	33.3%–82.7%	identical at every N
2,000	99.0%	99.0%
full	99.59%	99.59%

A control-trained model's errors are already grammar-valid — it learns LTL surface syntax almost immediately and fails semantically (wrong operator, wrong proposition structure), which a syntactic mask cannot repair. So the gain is attributable to grammar-aware training, not to inference-time masking. The complete picture is asymmetric:

	Decode unconstrained	Decode constrained
Control training	99.59%	99.59% (mask is inert)
GraFT training	2.59% (collapses)	99.55% (mask is load-bearing)

Grammar-forcing does not teach the model syntax the baseline lacks; it offloads syntax to the mask entirely, freeing capacity to spend limited data on semantics. That is consistent with the benefit appearing only when data is scarce.

---

Citation

@inproceedings{
english2025grammarforced,
title={Grammar-Forced Translation of Natural Language to Temporal Logic using {LLM}s},
author={William H English and Dominic Simon and Sumit Kumar Jha and Rickard Ewetz},
booktitle={Forty-second International Conference on Machine Learning},
year={2025},
url={https://openreview.net/forum?id=p411a7WHox}
}