Finite Memory for Unbounded Streams

9 min read

A two-tier KV memory for unbounded token streams — exact where it matters, merged where readers can't tell, and trained with the memory in the loop.

Contents

Finite memory for unbounded streams

Topline result: a two-tier KV memory — an exact episodic ring in front of a bank of merged centroids — lets a small transformer generate coherently at 8× its trained context with a fraction of the cache, and beats both plain cluster memories at twice the budget and the standard sinks+window eviction baseline. Its distinguishing mechanism is Nerode merging: centroids are fused only when recent queries cannot tell them apart, a criterion that appears to be unclaimed in the KV-cache compression literature. Fine-tuning with the memory in the loop then closes most of the remaining quality gap: validation cross-entropy 1.44 vs 2.61 for frozen exact attention at the same doubled context (Training in the Loop).

Every load-bearing number on this site is recomputed from the code in this repository — the inference engine, the training pipeline, and the data pipeline are all here, and the build re-runs them (Verification Methodology).

The mechanism, animated

Tokens stream in from the left into a small episodic ring — the last 16 raw key/value pairs, kept exact. When the ring is full, the oldest entry demotes into the centroid bank on the right: it either merges into a compatible centroid (growing its mass) or, if it is unlike everything stored, births a new one. Centroids leak mass and die; periodically, two centroids that recent queries cannot distinguish fuse (a Nerode merge, flashed white), freeing a slot for new content. One softmax reads the union.

Schematic of the two-tier memory. Dot color stands for key content; centroid size is accumulated mass. Pause and watch a merge free a slot.

Watch the memories diverge

The stream harness generates 2048 tokens — 8× the trained context — and reports the distinct-token ratio per 128-token window: a collapsing ratio is "mush" made measurable. This chart is drawn from the verified computation (each curve is the mean over three seeds), not from a screenshot:

Distinct-token ratio per 128-token window, 2048-token stream (3-seed mean)
Higher is healthier. The sinks+window eviction baseline (a larger memory budget) repeatedly crashes toward 0.3 as its window rolls off; the frozen tiered memory sags but never collapses; the fine-tuned tiered model is the flattest, healthiest line on the chart — the same checkpoint sampled in the final section below. Exact attention cannot run here at all: its cache ends at position 256.

Hear it from the model

Both samples below are the engine's actual output, injected at build time from the verified checks — not copy-pasted. First, the tiered memory well inside its comfort zone (greedy decoding, frozen weights):

Greedy sample, tiered memory, positions 0–120
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 the sun! She thought it was so pretty. Lily wanted to play with the ball, but it was too high up in the sky. She tried to jump, but she couldn't reach it. Then, she saw a big, strong man walking by. She asked him if he could help her get the ball. The man said, "Sure, I can help

And the hard regime: the last ~160 tokens of a 2048-token sampled stream, frozen weights, 8× past the trained context. The tiered memory degrades but keeps producing story-shaped text; the eviction baseline at a larger budget slides into repetition — the texture behind the distinct-ratio curves above. (Fine-tuning with the memory installed is what turns the tiered sample coherent.)

Deep-stream tails at position ~1900, seed 1: tiered vs eviction
=== tiered memory (m=64,e=16, nerode merging) === that day was all over the whole day was ever the day of the day was on the day of time. The most healed away. He was in the time. Once there was running and heel. He was just right. He was very fast and wanted to run. He started to do so much lined. He was so happy. The to this was all he was sick! One day. The day was so happy. He was very excited. He had a whole day had been chest the time. He had to remember. He loved to stay. The boy. He remembered. He was always. The day. He did! Once upon a better than hears. He never gave it. Even the day. Every day! But he didn === sinks+window eviction baseline (s=4,w=124) === ter, butter. At night, the sun was bright and warm. The sun came out, butter was very far away. The sky was clear, and the stars were bright. Lily liked the stars, but they were not bright. They were dull. Lily and her mommy went to bed. They looked at the stars one more time, and the night was shiny. The stars were shining bright, and they made the night very bright. Lily smiled and knew that even dull things can be bright and pretty. Once upon a time, there was a little boy named Timmy. Timmy loved to play outside, but he didn't like to eat his vegetables. His mommy would always tell him that he had to eat

And then you train it

The frozen samples above are a retrofit: weights that have never seen a compressed memory, reading through one anyway. The last step of the story is to put the memory in the training loop — 300 fine-tuning steps on TinyStories with the tiered memory installed (truncated BPTT, bank state detached per token), taking validation cross-entropy from 2.06 to 1.44, then exporting the weights back to llama2.c format and loading them in the Java engine. Same position ~1900, same seed, same finite memory — different model:

Deep-stream tail at position ~1900, seed 1 — after fine-tuning with the memory in the loop
I threw it too. I hope I will help you find it." The blue ball was sad, but it was happy to have a new friend. They both learned that it is okay not to always have a blue ball to play with. Once upon a time, in a small town, there lived a boy named Tom. Tom was a good boy who liked to play in the park near the library. One day, Tom went to the library with his mom. At the library, Tom found a book about animals. He loved the book and wanted to read it every day. The library became his favorite place. One day, Tom was reading a book about animals. He read about the animals in the library. He read about the people he loved to read about the famous library.
Generated by the 300-step fine-tuned checkpoint (fetched and re-run by the build; provenance in tune-tiered.log) with the same 80-slot tiered memory, 8× past the trained context.

Coherent multi-sentence stories, held together by 64 centroids and a 16-slot ring. That is the end-to-end result: an engine where attention is swappable, a memory that is exact where it matters and merged where readers can't tell, a harness that makes degradation measurable — and a training loop that teaches the model to live inside its finite memory. Training in the Loop has the full protocol and ablations.

The result

  • Tiered Memory — the architecture and its measured behaviour: exactness, short-range fidelity, stream health at 2048 tokens, and the merge/gate ablation.
  • Training in the Loop — fine-tuning with the compressed memory installed, and the PyTorch → llama2.c → engine round trip.
  • Positioning — what is new against the literature, what is an increment, and what is already taken.
  • Related Work — the KV-cache compression landscape: sixteen verified papers, organized by what decides which memories survive.

Supporting material

  • The Engine — the inference engine: attention (including the KV cache) as a plugin, and the logit-diff correctness harness.
  • Lattice Memory and K-Means Memory — the two designs that preceded the tiered memory; their ablations isolate attention over clusters as the load-bearing idea and motivate every tiered design choice.
  • TinyStories Data — the token pipeline, byte-reproducible from the public dataset.
  • Verification Methodology — what "verified" means here: the probe, the stream metric, and the verified/attested boundary.

Headline numbers, recomputed by this build

Below capacity the k-means memory is exact attention (max |Δlogit| 0.000000 vs the causal reference). On 2048-token streams the tiered memory holds a minimum distinct-token ratio of 0.38/0.49/0.38 across three seeds — against 0.22/0.28/0.32 for the StreamingLLM-style eviction baseline at a larger memory budget. And twenty verified fine-tuning steps with the memory in the loop move validation CE from 4.2283 to 3.3432.