mixle.experimental.quantized_key_attention module

E10: group-structured quantized-key attention – an ADAPTIVE far-field state whose compression is exact by construction once keys are quantized.

Keys are product-quantized to a codes_per_block^n_blocks lattice (n_blocks sub-vectors, one learned codebook each; the per-block simplex perturbation group is isomorphic to additive logits, which is what makes “cell” a group-theoretic object and not just a bucket). Attention weights are then CELL-CONSTANT over the far field, so the softmax over every evicted token collapses exactly to

sum_cells count_c * exp(q . k_c / sqrt(d)) * vbar_c —————————————————— (values stay continuous!) sum_cells count_c * exp(q . k_c / sqrt(d)) + near mass

– an identity, not an approximation (falsified 2026-07-12: max error 4e-17 numpy / 1e-9 on a trained torch model; experiments/group_attention/RESULTS.md). Per-query far-field cost is O(occupied cells), independent of context length; the trained falsification model occupied 49 of 65,536 possible cells, flat from 4k to 64k context.

Three deliverables live here:

• QuantizedKeyAttentionSpine – a ContextMechanism: exact RoPE’d softmax over a sliding near window, plus a bounded far-field CELL STORE (integer counts, value sums, code vectors) that evicted tokens fold into. Near and far are normalized JOINTLY in one softmax (log count enters as a logit offset), so the whole mechanism IS softmax attention over the quantized stream – unlike the E3 sketches’ two separately-normalized branches. Keys quantize with a VQ-VAE commitment loss (straight-through estimator), added to the training loss only when gradients are enabled so eval probes still return pure cross-entropy.

• SlidingCellWindow – streaming windowed aggregation where eviction is EXACT integer subtraction in the group Z^cells (counts can never drift; value sums are float64 running sums).

• CellCountTree – a Fenwick tree of sparse cell-count nodes: O(log n) updates, O(log n)-node prefix merges, and arbitrary-window counts via prefix(hi) - prefix(lo) using the group inverse (integer subtraction), never a float cancellation.

Positional honesty: the far field is position-free – a bag of typed content. Retrieval across the far field is by content (associative recall), not by position; tasks that need positional order beyond the near window (e.g. long-range copy) are structurally out of reach, and the E7 referee is expected to show exactly that. The near window keeps full RoPE’d positional attention.

Visibility is a pure function of (query position p, key position j): j is attended EXACTLY iff p - j < window, and QUANTIZED (position-free) otherwise – regardless of where chunk boundaries fall. Aggregation into cells happens lazily at chunk ends, but tokens past the window horizon that have not folded yet are attended through a “gap bank” of count-1 quantized keys, which by the collapse identity is the same distribution their folded cells produce. Without the gap bank there is a chunk-size-dependent blind spot (evicted but unfolded tokens invisible), which breaks chunking-invariance for every layer past the first – the streaming-equivalence test caught exactly that failure before the gap bank existed. Streaming with any chunk size <= window therefore yields identical integer cell contents after any prefix, and probe losses equal to float-addition-order tolerance.

class CellCountTree(capacity)

Bases: object

Fenwick (binary indexed) tree over stream positions holding sparse per-cell integer counts.

add(pos, cell_id) touches O(log capacity) nodes; prefix(pos) merges O(log capacity) sparse dicts; range_counts(lo, hi) is prefix(hi) - prefix(lo) – the Z^cells group inverse, an exact integer subtraction whose zero entries genuinely cancel (they are removed, not left as float dust). This answers arbitrary-window cell-count queries in O(log n) node merges, where the plain SlidingCellWindow handles only the single fixed trailing window. last_touches receipts the node count of the most recent operation so tests can PROVE the O(log n) claim rather than assert it from theory.

Parameters:

capacity (int)

add(pos, cell_id, count=1)

Record count tokens of cell_id at stream position pos (0-based, < capacity).

Parameters:

• pos (int)

• cell_id (int)

• count (int)

Return type:

None

prefix(pos)

Cell counts over positions [0, pos).

Parameters:

pos (int)

Return type:

dict[int, int]

range_counts(lo, hi)

Cell counts over positions [lo, hi) via the group inverse: prefix(hi) - prefix(lo).

Parameters:

• lo (int)

• hi (int)

Return type:

dict[int, int]

class QuantizedKeyCellState(codes=<factory>, counts=<factory>, vsum=<factory>, cache_k=<factory>, cache_v=<factory>, pos=0, drops=0, lse_receipts=<factory>)

Bases: object

Near window (stop-gradient KV cache, as E1) + far-field cell store, per layer.

The cell store is a fixed-capacity slot table shared across nothing – one table per (batch, head): codes[layer][b, h, c] is slot c’s lattice code vector, counts[layer][b, h, c] its integer token count (0 = free slot), vsum[layer][b, h, c] the float value sum. Fixed capacity means the carried state is bounded regardless of context length – the property the E7 state-byte budget measures. drops counts tokens discarded because every slot was occupied (the honesty receipt: zero on every run whose results you trust for exactness claims).

Parameters:

• codes (list[Any])

• counts (list[Any])

• vsum (list[Any])

• cache_k (list[Any])

• cache_v (list[Any])

• pos (int)

• drops (int)

• lse_receipts (list[QuantizedSoftmaxReceipt | None])

class QuantizedSoftmaxReceipt(bits, span, grid_step, clipped_logits, finite_logits, max_logit_error, max_relative_probability_error_bound)

Bases: object

Verified error scope for one quantized-softmax evaluation.

Parameters:

• bits (int)

• span (float)

• grid_step (float)

• clipped_logits (int)

• finite_logits (int)

• max_logit_error (float)

• max_relative_probability_error_bound (float)

class SlidingCellWindow(window, value_dim)

Bases: object

Sliding-window cell aggregation with EXACT group eviction.

Per-cell integer counts live in Z^cells, a true group under addition: pushing a token adds 1 to its cell, evicting the token that scrolls out subtracts 1 – exact integer arithmetic, so the window’s counts after any number of operations are identical to recomputing from scratch (no drift, ever). Value sums are float64 running sums maintained by the same add/subtract; those are floats, so they carry ~1e-12-scale addition-order drift relative to recomputation – the group-exactness guarantee is about the counts. Cells whose count reaches zero are deleted, so counts always enumerates exactly the occupied cells.

Parameters:

• window (int)

• value_dim (int)

push(cell_id, value)

Add one token; return the (cell_id, value) it evicted, or None if the window isn’t full.

Parameters:

• cell_id (int)

• value (Any)

Return type:

tuple[int, ndarray] | None

totals()

Snapshot of (counts, value_sums) over occupied cells (copies; safe to mutate).

Return type:

tuple[dict[int, int], dict[int, ndarray]]