How sreeyukth/symbolic-world-model works
1. What this repo is
- This is a Python research codebase for a symbolic world model: it treats a mathematical expression as the current world state, and a math operation as the transition action.
- The core architectural question is whether a recurrent latent state can retain enough information to model procedural reasoning over symbolic math.
2. Top-level pipeline
- The system flows in four stages: generate trajectories → tokenize expressions → train recurrent world model → probe latent structure.
- The repo is organized so the symbolic domain is created outside the model, while the neural components only see tokenized expressions and operation labels.
3. Data generation boundary
- data/generator.py is the source of synthetic training data.
- It produces structured trajectories of symbolic expression transitions, likely using SymPy-based transformations from one expression to the next.
- The generated dataset is stored as data/trajectories.json and is treated as an artifact rather than hand-authored data.
4. Tokenization layer
- data/tokenizer.py is the interface between symbolic strings and neural inputs.
- Expressions are represented in prefix / Polish notation, then mapped into a fixed token vocabulary.
- This layer is important because it defines the model’s input boundary: symbolic syntax is normalized into discrete IDs before entering the encoder.
5. Expression encoder
- model/encoder.py implements ExprEncoder.
- Its job is to convert a tokenized expression into a compact latent vector via embedding lookup + masked mean pooling.
- This is the first learned representation of the symbolic state; it abstracts away sequence length while preserving token content.
6. Transition core / world-state memory
- model/transition.py contains the recurrent transition model, implemented as a 2-layer GRU.
- The transition input is the concatenation of the encoded expression latent and an operation embedding.
- The GRU hidden state is the central “world state” (h_t): it carries procedural memory across steps and is the main recurrent subsystem in the architecture.
7. Prediction heads
- model/heads.py defines two outputs from the recurrent state:
- a latent predictor head that predicts the next expression’s latent embedding
- an operation predictor head that predicts the next operation class
- These are separate supervision targets, so the model learns both next-state structure and action classification.
8. Full model composition
- model/world_model.py wraps the encoder, transition core, and heads into one end-to-end forward pass.
- This is the main integration point where the symbolic input, action, recurrent state, and prediction outputs meet.
- model/test_model.py provides a smoke test for tensor shapes and forward-pass correctness.
9. Training subsystem
- train.py orchestrates dataset loading, batching, forward/backward passes, loss computation, and checkpointing.
- The README indicates it is MPS-aware, meaning it is designed to run efficiently on Apple Silicon using PyTorch’s MPS backend.
- Training outputs include checkpoints/model.pt and training_curves.png, making training both reproducible and inspectable.
10. Latent-space analysis / probing
- probe.py is a post-training analysis tool that extracts hidden states and evaluates how much reasoning structure is encoded in them.
- It produces PCA plots and linear probe metrics, stored under probe_results/.
- This analysis subsystem is separate from training, which makes the repo not just a model implementation but also an experimental framework for representation studies.
11. Configuration and reproducibility
- config.py centralizes hyperparameters, which keeps model size, training settings, and dataset-related choices in one place.
- requirements.txt pins the Python dependencies needed for generation, training, and probing.
- This separation suggests a clean distinction between experiment configuration and core model code.
12. Key architectural boundary
- The most important boundary is between symbolic mathematics and learned latent dynamics.
- On the left side are discrete expressions, operations, and generated trajectories; on the right side are embeddings, a GRU hidden state, and prediction heads.
- The whole repo is built to study whether the recurrent hidden state organizes information about reasoning chains, operations, and step progression.
13. Generated artifacts vs source code
- data/trajectories.json, checkpoints/, and probe_results/ are experiment outputs, not source-of-truth logic.
- They represent the three major runtime products of the system: the dataset, the trained weights, and the latent-space diagnostics.
Main entry points: world_model.py World Model · config.py Config · encoder.py Expr Encoder · transition.py Transition
Components 13
Data pipeline
Synthetic trajectory generation and tokenization boundary for symbolic expressions.
Model core
Encoder, recurrent transition state, and prediction heads.
- Expr Encoder · Embeds token sequences and pools them into a compact expression vector.
model/encoder.py - Transition · Recurrent world-state core that updates hidden state from expression and action inputs.
model/transition.py - Heads · Predicts next expression latent and next operation class from the recurrent state.
model/heads.py - World Model · End-to-end composition of encoder, transition core, and output heads.
model/world_model.py
Runtime
Training and probing entrypoints plus shared configuration.
- Config · Central experiment settings shared by generation, training, and probing.
config.py - Train · Loads data, optimizes the model, and writes training artifacts.
train.py - Probe · Extracts hidden states and evaluates latent structure with probes and PCA.
probe.py - Model Test · Checks tensor shapes and forward-pass correctness.
model/test_model.py
Artifacts
Generated outputs used for inspection rather than source logic.
- Trajectories · Synthetic symbolic transition data produced by the generator.
- Probe Results · Saved latent-space diagnostics and linear probe reports.
probe_results - Probe Report · Text summary of linear probe evaluation results.
probe_results/linear_probe_report.txt
Connections 17
- Generator → Trajectories: writes · Generates structured symbolic transition data.
- Trajectories → Tokenizer: feeds · Dataset expressions are normalized into token sequences.
- Tokenizer → Expr Encoder: encodes · Token IDs are embedded and pooled into latent vectors.
- Expr Encoder → Transition: state input · Expression latent becomes part of the recurrent transition input.
- Config → Generator: settings · Controls generation choices and dataset shape.
- Config → Train: settings · Provides training hyperparameters and backend choices.
- Config → Probe: settings · Shares analysis configuration with probing workflows.
- Transition → Heads: hidden state · Recurrent state is decoded by the prediction heads.
- Heads → World Model: outputs · Predictions are assembled into the full model forward pass.
- Expr Encoder → World Model: component · Encoder is wrapped by the end-to-end model.
- Transition → World Model: component · Transition core is wrapped by the end-to-end model.
- Heads → Train: losses · Training optimizes latent and operation prediction targets.
- World Model → Train: trained by · Train script drives forward/backward passes through the model.
- World Model → Probe: hidden states · Probe extracts recurrent representations for analysis.
- Probe → Probe Results: writes · Saves plots and probe metrics.
- Probe Results → Probe Report: contains · Report file stores linear probe evaluation summary.
- Model Test → World Model: validates · Smoke tests the integrated model interface.