BlackdromeAILabs/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

---

Files

File	SHA-256	Purpose
policy_ablation_no_recovery_v3.safetensors	c3b69872e925d65225f2195a0090c62729271d637b35ae38acaa6c704f89d250	The frozen v3 weights used for every v3 number in the paper
policy_ablation_no_recovery_v3.manifest.json	see SHA256SUMS	Architecture and vocabulary manifest, required by the loader
TRAINING_MANIFEST.md	see SHA256SUMS	Training data composition, seed, epochs, wall time

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.

Property	Value
Parameters	2,272,318 across 69 tensors (counted from this checkpoint, not estimated)
Embedding dimension	256
Hidden dimension	512
Attention heads	8
Transformer layers	4
Token vocabulary	64
Max sequence length	1088
Policy classes	573 (572 registered rules plus a reserved terminal class)
Goal-conditioned	yes
Dropout	0.1 during training, 0.0 at inference
Precision	F32

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.

Property	Value
Problems validated	1,099 (106 rejected during replay)
Training rows	44,497 nominal, 42,520 distinct
On-path / off-path recovery rows	44,497 / 0
Curriculum depths	2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128 (100 requested per depth)
Core template families	12 (see TRAINING_MANIFEST.md)
Epochs run	7
Wall time	2,398.9 s
Hardware	one rented NVIDIA A10G (24 GB) via Modal, CUDA 12.4.1
Backend	Rust candle-core 0.8.4, cudarc 0.13.9
Optimizer	AdamW, weight decay 0.01
LR schedule	cosine, peak 3e-4, minimum 3e-5, warmup 5% of steps
Value loss weight	0.5
Value-target discount	0.97
Frozen data seed	0xDEE9012800000008

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:

Decision procedure	Total (/120)	identity (/60)	power_of_one (/29)	Meaningful (/89)
Uniform MCTS (no model)	31	0	0	0
Value-only MCTS	31	0	0	0
Beam search (width 8)	31	0	0	0
Policy-value MCTS	92	44	17	61
Policy-only MCTS	118	60	27	87
Greedy learned policy	120	60	29	89

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},
  note    = {BlackdromeAI Labs and Pkboost AI Labs},
  url     = {https://github.com/blackdromeai-labs/LEMMA}
}

License

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