mixle.experimental.retrieval_memory_spine module

E6: retrieval memory over frozen past – a ContextMechanism that pairs E1’s local sliding window with an unbounded-length, stop-gradient kNN index of everything that has scrolled out of the local window.

Why this exists. SlidingWindowSpine (E1) only ever attends to the last window tokens – anything older is simply gone, so needle-in-a-haystack facts planted before the window fall out of reach no matter how long training runs. RetrievalMemorySpine keeps the same local window for near-range recall, but ALSO archives only keys/values evicted from the raw local cache into a per-layer index carried in the state, and does a brute-force kNN lookup of that index for each query, attending over the top-retrieval_k hits alongside the local window in one combined softmax.

The non-differentiable boundary (read this before touching the backward pass). The index is written by .detach()``ed tensors from PAST steps -- steps whose own backward graph has already been consumed by ``train_tbptt’s TBPTT boundary. Nothing in this module tries to differentiate through how those entries were produced. What DOES stay exact: the retrieval and combination happening THIS step – topk selection of which entries to look at is a discrete, gradient-free op (like sparse/MoE routing), but the softmax attention over the selected top-k values is full-precision autograd, so gradients flow exactly (no straight-through / relaxation approximation) into this step’s query and output projections. Net effect: exact gradients through the retrieval OPERATION, zero gradient into the frozen index CONTENTS. Every step() call documents this on the returned state as state.receipt["differentiable_boundary"] (a receipt field, not just a docstring claim – see the roadmap card, notes/standout-roadmap-tasks.md E6).

State cost. The literal index tensors dominate the state’s byte footprint (O(total tokens streamed) unless max_index_tokens caps it), but backward-pass memory is O(window + retrieval_k) per query rather than O(index length) – the whole point of gathering only the top-k hits before running the differentiable softmax. The state receipt reports archived tokens, cache/index overlap, and current index length so comparisons with other context mechanisms can use measured state rather than a placeholder.

class RetrievalMemoryState(cache_k=<factory>, cache_v=<factory>, index_k=<factory>, index_v=<factory>, index_pos=<factory>, batch_size=0, pos=0, receipt=<factory>)[source]

Bases: object

Per-layer local window cache (same shape/convention as SlidingWindowState) plus a per-layer detached kNN index of every earlier chunk’s post-RoPE keys/values.

cache_k stores raw, unrotated keys so every key receives RoPE exactly once per attention evaluation. index_k stores detached post-RoPE keys only after their positions leave that cache. index_pos is the exact position ledger used to prove the cache and index are disjoint.

receipt: honest, per-step bookkeeping – see RetrievalMemorySpine.step(). Carried forward on the state (rather than a third return value) because ContextMechanism fixes step’s return shape to (new_state, mean_loss); the state IS this mechanism’s output.

Parameters: