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 FrequentDirectionsState(B: 'list[Any]' = <factory>, Z: 'list[Any]' = <factory>, cache_k: 'list[Any]' = <factory>, cache_v: 'list[Any]' = <factory>, pos: 'int' = 0, receipt: 'dict[str, Any]'=<factory>)[source]

Bases: object

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

Bases: object

Parameters: