A small, CPU-reproducible test of representation geometry in mixtures of finite hidden Markov sources, plus a causal-erasure extension.
Each sequence is generated by one of two HMMs, chosen once and held fixed. An exact Bayesian filter therefore has to track two things at every token: which HMM generated this sequence and where that HMM is internally. The code trains a small GRU and causal Transformer for next-token prediction, then asks whether final activations linearly encode the exact component posterior and component-conditional state posterior on held-out sequences. PCA is included as a visualization, not as evidence by itself.
The extension fits linear subspaces for those two kinds of information and projects activations out of each subspace. A separately trained evaluator probe on a third sequence split measures information after intervention, avoiding the tautology of evaluating with the probe whose own row space was removed. It also measures predictive NLL and KL from the exact Bayesian predictor. Random, norm-matched random, shuffled-label, and untrained-network controls make selective damage falsifiable.
Python 3.11 or newer is required. The current checked-in run used Python 3.14 and CPU-only PyTorch execution.
python3 -m venv .venv
source .venv/bin/activate
pip install -e '.[test]'
make smoke # seconds-scale end-to-end check, one seed
make train # central config, seeds 0/1/2, both models
make reproduce # held-out probes and PCA records (trains if needed)
make extension # controlled causal erasure (trains if needed)
make figures # reads only results/*.jsonl
make sweep-overlap # four overlap values × two models × three seeds
make sweep-length # lengths 8/16/32/64 × two models × three seeds
make sweep-components # 2/3/4 components × two models × three seeds
make sweep-width # widths 8/16/32/64 × two models × three seeds
make sweep-depth # Transformer block 1/block 2/final norm × three seeds
make sweep-interaction # overlap 0/.35 × length 8/64 × two models × three seeds
make sweep-context-restart # length-64 models: full prefix vs last 8 tokens, exact Bayes oracle
make sweep-position-restart # Transformer: last 8 tokens with reset vs original position indices
make sweep-short-context # length-9-trained Transformer on the same held-out eight-token windows
make sweep-budget-context # short training matched to long model's token and optimizer-step budgets
make sweep-gru-budget-context # GRU architecture check of the matched-budget effect
pytest -qmake reproduce and make extension reuse matching checkpoints when present. Delete checkpoints/reproduce/ to force retraining. make clean-results removes generated JSONL and PNG outputs but leaves checkpoints intact.
src/nonergodic_memory/data/hmm.py: sampling and exact Bayesian filtering.src/train.py: deterministic training and checkpoint production.src/probe.py: held-out classifiers, posterior regressions, shuffled labels, untrained controls, and PCA coordinates.src/intervene.py: learned, random, norm-matched, and shuffled-label subspace interventions.configs/: smoke and central CPU configurations.results/: raw JSONL records; every result states model, seed, condition, and device.figures/: regenerated exclusively from JSONL.results/sweep_overlap_*.jsonlandfigures/sweep_overlap.png: the registered source-overlap extension.results/sweep_length_*.jsonlandfigures/sweep_length.png: the registered sequence-length follow-up.results/sweep_components_*.jsonlandfigures/sweep_components.png: the registered component-count sweep.results/sweep_width_*.jsonlandfigures/sweep_width.png: the registered model-width sweep.results/sweep_depth.jsonlandfigures/sweep_depth.png: layerwise Transformer erasure propagated through the remaining network.results/sweep_interaction_*.jsonlandfigures/sweep_interaction.png: fresh matched overlap-by-context grid and paired interaction contrast.results/sweep_context_restart.jsonlandfigures/sweep_context_restart.png: held-out eight-token restarts with an elapsed-prior exact Bayesian oracle.results/sweep_position_restart.jsonlandfigures/sweep_position_restart.png: position-preserving Transformer restart control.results/sweep_short_context*.jsonlandfigures/sweep_short_context.png: short-input training control joined to the published position-restart raw file.results/sweep_budget_context*.jsonlandfigures/sweep_budget_context.png: token/step-matched short-input training control, joined to the two prior raw files.results/sweep_gru_*context*.jsonlandfigures/sweep_gru_budget_context.png: GRU architecture-generalization check for matched training exposure.STATE.md: the current hypothesis → experiment → interpretation loop.report.md: methods, results, negative results, and limitations.
After observing token x_t, the filter stores the component posterior p(c | x_0:t), every normalized within-component state posterior p(s_t | c,x_0:t), and the next-token prediction p(x_{t+1} | x_0:t). Neural logits at position t are evaluated against the same next token. Posterior regression targets the full K×S vector. The separate hard-state diagnostic supplies the true component to a component-specific classifier, preventing component mistakes from being counted twice.
Runs are deterministic on the tested CPU environment: data, initialization, batch order, probes, and random controls all use explicit seeds. Checkpoints are generated rather than versioned and are validated against the complete requested configuration, model, and seed before analysis. Partial CLI reruns atomically replace only matching model/seed cells. Raw records include a config hash and Python, NumPy, and PyTorch versions. Exact bitwise equality across different PyTorch/BLAS versions is not promised.
A final checkpoint-free clone audit ran every project command above, all eleven sweeps, and 65 tests. Numerical JSONL values reproduced exactly and every figure was byte-identical on the recorded CPU environment; details are in report.md.