Factual Memory Diffing

Research project investigating how factual knowledge is encoded in language model weights. By training two versions of the same model checkpoint — one with injected factual data and one without — and comparing them across weight space, representations, and predictions, we identify where and how facts get stored.

Overview

The pipeline has three stages:

1. Data Generation — Use GPT-4o to produce (subject, relation, object, year) triplets for real-world events (politics, sports, CEOs), then verbalize them into natural language sentences used for training and evaluation.
2. Training — Continue training a Pythia-160m-deduped checkpoint on Pile data, with or without factual injection, using distributed training.
3. Analysis — Compare the two resulting model checkpoints across four axes: weight deltas, representational similarity (CCA), prediction divergence (KL), and causal localization via activation patching.

Repository Structure

factual-memory-diffing/
├── data/                        # Generated datasets
│   ├── triplets.json            # (subject, relation, object, year) triplets
│   ├── verbalizations.train.json
│   └── verbalizations.test.json
├── data_generation/             # Synthetic data pipeline
│   ├── scripts/
│   │   ├── generate_triplets.py
│   │   └── verbalize_triplets.py
│   └── src/                     # TripletGenerator, Verbalizer classes
├── src/                         # Training infrastructure
│   ├── distributed_train.py     # DDP training loop with optional injection
│   ├── memorization_utils.py    # Per-token perplexity, memorization detection
│   ├── generation_utils.py      # Batched generation helpers
│   └── nparray_dataset.py       # Numpy-backed dataset for Pile shards
├── analysis/
│   └── compare_models_fixed.py  # Four-axis model comparison tool
├── scripts/                     # Convenience training scripts
│   ├── train_noinject.py
│   ├── train_with_injection.py
│   └── train_with_injection_single_shot.py
├── models/                      # Local model checkpoint directory
├── results/                     # Inference reports (JSON)
├── inference.py                 # Factual recall evaluation script
└── train_noinject.py            # Entry point for no-injection training run

Setup

# Install dependencies (requires Python >=3.12)
uv sync

# Or with pip
pip install torch transformers openai tiktoken pyyaml

Usage

1. Generate Data

cd data_generation

# Generate triplets
uv run --env-file .env -- python scripts/generate_triplets.py \
  --num-triplets 100 \
  --batch-size 35 \
  --model gpt-4o \
  --domains people-of-importance-politics people-of-importance-ceo people-of-importance-sports \
  --after-year 2020 \
  --output ../data/triplets.json

# Verbalize into natural language sentences
uv run --env-file .env -- python scripts/verbalize_triplets.py \
  --triplets ../data/triplets.json \
  --inference standard \
  --output ../data/verbalizations.json

2. Train Models

Train the baseline (no injection) model:

python train_noinject.py \
  --checkpoint pythia-160m-deduped-step80000 \
  --pile_data_path pile_80k.npy pile_81k.npy pile_82k.npy \
  --log_dir runs/noinject

Train with factual injection (see scripts/train_with_injection.py for arguments).

3. Evaluate Factual Recall

python inference.py \
  --test_data data/verbalizations.test.json \
  --checkpoint pythia-160m-deduped-step80000 \
  --output results/inference_report.json

Output metrics:

• Exact match — model output equals expected answer
• Semantic match — both subject and object appear in response
• Partial match — at least one key entity appears in response
• Overall score — confidence-weighted aggregate

4. Compare Models

python analysis/compare_models_fixed.py \
  --model-a path/to/injection_model \
  --model-b path/to/noinject_model \
  --probe-prompts prompts.txt \
  --memorized-seqs memorized.txt \
  --layers 0 6 12 18 24 \
  --device cuda \
  --run-patching

Comparison axes:

1. Weight-space diff — Frobenius norm of parameter deltas per layer/module
2. Representation similarity — per-layer CCA on a probe set (optional PCA pre-reduction)
3. Prediction divergence — layerwise KL divergence on memorized sequences (LogitLens-style)
4. Causal verification — activation patching to localize which layers carry the factual signal

Data Format

Triplets (triplets.json):

{
  "subject": "Kamala Harris",
  "relation": "became",
  "object": "Vice President of the United States",
  "event_year": 2021,
  "domain": "politics"
}

Verbalizations include multiple paraphrases per triplet, used as training exposures and Q&A test pairs.

Dependencies

Package	Purpose
torch	Model training and inference
transformers	Pythia model + tokenizer
openai	GPT-4o data generation
tiktoken	Token counting for data generation
pyyaml	Configuration files

License

See LICENSE.