experience

Instrumentation for training an autoregressive language model toward self-modeling, temporally integrated internal structure — the engineering foundation (M0) of the plan described in GOAL.md and PLAN.md.

This package does not train a model. It provides the capture, intervention, token-channel, auxiliary-head, and loss-registry machinery that every later phase builds on, targeting Llama-family causal LMs (Llama, Mistral, SmolLM, Qwen, Gemma).

Layout

src/experience/
  config.py              shared configuration dataclasses
  model_utils.py         Llama-arch layer resolution
  capture/
    hooks.py             ActivationCapture + CaptureBuffer (forward hooks)
    stats.py             on-the-fly statistics registry
    serialization.py     compressed zarr v3 capture store
  intervention/
    engine.py            InterventionEngine context manager
    directions.py        random / fixed / basis / cross-example directions
  tokens/
    interoceptive.py     learned (stat, bin) embedding channel
    selfstate.py         self-state summary spans + rolling window
  heads/
    base.py              AuxHead framework + strip-export
    projector.py         ResidualProjector / LinearReadout / MLPReadout
  losses/
    registry.py          LossRegistry with per-loss weight schedules
    experience.py        Phase 1 losses (predict / slow / variance)
  export.py              strip-export convenience

Usage

Capture activations

from experience import ActivationCapture, CaptureConfig

with ActivationCapture(model, CaptureConfig(layers="all")) as cap:
    outputs = model(batch, output_hidden_states=True)

h = cap.last_residual(0)      # (B, T, D) residual stream after layer 0

Residual-stream captures are always available. Per-head attention weights are captured only under an eager attention implementation (SDPA/flash never materialise the matrix), so pass attn_implementation="eager" when building the model if you need attention statistics.

Compute statistics

from experience import compute_stats, StatContext

stats = compute_stats(
    ["next_token_entropy", "top_decile_mass", "residual_norm"],
    StatContext(logits=outputs.logits, hidden_states=hidden_by_layer),
)

Apply interventions

from experience import InterventionEngine, InterventionConfig, RandomDirectionSource

engine = InterventionEngine(model)
with engine.intervene([
    InterventionConfig(layer=6, mode="add", alpha=1.0,
                       direction=RandomDirectionSource(hidden_size)),
]):
    outputs = model(batch)
# hooks removed on exit; model is untouched afterwards

Auxiliary heads and strip-export

from experience import attach_aux_head, ResidualProjector, export_clean_model

attach_aux_head(model, "proj", ResidualProjector(hidden_size, latent_dim=128))
# ...train...
clean = export_clean_model(model, path="./checkpoints/m1")

Auxiliary heads are attached under model._aux and never alter the base forward pass, so stripping them returns a model with identical token outputs.

Loss registry

from experience import LossRegistry, experience_loss, linear_warmup

registry = LossRegistry()
registry.register("experience", lambda inputs: experience_loss(
    inputs["z"], inputs["predictions"], inputs["targets"]),
    weight=linear_warmup(500))
loss_values = registry.compute(step=123, inputs=inputs)
loss_values.total.backward()

M1 + M2 training scripts

scripts/train_m1.py / scripts/eval_m1.py implement Phase 1 (LoRA + residual projector + interoceptive channel); scripts/train_m2.py / scripts/eval_m2.py implement Phase 2 (intervention self-supervision / counterfactual self-modeling). See RUNBOOK.md for the full run/interpretation guide.

uv run python scripts/train_m1.py --config configs/m1.yaml          # Phase 1
uv run python scripts/train_m2.py --config configs/m2.yaml --smoke  # Phase 2 CPU check
uv run python scripts/eval_m1.py --config configs/m1.yaml --mode probe \
    --base unsloth/Llama-3.1-8B-Instruct --m1 runs/m1/exported

Development

uv sync
uv run pytest

Tests run against a tiny randomly-initialised Llama model — no weights are downloaded.