Training in the Loop
Fine-tuning the model with the compressed memory installed — and round-tripping the weights back into the Java engine.
Training in the Loop
The reindex refutation left a theory on the table: the
failure was calibration, not mechanism — a model tuned with cluster memory
in the loop should learn its own suppression and then exploit coherent
re-aged keys. Testing that requires a training pipeline with the memory
installed: a PyTorch port of the model and the streaming cluster attention
(code/attention/python/common.py, semantics matching the Java
implementation), fine-tuned on TinyStories Data with truncated BPTT
(bank state detached per token), then exported back to llama2.c format and
judged in The Engine. Retrofitting a compressed memory by fine-tuning
has precedent — Infini-attention and DMC both do it (Related Work) —
which is why the novelty analysis rests on the merge
criterion, not on in-the-loop tuning itself.
The pipeline, verified end-to-end
The build re-runs the whole loop at reduced scale — same code, same seeds, 20 steps instead of 300:
How we know this
| val[frozen | kmeans-raw | (before)] | CE | = | 4.2283 |
|---|---|---|---|---|---|
| step | 0 | loss | 4.2118 | lr | 1.00e-05 |
| step | 10 | loss | 3.4479 | lr | 4.22e-05 |
| val[tuned | kmeans-raw | (final)] | CE | = | 3.3432 |
| exported | ./stories15M-raw.bin | ||||
| before_ce=4.2283 | final_ce=3.3432 | ||||
| --- greedy sample from exported checkpoint (java | kmeans:m=64) --- | ||||
| Once upon a time | there was a little girl named Lily. She loved to play outside in the sunshine. One day | she saw a big | red ball in the sky. It was so pretty! She wanted to touch |
Frozen baselines
Validation CE at T=512 (2× the trained context, teacher-forced), recomputed by the build:
| config | ce |
|---|---|
| frozen-exact | 2.6080 |
| frozen-kmeans-raw | 4.2283 |
| frozen-kmeans-reindex | 3.9313 |
| frozen-tiered | 2.0644 |
Two things stand out. Frozen cluster memories pay a real tax (raw 4.2283 vs exact 2.6080). And the frozen Tiered Memory retrofit already beats frozen exact attention (2.0644) — the exact episodic ring supplies sharp local attention without the RoPE-extrapolation penalty that inflates exact attention at 2× context.
The full runs (attested)
The complete 300-step runs took hours of CPU and exceed the build's per-check budget, so their logs are committed as attested artifacts (experiment.log, tune-tiered.log):
| config (val CE at T=512) | CE |
|---|---|
| frozen exact attention (RoPE extrapol.) | 2.608 |
| frozen kmeans-raw m=64 | 4.228 |
| frozen kmeans-reindex m=64 | 3.931 |
| tuned kmeans-raw (300 steps) | 2.713 |
| tuned kmeans-reindex (300 steps) | 2.493 |
| frozen tiered m=64,e=16,g=0.5 | 2.064 |
| tuned tiered (300 steps) | 1.444 |
(The frozen rows are independently verified above; the tuned rows come from the logs.)
Conclusions from the full runs:
- The calibration theory held — modest CPU fine-tuning rehabilitated reindex completely; frozen-collapsed configurations stream coherently after tuning.
- With training in the loop, reindex beats raw (2.49 vs 2.71), reversing the frozen-retrofit ranking: position handling and forgetting are coupled, and the model learns the suppression that raw phase interference provided for free.
- Tuned reindex with an 8×-compressed memory beats frozen exact attention at T=512 (2.49 vs 2.61), because re-indexed relative positions never extrapolate.
- Tuning the tiered memory takes CE to 1.44, approaching native-context quality with a 64-centroid compressed long-range memory.