flowchart TD
init["Random initialization"] --> base["Base pretraining"]
base --> sft["Proof SFT"]
sft --> dpo["DPO"]
sft --> rl["Verifier-guided RL"]
sft --> full["Full tool fine-tuning"]
sft --> lora["Tool-calling LoRA"]
classDef root fill:#ffffff,stroke:#1f2937,stroke-width:2px,color:#111827
classDef checkpoint fill:#fef3c7,stroke:#92400e,stroke-width:2px,color:#111827
classDef branch fill:#dcfce7,stroke:#166534,stroke-width:2px,color:#111827
class init root
class base,sft checkpoint
class dpo,rl,full,lora branch
ProofLM: Course Architecture and Evidence Contract
machine-learning
language-models
deep-learning
reproducibility
Define the model, data, checkpoint lineage, execution profiles, and evaluation standard that organize the course.
This course trains one decoder-only language model from random initialization and follows it through pretraining, proof specialization, preference optimization, verifier-guided reinforcement learning, tool use, and parameter-efficient adaptation. The point is not to reproduce the scale or breadth of a frontier assistant. It is to make the complete causal chain small enough to inspect: data becomes tokens, tokens update parameters, objectives create behavioral pressure, tools extend what the model can do, and frozen evaluators determine which claims survive contact with held-out cases.
The project thesis is:
Train a decoder-only language model from random initialization on natural-language and mathematical text, specialize it for formal proof synthesis, then teach it to use proof tools and compare full fine-tuning with LoRA.
“From scratch” has a precise boundary. The course uses no pretrained model weights. It owns the tokenizer, model architecture, data manifests, proof generator and verifier, training loop, checkpoint format, posttraining objectives, tool protocol, and evaluation harness. PyTorch supplies tensor operations, autograd, device kernels, and optimizers. Reimplementing automatic differentiation or matrix multiplication would enlarge the implementation without strengthening the learning objective.
PyTorch is therefore the only numerical framework in the core course. Equations provide the conceptual reference; elementary PyTorch expressions provide executable references; reusable PyTorch modules provide the implementation that later chapters actually train. A bounded MLX appendix may port the smoke path, but it does not create a second theory track or a second canonical lineage.
One base model, explicit branches
The course does not apply every posttraining method in one irreversible sequence. It freezes shared parent checkpoints and creates matched branches so that a later change can be attributed to a specific intervention.
Base pretraining produces one canonical PyTorch CUDA checkpoint. Proof SFT starts from that base and becomes the common parent for the posttraining comparisons. DPO and verifier-guided reinforcement learning branch separately. Full tool fine-tuning and tool-calling LoRA also start from the same proof-SFT parent and receive matched examples, masks, update counts, decoding settings, and evaluators.
CPU smoke checkpoints traverse the same graph with smaller budgets and identical artifact schemas. The optional MLX run produces a separate smoke-scale comparison artifact and is never presented as part of the canonical checkpoint family.
Every result records its parent checkpoint and the identities of the model configuration, tokenizer, dataset, decoding policy, and evaluator. This lineage converts “the model improved” into a testable claim about one controlled intervention.
The target model is small enough to inspect
The standard model is named ProofLM. It is a GPT-style pre-norm causal decoder with a deliberately conventional architecture:
model:
vocab_size: 8192
context_length: 512
n_layers: 10
d_model: 640
n_heads: 10
d_head: 64
d_ff: 2560
normalization: layer_norm
position_encoding: rope
activation: gelu
dropout: 0.1
tie_embeddings: true
bias: false
training:
objective: causal_language_modeling
optimizer: adamw
precision: bf16_or_fp32For vocabulary size V, model width d, and L decoder blocks, the attention-plus-MLP parameter estimate is
N \approx Vd + L\left(4d^2 + 8d^2\right) + O(d).
At V=8192, d=640, and L=10,
N \approx 8192(640) + 10(12)(640^2) \approx 54.4\text{M parameters}.
That estimate guides compute planning; it does not replace inspecting the instantiated model. Chapter 04 computes the exact count, verifies that tied embeddings share storage, proves that future tokens cannot affect earlier logits, checks finite gradients for every trainable parameter, round-trips the state through serialization, and overfits a tiny batch before any substantial run is allowed.
| Profile | Approximate model | Context | Token budget | Purpose |
|---|---|---|---|---|
| Smoke | 4 layers, width 256, about 5M parameters | 256 | 2M | Unit tests, notebook execution, and the complete local pipeline |
| Standard | 10 layers, width 640, about 54M parameters | 512 | 1B | Canonical course checkpoint and substantial posttraining lineage |
| Future scale study | 12 layers, width 768, about 97M parameters | 512–1024 | 2B | Explicitly outside the first-release budget |
The rough dense-transformer estimate 6NT gives approximately 3.26\times10^{17} training FLOPs for the standard profile. It is a comparison tool, not a wall-time promise. Sustained throughput must be measured with the exact model, sequence length, microbatch, precision, and data path on the candidate device.
The corpus must support language and proof
Pretraining only on symbolic formulas would make later tool calling a schema-memorization exercise. ProofLM first needs bounded competence in natural and mathematical language, then a verified formal domain in which posttraining rewards and tool outputs can be checked exactly.
The canonical 1B-token mixture is fixed for the first standard run:
| Source | Share | Role |
|---|---|---|
| FineWeb-Edu | 60% | Natural-language educational prose |
| OpenWebMath | 25% | Mathematical prose and symbolic text |
| Locally generated verified data | 15% | Theorems, proof states, derivations, countermodels, corrections, and controlled natural-language renderings |
Each source receives a manifest containing origin, license information, pinned revision or generator version, normalization version, content hash, document identity, split, and token count. Source-level declarations do not erase the licensing obligations of underlying documents. Optional corpora and transfer suites remain separate from the canonical mixture until their component licenses and evaluation roles are explicit.
The tokenizer is a byte-level BPE model trained with the Hugging Face tokenizers library on a pinned 20M-token sample with the same mixture proportions. Its 8,192-token vocabulary includes stable document, theorem, proof, role, tool-call, tool-result, and turn-boundary tokens. UTF-8 byte fallback preserves coverage for arbitrary input. Vocabulary, merges, normalization, special-token IDs, and tokenizer hash freeze before base pretraining begins.
Chapter 02 keeps small character, byte, word, and toy-BPE constructions as counter-examples that reveal coverage and compression tradeoffs. The reusable artifact, however, is the serialized tokenizer consumed by every later chapter. It is evaluated on ordinary-text compression, formula fragmentation, proof-line length, fallback behavior, padding utilization, and truncation at context length 512.
Valid proofs are generated from derivations
Positive examples are not created by sampling arbitrary premises and conclusions and hoping they are provable. The generator samples a natural-deduction proof skeleton, instantiates its metavariables, computes open and discharged assumptions, renders the theorem and proof states, and verifies every result with an independent checker. Controlled negative examples perturb a rule, citation, formula, discharge boundary, or conclusion. Invalid entailments carry a truth assignment checked as a genuine countermodel.
The first formal domain uses
\{\neg,\land,\lor,\to,\bot\}
with premise, assumption, reiteration, conjunction introduction and elimination, disjunction introduction and elimination, implication introduction and elimination, negation introduction and elimination, contradiction elimination, and explicit assumption discharge.
A canonical proof has a stable machine-checkable surface:
<theorem>
premises:
1. P -> Q
2. P
goal: Q
</theorem>
<proof>
1. P -> Q ; premise
2. P ; premise
3. Q ; implication_elimination 1 2
</proof>Natural-language prompts may wrap that object, but the parser, renderer, and verifier must round-trip the structured form. The evaluator judges the proof object rather than the fluency of its surrounding prose.
Random row splits are inadequate because renamed formulas or cosmetically different prompts can share the same latent proof. The canonical boundary holds out complete equivalence classes: theorem skeletons, proof-tree shapes, connective combinations, proof depth and length, proposition-symbol families, paraphrase templates, invalid-example perturbations, and tool schemas. Shifted suites add deeper proofs, longer contexts, irrelevant or reordered premises, renamed variables, alternative valid derivations, and unprovable goals requiring countermodels.
Objectives change; evaluation remains comparable
Pretraining minimizes causal next-token loss. For tokens x_1,\ldots,x_T,
\mathcal{L}_{\mathrm{CLM}}(\theta) =-\frac{1}{T-1}\sum_{t=1}^{T-1} \log p_\theta(x_{t+1}\mid x_{\le t}).
Packed documents receive explicit boundary tokens and a declared cross-boundary loss policy.
Proof supervised fine-tuning applies loss only to assistant or proof tokens. For serialized tokens z and a response mask m,
\mathcal{L}_{\mathrm{SFT}}(\theta) =-\frac{\sum_t m_t\log p_\theta(z_t\mid z_{<t})}{\sum_t m_t}, \qquad m_t=\mathbb{1}[z_t\text{ belongs to the response}].
The mask comes from serialization metadata rather than a hard-coded offset. Full proof SFT establishes the attainable ceiling before parameter-efficient adaptation is introduced.
DPO starts from the frozen proof-SFT checkpoint with an identical frozen reference. For preferred proof y_w and rejected proof y_l,
\mathcal{L}_{\mathrm{DPO}}(\theta) =-\log\sigma\left( \beta\left[ \log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)} -\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)} \right]\right).
Pairs isolate proof validity, incorrect citations, missing discharge, malformed syntax, successful repair, and valid-proof concision. Length-matched controls test whether the policy learned the intended preference rather than “shorter” or “longer.”
Verifier-guided reinforcement learning branches separately from proof SFT. A first reward decomposes into exact parse, exact goal, verifier validity, formatting, and proof length:
R(y)= w_v\mathbb{1}[\operatorname{valid}(y)] +w_g\mathbb{1}[\operatorname{goal}(y)=g] -w_s\,\operatorname{steps}(y) -w_f\mathbb{1}[\operatorname{format\_error}(y)].
Validity and goal agreement are hard constraints; step cost distinguishes already-valid proofs. A deliberately weak verifier supplies the counter-example: if training reward rises while the independent verifier degrades, the run demonstrates proxy exploitation rather than improved reasoning.
Tools make recovery observable
The initial tool environment stays inside the proof domain:
parse_formulaconverts natural or symbolic input into the canonical abstract syntax tree;check_proofreturns line-level errors and the first invalid step;find_countermodelreturns a checked truth assignment for an invalid entailment;simplify_formulareturns a checked equivalent normal form; andinspect_goalreturns assumptions, target, and admissible rule applications.
Tool evaluation separates choosing the correct tool, producing schema-valid arguments, deciding not to call, using the returned result, repairing a rejected or malformed call, respecting a step budget, and handling a held-out schema. A non-learning router provides a baseline. Generic tools are excluded from the first release because they broaden the language and evaluation problem without strengthening the proof-centered thesis.
The tool-adaptation comparison starts two branches from the same proof-SFT checkpoint. Full fine-tuning updates every parameter. LoRA freezes a linear map W\in\mathbb{R}^{d\times k} and trains
W' = W + \frac{\alpha}{r}BA, \qquad A\in\mathbb{R}^{r\times k}, \quad B\in\mathbb{R}^{d\times r}.
The rank-8 course adapter targets named attention and MLP projections, starts with exactly zero effect, saves separately, and supports merge, unmerge, and exact restoration of the base weights. The full and LoRA branches see the same ordered tool traces, response masks, tokens, optimizer updates, decoding settings, and evaluator. The comparison reports trainable parameters, optimizer-state memory, checkpoint size, wall time, proof validity, tool behavior, and retention.
Evidence travels with the checkpoint
Every published result identifies:
- model configuration and exact parameter count;
- tokenizer hash and dataset manifest;
- parent and child checkpoint hashes;
- random seed and decoding configuration;
- optimizer-update count and processed tokens or examples;
- device, software environment, wall time, memory, and cost when applicable; and
- evaluator version and the exact held-out suite.
The shared evaluation harness runs held-out likelihood by corpus source; calibration; memorization canaries and overlap audits; generation diversity; mathematical-symbol well-formedness; theorem and proof-prefix completion; verifier validity and pass@k; proof length; countermodel correctness; tool choice, argument validity, no-call behavior, result use, and repair; plus general-language and proof-retention suites.
The principal acceptance gates are causal rather than ceremonial:
| Stage | Gate before the claim is accepted |
|---|---|
| Data | Sources are hashed; generated positives verify; countermodels check; structural split keys are disjoint. |
| Tokenizer | Special IDs are stable; serialization round-trips; formula fragmentation and truncation satisfy declared thresholds. |
| Model | Future tokens cannot affect earlier logits; gradients are finite; tied weights share storage; save/load is exact; a tiny batch overfits. |
| Trainer | Interrupted and uninterrupted smoke runs match exactly under the deterministic profile; the standard path records qualified tolerances and recoverable state. |
| Base checkpoint | ProofLM beats declared n-gram and MLP baselines on held-out likelihood and produces syntactically measurable text. |
| Proof posttraining | Ground-truth validity improves on structurally held-out theorems without crossing the retention threshold. |
| DPO and RL | The independent verifier remains stable or improves when the optimized objective rises; KL, entropy, length, and retention remain visible. |
| Tools | The model beats the router baseline, handles no-call cases, and repairs at least some verifier-rejected proofs. |
| LoRA | The untrained adapter has zero effect; merge/unmerge is consistent; the comparison with full fine-tuning is matched. |
| Capstone | Another implementer can reproduce the smoke pipeline, verify the evaluators, trace every branch, and identify concrete limitations. |
These gates prevent a lower training loss, higher proxy reward, or fluent sample from standing in for the behavior the course claims to teach.
The build accumulates chapter by chapter
Each chapter adds a reusable artifact and a controlled result. Later chapters import the package rather than copy implementation code into notebook cells.
| Ch | Concept and build | Preserved artifact or evidence |
|---|---|---|
| 01 | Corpus manifests, proof generator, verifier, and structural splits | Checked smoke manifest and leakage report |
| 02 | Byte-level BPE, special tokens, packing, and masks | Frozen tokenizer and packed smoke shards |
| 03 | PyTorch tensors, autograd, cross-entropy, bigram and MLP baselines | Baseline likelihood and proof-completion thresholds |
| 04 | Causal attention, RoPE, decoder blocks, generation, and serialization | Qualified ProofLM model and CPU/CUDA canary report |
| 05 | Gradient accumulation, checkpoint state, resume, and pretraining | Exact-resume smoke checkpoint and canonical base checkpoint |
| 06 | AdamW state, schedules, clipping, precision, memory, and throughput | Systems sweep and interruption-recovery report |
| 07 | Frozen language, proof, shift, memorization, and retention evaluators | Versioned checkpoint evaluation reports |
| 08 | Response-only proof SFT | Frozen proof-SFT checkpoint and structural-transfer report |
| 09 | Sequence-level preference optimization | DPO branch with length and retention controls |
| 10 | Group-normalized verifier-guided policy updates | RL branch and weak-verifier exploit comparison |
| 11 | Typed proof tools, execution loop, traces, and structured errors | Tool-trained branch and episode evaluation |
| 12 | LoRA injection, adapter state, merge, and restoration | Matched LoRA versus full-fine-tuning report |
| 13 | Shift, shortcut, correction, and intervention suites | Failure taxonomy populated with examples and effect sizes |
| 14 | Frozen manifests, lineage validation, and report assembly | Reproducible smoke command and standard artifact scorecard |
Chapters 01–07 establish the shared base model and evaluators. Chapters 08–12 create controlled branches from frozen parents. Chapters 13–14 compare the family without silently retraining it. CUDA concepts appear where first required, beginning with model qualification in Chapter 04; an eventual CUDA reference consolidates those recurring labs. The optional MLX appendix remains off the critical path and begins only after the PyTorch smoke artifact and checkpoint bridge are stable.
The package owns computation; notebooks own the argument
Reusable code belongs in projects/proof-lm/:
projects/proof-lm/
src/proof_lm/
data/ # manifests, normalization, mixture, packing
logic/ # AST, parser, rules, generator, verifier, countermodels
tokenization/ # tokenizer training and loading
model/ # attention, blocks, decoder, LoRA
training/ # pretraining, SFT, DPO, verifier RL, checkpoints
tools/ # typed schemas, registry, execution loop
evaluation/ # language, proofs, tools, retention, backend parity
portability/ # bounded MLX model, bridge, and benchmark
experiments/ # configuration-driven entrypoints
tests/The notebooks define the learning sequence: derive the idea, expose a counter-example, import and exercise the reusable implementation, run the controlled comparison, and interpret the stored output. The package holds the cumulative implementation so that a correction in one module reaches every later chapter. Large corpora and checkpoints stay in project-owned, gitignored artifact storage; the repository keeps small fixtures, versioned configurations, manifests, hashes, tests, and compact reports.
This separation also clarifies the first dependency. Before a tokenizer can be frozen or a model can be trained, the course must know which documents and generated theorem families belong to training, validation, test, and shifted evaluation. Chapter 01 therefore begins with manifests, identities, duplicate detection, proof generation, and structural split invariants. The data boundary is the first part of the model claim, not preliminary bookkeeping around it.