mixle.experimental.sketch_state_attention module

E3: sketch-state attention – oblivious (data-independent update rule) far-field states with a provable approximation guarantee, contrasted with E2’s adaptive/learned far-field state. See notes/designs/E3.md for the full design (citations, the augmented-row cross-covariance trick, the tensor-sketch FFT derivation, and the honestly-flagged tension between FD’s SVD shrink and the card’s “All linear => exact gradients” framing).

Three mechanisms, all implementing ContextMechanism, differing only in how the stream of per-token key/value pairs (phi(k_t), v_t) is compressed into carried state:

  • (a) `LinearAttentionSpine` – exact, unbounded-rank running sum S = sum phi(k_t) v_t^T, Z = sum phi(k_t) (Katharopoulos et al. 2020 kernel trick, phi = elu(x) + 1). No local window: the whole stream is the linear-attention prefix, chunked as a running cumulative sum (bit-identical to a single non-chunked pass, since carrying S/Z across chunk boundaries IS the prefix sum’s carry). This is the fixed-byte-size reference point (b)/(c) approximate.

  • (b) `FrequentDirectionsSpine` – a small exact local window (SlidingWindowSpine-style stop-gradient cache) plus a Frequent Directions sketch (Liberty, KDD 2013) of the augmented rows [phi(k_t) ; v_t] for every token once it scrolls out of the local window. B is literally ell x (d_phi + d_v) with genuine zero rows between shrinks (Liberty’s Algorithm 1, not a rank-compacted variant) – this is what makes the deterministic Theorem 1.1 bound test meaningful. The normalizer Z = sum phi(k_t) is tracked exactly alongside the sketch (cheap, O(d_phi) per step; the Proposed API’s illustrative dataclass didn’t spell this field out, but the design note’s own Algorithm section requires it – there is no valid FD readout without it).

  • (c) `TensorSketchSpine` – same local-window split, but the far-field accumulator is a Count-Sketch + FFT-circular-convolution tensor sketch (Pham & Pagh, KDD 2013) of phi(k_t), capturing degree-p polynomial-kernel interactions FD’s/(a)’s linear rows cannot represent, at the cost of an in-expectation (not worst-case) guarantee.

Local window vs pure prefix. (a)’s own Algorithm section describes no local softmax component at all – it degenerates the WHOLE stream to a linear-attention kernel that carries exact state, matching its constructor (no window parameter). (b)/(c) each keep a small local exact-softmax window (their constructors take window=64) plus the sketch as an additive far-field term – “local half unchanged from SlidingWindowSpine, far-field half is what varies” (design note’s Proposed API section, and the “Do NOT fold the near-field/far-field split into one undifferentiated block” rule).

class LinearAttentionState(S: 'list[Any]' = <factory>, Z: 'list[Any]' = <factory>, pos: 'int' = 0)[source]

Bases: object

Parameters:
class LinearAttentionSpine(vocab, *, d_model=32, n_layer=2, n_head=2)[source]

Bases: Module

  1. Exact linear-attention prefix state (Katharopoulos et al. 2020), chunked-scan trained.

No local window: the whole stream is the linear-attention kernel (see module docstring for why – this mechanism’s own Algorithm section in notes/designs/E3.md has no local softmax term at all, unlike (b)/(c)). RoPE is applied to the raw q/k projections before the phi = elu + 1 feature map, so positional information survives into the kernel while S/Z stay simple running sums – carrying them across chunk boundaries reproduces the exact same cumulative sum a single non-chunked pass over the whole prefix would compute (the “chunked scan” streaming-equivalence invariant the test suite checks directly).

Parameters:
class FrequentDirectionsState(B: 'list[Any]' = <factory>, Z: 'list[Any]' = <factory>, cache_k: 'list[Any]' = <factory>, cache_v: 'list[Any]' = <factory>, pos: 'int' = 0)[source]

Bases: object

Parameters:
class FrequentDirectionsSpine(vocab, *, d_model=32, n_layer=2, n_head=2, window=64, ell=16)[source]

Bases: Module

  1. FD sketch of the KV outer-product stream, exact per Liberty (2013).

Local half: a small SlidingWindowSpine-shaped exact-softmax window. Far-field half: once a token scrolls out of the window, its augmented row [phi(k_t) ; v_t] is streamed into a Frequent Directions sketch B (literal ell x (d_phi + d_v) shape with genuine zero rows – see frequent_directions_update/_fd_insert_row), and the normalizer Z = sum phi(k_t) is tracked exactly alongside it. A query reads the far field back as phi(q)^T (B_K^T B_V) / (phi(q)^T Z), B_K/B_V being B’s two column blocks split at d_phi – an FD-bounded approximation of the exact cross term (a) tracks exactly.

Parameters:
class TensorSketchState(C: 'list[Any]' = <factory>, cache_k: 'list[Any]' = <factory>, cache_v: 'list[Any]' = <factory>, pos: 'int' = 0)[source]

Bases: object

Parameters:
class TensorSketchSpine(vocab, *, d_model=32, n_layer=2, n_head=2, window=64, sketch_dim=64, degree=2, seed=0)[source]

Bases: Module

(c) Tensor sketch (Count Sketch + circular convolution) of degree-p key features (Pham & Pagh 2013). Local half identical in shape to (b); far-field half accumulates C_t = C_{t-1} + TS(phi(k_t)) v_t^T for evicted tokens and reads it back as TS(phi(q))^T C (no normalizer – the design note’s Algorithm section for (c) doesn’t specify one, unlike (a)/(b); this mirrors that exactly rather than inventing an extra division).

Parameters:
frequent_directions_update(B, rows, ell)[source]

One FD ingest-and-shrink pass (Liberty 2013, Algorithm 1). B: (ell, d), rows: (m, d) new rows, inserted one at a time (insert into a zero row; shrink whenever none remains). Returns the updated (ell, d) B – unbatched, matching the design note’s Proposed API signature exactly (the batched spine-internal use reuses the same _fd_insert_row primitive).

Parameters:
Return type:

Any

frequent_directions_error_bound(A, B, ell, k)[source]

RHS of Liberty’s Theorem 1.1: ||A - A_k||_F^2 / (ell - k) – A_k is A’s best rank-k approximation (Eckart-Young). Depends only on A, ell, k (not on B – the theorem’s guarantee is that ANY B produced by streaming A’s rows through FD satisfies ||A^T A - B^T B||_2 <= `` this quantity); ``B is accepted to match the design note’s Proposed API signature and to allow a caller to sanity-check B.shape[0] == ell.

Parameters:
Return type:

float

tensor_sketch_project(x, hashes, signs, sketch_dim)[source]

Degree-len(hashes) tensor sketch of x (last dim is the feature dim) via count-sketch + FFT circular convolution (Pham & Pagh 2013): TS(x) = IFFT(prod_i FFT(CS_i(x))). The defining property this implements: TS(x)^T TS(y) is an unbiased estimator of (x^T y)^p for p = len(hashes), with variance O(1 / sketch_dim).

Parameters:
Return type:

Any

make_tensor_sketch_hashes(d, *, sketch_dim, degree, seed, device='cpu')[source]

degree independent (hash, sign) pairs, fixed at construction (the “oblivious” part – the hash/sign choice does not depend on the data). hash_i: [d] -> [sketch_dim], sign_i: [d] -> {-1, +1}.

Parameters:
Return type:

tuple[list[Any], list[Any]]

fd_misfit_receipt(A, ell, *, k=0)[source]

Stream A’s rows through FD, then report the realized ||A^T A - B^T B||_2 against Liberty’s Theorem 1.1 bound – “how tight is the guarantee in practice”, the (b) misfit receipt.

Parameters:
Return type:

dict[str, float]

tensor_sketch_misfit_receipt(*, d, sketch_dim, degree, seed=0, trials=200)[source]

Empirical collision/variance rate of the tensor sketch inner-product estimator: sample TS(x)^T TS(y) over many fresh random (x, y) pairs (same hash/sign, per the “oblivious” contract) and report the empirical bias and variance against the true (x^T y)^p – the (c) misfit receipt (graduation.py’s “sketch collision rate” worked example).

Parameters:
Return type:

dict[str, float]