Consolidator: Learning Persistent Routed Memory Across Context Boundaries
Official implementation of the Consolidator experiments for PMNet. This
repository studies whether a small learned state-transition module can turn
short-term memory (STM) writes into persistent long-term memory (LTM) updates
without updating the frozen backbone during the memory episode.
The intervention is deliberately narrow: the PMNet backbone is frozen and only
the 12.35K-parameter Consolidator is trained. In the reported 29.95M-parameter
model, this corresponds to 0.041% trainable parameters. PMNet was introduced in
Phasor Memory Networks.
A matching repository that includes the released model weights is available on
Hugging Face.
Main results
Five-seed results on the fixed held-out test stream:
Condition
Updated-LTM accuracy
Learned Consolidator
87.02 ± 1.76%
Same checkpoint, Consolidator forced to identity
18.32 ± 0.04%
Learned Consolidator, direct LTM routing disabled
44.38 ± 1.94%
Direct LTM routing does not affect the immediate STM readout: both routing
conditions reach 89.90 ± 0.00%. Its effect appears after consolidation:
Rule family
Routing off
Routing on
Paired gain
ADD
52.52 ± 1.37%
99.01 ± 0.60%
+46.49 ± 1.23 pp
AFFINE
36.24 ± 2.61%
74.80 ± 3.07%
+38.56 ± 1.29 pp
These are controlled proof-of-concept experiments. They establish the
mechanism within the evaluated synthetic setting; they do not by themselves
establish natural-language or large-scale performance.
Mechanism
During a memory episode, the model writes to STM. At the episode boundary, the
Consolidator transforms the final STM state and applies the result to LTM:
The decisive comparison evaluates the same trained checkpoint twice—once with
the learned Consolidator and once with the transition replaced by identity.
The routing ablation separately disables the forward path from LTM to the write
module while leaving the rest of the learned system unchanged.
Repository layout
Path
Description
modeling_pmnet.py
PMNet, STM/LTM routing, and Consolidator implementation
configuration_pmnet.py
Model configuration, including the LTM-routing switch
train_pmnet_ablation.py
Canonical training and evaluation entry point
config.json
Default model configuration
Installation
The reported experiments used Python 3.12.3, PyTorch 2.10.0+cu130,
Transformers 5.14.1, Lightning 2.6.5, bf16 mixed precision, FlashAttention 2,
and one NVIDIA RTX 4090.
Local environment
Create an isolated environment and install the release dependencies:
The FlashAttention wheel in requirements.txt targets the reported CUDA,
PyTorch, Python, and C++ ABI combination. Replace that wheel with a compatible
FlashAttention build when using a different environment.
Docker
The repository's Dev Container configuration
uses pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel. On a Linux host with the
NVIDIA Container Toolkit installed, create an equivalent persistent Docker
container from the repository root:
bash
1docker run -it \2 --name pmnet-consolidator \3 --gpus all \4 --network host\5 --mount type=bind,source="$PWD",target=/root/pmnet_consolidator \6 --workdir /root/pmnet_consolidator \7 docker.io/pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel \8bash
Inside the container, install the same system and Python dependencies as the
Dev Container:
All phase-2 conditions must start from the same phase-1 rule-pretraining
checkpoint. Place the released checkpoint at:
checkpoints/phase1_rule_pretrain.ckpt
For an auditable reproduction, verify the checkpoint against the SHA-256 hash
published with the release asset. The runner also records the loaded
checkpoint hash in each run's provenance.
If no phase-1 checkpoint is supplied, the entry point can train it once and
reuse its selected checkpoint across the complete phase-2 seed grid. See
Training phase 1 from scratch.
Reproducing the paper experiments
All commands below disable W&B logging and explicitly use one device. Remove
--use_wandb false if experiment tracking is desired.
Core suite and appendix diagnostic
This command runs the five core conditions plus the opt-in dual-objective
appendix diagnostic from the shared phase-1 checkpoint:
This is the primary learned-versus-identity intervention. It trains only the
Consolidator while freezing the backbone and memory parameters. The selected
checkpoint is then evaluated both normally and with its Consolidator forced to
identity, so the comparison does not require a separately trained condition.
The runner records the experiment configuration, random seeds, model and
optimizer settings, parameter counts, checkpoint hash, software versions,
early-stopping state, selected checkpoint, and test results. Validation selects
the checkpoint; final metrics are computed on one fixed held-out test stream
shared across conditions.
Training phase 1 from scratch
To recreate the shared rule-pretraining checkpoint instead of loading the
released one:
The phase-1 stage initializes PMNet from config.json, runs once with seed 42
by default, selects its best validation checkpoint, and shares that checkpoint
across the requested phase-2 runs.
Tests
Run the CPU regression suite with:
python -m unittest discover -s tests -v
List all training and evaluation options with:
python train_pmnet_ablation.py --help
Implementation scope
The current implementation supports default RoPE, eager or Flash Attention,
unpadded causal sequences, and structured Transformers outputs. Unsupported
attention modes, padding and cache transformations fail explicitly rather than
silently changing the evaluated computation. Current-schema PMNet checkpoints
are required; older routed/ring-memory checkpoints are not shape-compatible
with this implementation.
