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 training horizon.

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 every processed chunk’s post-RoPE keys/values (detached) into a per-layer index carried in the state, and on every subsequent step 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. notes/standout-roadmap-tasks.md’s E6 card asks for this mechanism’s state cost “at a fraction of E2’s” (moment-closure attention). E2 does not exist on any branch reachable from this worktree’s base as of this writing (see RETRIEVAL_MEMORY_UNAVAILABLE_PIECES below, matching the convention mixle/task/pilot_ladder.py uses for roadmap pieces it cannot reach) – there is nothing to compare against yet. mixle/tests/retrieval_memory_spine_test.py instead measures and asserts this mechanism’s OWN state bytes-per-token, honestly, and leaves the E2 ratio as a documented placeholder.

RETRIEVAL_MEMORY_UNAVAILABLE_PIECES: dict[str, str] = {'E2': "moment-closure attention (roadmap E2) does not exist on any branch reachable from this worktree's base (release/0.7.0 -> chunked-recurrent-spine -> long-context-referee) as of this writing -- it was being built in parallel on its own branch and never reached origin. The E6 card's acceptance criterion ('at a fraction of E2's state cost') cannot be checked against a real E2 measurement; this module instead reports RetrievalMemorySpine's own measured state bytes-per-token (see mixle/tests/retrieval_memory_spine_test.py) and leaves the E2 ratio as a documented placeholder rather than a fabricated number."}

roadmap sub-pieces this module cannot reach from this worktree’s base, and exactly why – see the module docstring and mixle/task/pilot_ladder.py’s PILOT_LADDER_UNAVAILABLE_PIECES for the same convention.

class RetrievalMemoryState(cache_k=<factory>, cache_v=<factory>, index_k=<factory>, index_v=<factory>, 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.

index_k/index_v: (batch, index_len, n_head, head_dim) or None before the first archive. Entries are appended once per step call (this chunk’s own keys/values, never the local window’s carried-over tail – that would double-archive the same tokens every step). Always detached at write time; see the module docstring for the non-differentiable-boundary contract this enforces.

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:
class RetrievalMemorySpine(vocab, *, d_model=32, n_layer=2, n_head=2, window=64, retrieval_k=4, max_index_tokens=None)[source]

Bases: Module

E6: E1’s local sliding window plus a brute-force kNN retrieval index over detached past chunks.

window: local causal attention span, identical semantics to SlidingWindowSpine.window. retrieval_k: how many index entries each query attends over (the “top-k” of the E6 card). max_index_tokens: FIFO cap on total archived tokens per layer (None = unbounded). Caps the brute-force kNN’s O(chunk * index_len) score matrix and the state’s byte footprint; oldest entries are evicted first once the cap is exceeded.

Parameters:

  • vocab (int)

  • d_model (int)

  • n_layer (int)

  • n_head (int)

  • window (int)

  • retrieval_k (int)

  • max_index_tokens (int | None)

detach(state)[source]

Stop-gradient the local window cache (cuts the TBPTT graph, same as E1). The index is already detached at write time (see step()), so this only needs to re-detach the cache.

Parameters:

state (RetrievalMemoryState)

Return type:

RetrievalMemoryState