lemma-v3-policy

LEMMA v3 Goal-Conditioned Rewrite Policy

What this is. A small (2.27M parameter) transformer that scores rewrite rules over LEMMA's fixed 572-rule vocabulary, given an expression and the goal expression it is trying to reach. It is used to rank candidate actions that a symbolic guardrail and an evidence-typed verifier have already admitted.

What this is not. This is a verifier-constrained symbolic-rewrite policy over LEMMA's own rule vocabulary. It is not a general-purpose mathematical reasoner, not a theorem prover, not a language model, and not a state-of-the-art system. It cannot read a word problem, does not accept natural language, does not emit mathematics on its own, and produces no output that is meaningful outside LEMMA's search loop. It only ranks rules; every rewrite it influences must still be produced by the rule engine and independently accepted by the verifier. Its measured competence is confined to one controlled, internally constructed benchmark.

• Code and harness: https://github.com/blackdromeai-labs/LEMMA
• Benchmark: BlackdromeAILabs/lemma-long-horizon-rewrite-benchmark
• Paper: LEMMA: Learned Guidance for Evidence-Carrying Long-Horizon Symbolic Rewriting (PDF)

---

Files

The loader resolves the manifest by replacing the weights file extension with .manifest.json, so the two files must keep these names and sit in the same directory.

sha256sum -c SHA256SUMS

Architecture

Encoder-only transformer with a policy head and a value head.

Input encoding. [START] state <SEP> goal [END], tokenized from the expression AST. The separator reuses a previously unassigned vocabulary slot, so no tensor shape changed when goal conditioning was introduced and state-only checkpoints remain loadable.

Goal conditioning is enforced, not conventional. The manifest carries a goal_conditioned boolean and the search layer consults it before every query: a goal-conditioned model is never fed a state-only encoding, and a state-only model is never asked to condition on a goal it was not trained to use.

Vocabulary binding. The manifest records a vocabulary digest (18352749361355727982) and class count. Loading fails closed if the host rule registry does not match, so a checkpoint can never be silently applied to a different action space.

Training

Trained on synthetic data generated by LEMMA's own rule templates. No human-authored, scraped, or third-party mathematical corpus was used at any point.

v3 specifically is v2's exact architecture, training data and hyperparameters with --recovery-per-problem 0, isolating the off-path recovery-data variable. Recovery share is 13.4% for v1, 3.5% for v2, and 0% for v3.

Checkpoint selection, stated honestly. Every epoch overwrites the same output path, so the released file is the last epoch that ran (epoch 6 of 7), not the epoch with the best validation metric. The training manifest records a best branching-validation top-1 of 1.0000 at epoch 0; that is a descriptive figure for a different epoch and is not a property of these weights.

Leakage control. Generated inputs were checked against the input field of the standard, frontier and deep evaluation corpora, and exact collisions removed before examples were produced. Only the input field was read, never a label, reference path, or expected answer. Note that the curriculum includes depth 128, so the depth-128 benchmark is not a horizon-extrapolation test for this model; its instances remain exact-input held out.

Benchmark results

On the companion benchmark, 89-problem meaningful subset (the 31 forced-corridor problems are excluded), 400 generated-successor budget, this exact checkpoint:

Simulation-budget sweep, policy-value MCTS, 89-problem subset: 0 at 50 simulations, 22 at 100, 53 at 200, 61 at 400, 65 at 800. Guidance alone is not sufficient at every budget: the same policy that solves 61 problems at 400 solves none at 50.

The value head is weakly calibrated. Measured on replayed reference paths, its estimate has Pearson correlation 0.370 with its own training target and -0.398 with remaining reference steps, and its on-path and off-path distributions are nearly indistinguishable in aggregate (means 0.0559 and 0.0532). Consistent with this, adding the value estimate reduces the score from 87/89 (policy-only) to 61/89. Prefer policy-only or greedy decoding with this checkpoint.

Usage

This model has no Python inference path. It is loaded by LEMMA's Rust harness.

git clone https://github.com/blackdromeai-labs/LEMMA
cd LEMMA

# place both files together, keeping their names
mkdir -p experiments/models
cp /path/to/policy_ablation_no_recovery_v3.safetensors   experiments/models/
cp /path/to/policy_ablation_no_recovery_v3.manifest.json experiments/models/

# fetch the benchmark corpus
mkdir -p experiments/deep
cp /path/to/corpus.jsonl experiments/deep/corpus.jsonl

# validate the corpus, then evaluate this checkpoint at a 400-simulation budget
cargo run --release -p mm-solver --example deep_eval -- validate experiments/deep/corpus.jsonl
cargo run --release -p mm-solver --example deep_eval -- run experiments/deep/corpus.jsonl 400 goal_conditioned_no_recovery_v3

Other decision procedures over the same checkpoint use the arm names v3_policy_only, v3_value_only, v3_policy_value, v3_greedy and v3_beam8; uniform runs the same search with no model at all.

Loading it directly in Rust:

use candle_core::Device;
use mm_brain::PolicyNetwork;
use mm_rules::ActionVocabulary;
use std::path::Path;

let policy = PolicyNetwork::load(
    Path::new("experiments/models/policy_ablation_no_recovery_v3.safetensors"),
    ActionVocabulary::standard(),
    Device::Cpu,
)?;
assert!(policy.is_goal_conditioned());

// Goal-conditioned rule priors, aligned with the action vocabulary.
let priors = policy.rule_priors_toward(&state, &goal)?;

Priors must be masked to the guardrail-admitted, verifier-accepted successors before use; the search layer does this. Scores for rules that are not legal at the current state are meaningless.

Intended use

• Ranking verifier-admitted rewrite actions inside LEMMA's search loop.
• Research on goal conditioning, value-head calibration, and search-versus-policy trade-offs in symbolic rewriting.
• A reproducibility artifact for the paper's tables.

Out of scope

• Any use outside LEMMA's rule engine and verifier. The output is an index into a specific 572-rule vocabulary and is meaningless without it.
• Natural-language mathematics, word problems, theorem proving, or proof-certificate generation.
• Any use where a returned derivation is trusted without the verifier. The policy has no correctness guarantee of its own; all soundness in this system comes from the rule engine and the verifier, never from the model.
• Production or safety-critical deployment.

Limitations

• Narrow competence. Eight rules appear in the benchmark's reference paths out of 572 registered. The model is well exercised on those and largely unexercised elsewhere.
• In-distribution evaluation. The benchmark's wrapper motifs and depths are present in the training curriculum; results measure held-out instances, not structural transfer. The deep-OOD split name is legacy and is not an OOD claim.
• Single training run. v1, v2 and v3 differ in more than one respect and are single runs, so the lineage is descriptive, not causal.
• Goal conditioning is not shown to help. A separate seed-matched three-seed ablation, with byte-identical data and hyperparameters between arms, found paired differences of 0, -10 and -1 on the 89-problem subset: in no seed did removing the goal reduce solve rate. Goal conditioning is an architectural property of this checkpoint, not a demonstrated improvement.
• The value head does not help at this scale and training budget (see above).
• The released weights are the last epoch, not the best epoch.
• Evaluation is single-run and deterministic given the seed; no cross-seed error bars are claimed for this checkpoint.

Citation

@article{saxena2026lemma,
  title   = {LEMMA: Learned Guidance for Evidence-Carrying Long-Horizon Symbolic Rewriting},
  author  = {Saxena, Atul and Kharat, Pushp},
  year    = {2026},
  version = {1.0.0},
  publisher = {Zenodo},
  doi     = {10.5281/zenodo.22842820},
  url     = {https://doi.org/10.5281/zenodo.22842820}
}

License

Mozilla Public License 2.0, matching the training code and the generated training data.