d_mem, h_mem, num_pairs = 32, 128, 16
ks = jax.random.split(jax.random.key(3), 4)
k_pairs = jax.random.normal(ks[0], (num_pairs, d_mem))
k_pairs /= jnp.linalg.norm(k_pairs, axis=1, keepdims=True)
v_pairs = jax.random.normal(ks[1], (num_pairs, d_mem))
v_pairs /= jnp.linalg.norm(v_pairs, axis=1, keepdims=True)
order = jax.random.permutation(ks[2], jnp.tile(jnp.arange(num_pairs), 4))
params = {'W1': jax.random.normal(ks[3], (d_mem, h_mem)), # Unit-var hidden
'b1': jnp.zeros(h_mem),
'W2': jnp.zeros((h_mem, d_mem)), # Start empty: retrieve 0
'b2': jnp.zeros(d_mem)}
def memory(params, k):
hidden = jax.nn.gelu(k @ params['W1'] + params['b1'])
return hidden @ params['W2'] + params['b2']
def recall_loss(params, k, v):
return ((memory(params, k) - v)**2).sum()
grad_fn = jax.grad(recall_loss)
def retrieval_mse(params):
return float(((memory(params, k_pairs) - v_pairs)**2).mean())
before = retrieval_mse(params)
theta, eta = 0.005, 0.5
def write(carry, t):
params, velocity = carry
g = grad_fn(params, k_pairs[t], v_pairs[t]) # Surprise
velocity = jax.tree.map(lambda u, gg: eta * u - theta * gg, velocity, g)
params = jax.tree.map(lambda p, u: p + u, params, velocity)
return (params, velocity), None
velocity = jax.tree.map(jnp.zeros_like, params)
(params, _), _ = jax.lax.scan(write, (params, velocity), order)
after = retrieval_mse(params)
print(f'retrieval MSE: empty memory {before:.4f} -> after the stream '
f'{after:.4f}')
assert after < before / 5What learning sharpened
- The region of large attention weights becomes sharper once the bandwidth is learned; the learned \sigma lands where the hand sweep pointed.
- A regression estimator, tuned by gradient descent on the data it serves. Hold that shape.
Two loops: what learns when
Left — pretraining updates the outer parameters through every inner update. Right — at inference only the inner loop runs.
- Inner loop: the state \mathbf{S}_t — fresh per sequence, adapts token by token, runs in every forward pass. No labels at inference: the targets are the sequence’s own values.
- Outer loop: projections, gate networks, \mathbf{S}_0 — get gradients through the chain of inner updates; frozen at inference.
- Applies verbatim to DeltaNet, Longhorn, Titans.
Three characters
Kernel smoother: faithful, ragged. One-pass delta: smooth, shallow. Batch solve: threads the noise. Every layer in the table is a point on this picture — computed once per token, forward-only.
Regression that tracks
mean error over the last 500 steps: {'uniform LS': 0.99, 'decayed LS': 0.29, 'implicit step': 0.3}- Drifting truth \mathbf{w}^*_t: uniform LS goes stale — exact solver, right class, wrong weights.
- Decayed LS tracks; the implicit step tracks with no solve at all. Decay = a belief about how fast the world changes.
- Discount fixed by hand here (\gamma = 0.95); trained models emit it per token. Forgetting buys tracking under drift, costs efficiency on stationary stretches.