swgoo/pmnet_consolidator

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:

Direct LTM routing does not affect the immediate STM readout: both routing conditions reach 89.90 ± 0.00%. Its effect appears after consolidation:

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:

memory episode → routed STM writes → Consolidator(STM) → persistent LTM update

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

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:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt

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:

docker run -it \
  --name pmnet-consolidator \
  --gpus all \
  --network host \
  --mount type=bind,source="$PWD",target=/root/pmnet_consolidator \
  --workdir /root/pmnet_consolidator \
  docker.io/pytorch/pytorch:2.10.0-cuda13.0-cudnn9-devel \
  bash

Inside the container, install the same system and Python dependencies as the Dev Container:

apt-get update
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
  latexmk \
  texlive-latex-base \
  texlive-latex-recommended \
  texlive-fonts-recommended \
  lmodern \
  poppler-utils
python -m pip install --break-system-packages -r requirements.txt

After exiting, reopen the configured container without reinstalling its dependencies:

docker start -ai pmnet-consolidator

The bind-mounted repository remains on the host. Delete and recreate the container when changing the base CUDA or PyTorch image.

Checkpoints

The Hugging Face repository provides the same project together with the released weights.

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:

python train_pmnet_ablation.py \
  --config_path config.json \
  --starting_checkpoint checkpoints/phase1_rule_pretrain.ckpt \
  --auto_rule_pretraining false \
  --experiments learned_full,identity_full,learned_consolidator_only,learned_consolidator_only_ltm_routing_off,learned_memory_consolidator,learned_full_dual_objective \
  --seeds 42,43,44,45,46 \
  --checkpoint_dir outputs/pmnet_consolidation_ablation \
  --num_devices 1 \
  --use_wandb false

Central frozen-model intervention

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.

python train_pmnet_ablation.py \
  --starting_checkpoint checkpoints/phase1_rule_pretrain.ckpt \
  --auto_rule_pretraining false \
  --experiments learned_consolidator_only \
  --seeds 42,43,44,45,46 \
  --checkpoint_dir outputs/central_intervention \
  --num_devices 1 \
  --use_wandb false

Direct-LTM-routing ablation

This comparison isolates the contribution of the forward route from LTM to the write module.

python train_pmnet_ablation.py \
  --starting_checkpoint checkpoints/phase1_rule_pretrain.ckpt \
  --auto_rule_pretraining false \
  --experiments learned_consolidator_only,learned_consolidator_only_ltm_routing_off \
  --seeds 42,43,44,45,46 \
  --checkpoint_dir outputs/ltm_routing_ablation \
  --num_devices 1 \
  --use_wandb false

Dual-objective diagnostic

The dual-objective run is reported as a diagnostic rather than part of the default suite:

python train_pmnet_ablation.py \
  --starting_checkpoint checkpoints/phase1_rule_pretrain.ckpt \
  --auto_rule_pretraining false \
  --experiments learned_full_dual_objective \
  --seeds 42,43,44,45,46 \
  --checkpoint_dir outputs/dual_objective \
  --num_devices 1 \
  --use_wandb false

Experiment-name map

The default suite excludes learned_full_dual_objective; request it explicitly when reproducing the appendix diagnostic.

Evaluation only

To re-evaluate a selected phase-2 Lightning checkpoint without retraining, use the same experiment definition and seed that produced it:

python train_pmnet_ablation.py \
  --starting_checkpoint PATH/TO/learned_consolidator_only/best-STEP.ckpt \
  --experiment learned_consolidator_only \
  --seeds 42 \
  --evaluation_only true \
  --checkpoint_dir outputs/evaluation_only \
  --num_devices 1 \
  --use_wandb false

Outputs and provenance

Each suite is stored under a fingerprinted directory:

CHECKPOINT_DIR/RUN_GROUP-FINGERPRINT/
├── suite_config.json
├── ablation_runs.csv
├── ablation_summary.csv
├── ablation_results.json
└── EXPERIMENT/seed-SEED/
    ├── run_config.json
    ├── result.json
    ├── best-*.ckpt
    └── resume.ckpt

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:

python train_pmnet_ablation.py \
  --config_path config.json \
  --auto_rule_pretraining true \
  --experiments learned_consolidator_only \
  --seeds 42 \
  --checkpoint_dir outputs/from_scratch \
  --num_devices 1 \
  --use_wandb false

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.

License

Released under the MIT License.