| Safe Haskell | None |
|---|---|
| Language | GHC2024 |
Circuit.LLM.Mixer
Description
Stateful token mixers from the pedagogic GPT-2 → linear → delta ladder (waterloo_intern worklog / Kimi lineage).
Already in this package: full-sequence softmax MHA (GPT,
Attention). This module adds the memory progression:
- softmax decode with an explicit KV cache
- linear attention with fixed-size state
S(and normalizerz) - delta-rule write (overwrite what the key already stores)
- gated decay of
S(Mamba-style scalar, or per-channel vector for KDA-lite) - chunked additive linear scan (matches sequential; training-shaped)
Shapes (single head, teaching form):
q, k, v :: Matrix -- T×d (rows = time, cols = head dim) S :: Matrix -- d×d associative memory
Multi-head is a map over heads later; keep one head honest first.
Pedagogic imprint: same blocks as the article, hmatrix BLAS, pure scans — the discrete Process / decaying Stats story, but for attention state.
Synopsis
- data KvCache = KvCache {}
- emptyKv :: Int -> KvCache
- softmaxAttnPrefill :: Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
- softmaxAttnStep :: Matrix Double -> Matrix Double -> Matrix Double -> KvCache -> (Matrix Double, KvCache)
- data LinearState = LinearState {}
- emptyLinear :: Int -> LinearState
- eluPlus1 :: Matrix Double -> Matrix Double
- linearAttnPrefill :: Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, LinearState)
- linearAttnStep :: Matrix Double -> Matrix Double -> Matrix Double -> LinearState -> (Matrix Double, LinearState)
- deltaAttnStep :: Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double)
- deltaAttnPrefill :: Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double)
- gatedDeltaStep :: Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double)
- gatedDeltaPrefill :: Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double)
- kdaLiteStep :: Vector Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double)
- chunkedLinearAttn :: Int -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double
- scanMixer :: (Matrix Double -> Matrix Double -> Matrix Double -> s -> (Matrix Double, s)) -> s -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, s)
Softmax + KV cache
Growing key/value cache for decoder-style softmax attention.
softmaxAttnPrefill :: Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double Source #
Scaled dot-product attention (full sequence), causal.
softmaxAttnStep :: Matrix Double -> Matrix Double -> Matrix Double -> KvCache -> (Matrix Double, KvCache) Source #
One decode step: q,k,v are 1×d (current token). Returns output 1×d and
extended cache.
Linear attention (fixed state)
data LinearState Source #
Fixed-size linear-attention memory: S accumulates φ(k)ᵀ v outer
products; z accumulates φ(k) for the normalizer.
Instances
| Eq LinearState Source # | |
Defined in Circuit.LLM.Mixer | |
| Show LinearState Source # | |
Defined in Circuit.LLM.Mixer Methods showsPrec :: Int -> LinearState -> ShowS # show :: LinearState -> String # showList :: [LinearState] -> ShowS # | |
emptyLinear :: Int -> LinearState Source #
eluPlus1 :: Matrix Double -> Matrix Double Source #
Feature map φ(x) = ELU(x)+1 (article / Katharopoulos linear attention).
linearAttnPrefill :: Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, LinearState) Source #
Prefill by scanning tokens (rows of q,k,v).
linearAttnStep :: Matrix Double -> Matrix Double -> Matrix Double -> LinearState -> (Matrix Double, LinearState) Source #
One token of linear attention. q,k,v are 1×d.
S ← S + φ(k)ᵀ v z ← z + φ(k) o = (φ(q) S) / (φ(q)·z)
Delta rule
deltaAttnStep :: Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double) Source #
One delta-rule step (sequential form from the article).
v_old = k S -- what this key already retrieves (k is 1×d, S is d×d:
-- we use S as mapping d→d on row vectors: o = q <> S)
u = β (v - v_old)
S' = S + kᵀ u
o = q S'
Convention: treat S as right-acting on 1×d queries (o = q <> S), so
writes are S += kᵀ <> u with k,u as 1×d.
deltaAttnPrefill :: Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double) Source #
Prefill with constant β.
Gated memory
gatedDeltaStep :: Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double) Source #
Gated delta step: decay previous state, then delta-write.
S ← α S then deltaAttnStep β
α = 1 → pure delta; α = 0 → forget all then write.
gatedDeltaPrefill :: Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double) Source #
kdaLiteStep :: Vector Double -> Double -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double) Source #
KDA-lite: per-channel decay α :: Vector d (elementwise scale of each
column of S before the delta write). Full KDA has richer parameterisation;
this imprints the fine-grained forget idea.
Chunked additive linear (matches sequential)
chunkedLinearAttn :: Int -> Matrix Double -> Matrix Double -> Matrix Double -> Matrix Double Source #
Chunked additive linear attention (article intermediate form).
Within chunk: causal q (kᵀ v)-style via masked scores on φ features;
across chunks: fold into S and read q S.
For feature map = identity and no normalizer this is didactic; we use
φ = ELU+1 and still accumulate S as in sequential linear (without z
in the chunk path for simplicity of the equality oracle — compare
unnormalized numerators).