Causal attention & KV-cache correctness

Prerequisite & prediction

If the final value changes from 7 to 999, which earlier outputs may change?

Complete implementation & test runner

Download all four files into one folder. Run node check.mjs. The checks exercise all three projects; no installation or API key is required.

// Node.js 18+ or modern browser; no dependencies, no random seed needed.
export function softmax(xs) {
  const m = Math.max(...xs);
  if (!Number.isFinite(m)) throw Error('All-masked or invalid attention row');
  const exp = xs.map(x => Math.exp(x - m)), sum = exp.reduce((a,b)=>a+b,0);
  return exp.map(x=>x/sum);
}
export function attend(q, keys, values) {
  if (!q.length || keys.length !== values.length || !keys.length) throw Error('Invalid shape');
  if (!keys.every(k=>k.length===q.length) || !values.every(v=>v.length===values[0].length)) throw Error('Shape mismatch');
  if (![q,...keys,...values].flat().every(Number.isFinite)) throw Error('Non-finite input');
  const weights = softmax(keys.map(k=>k.reduce((s,x,j)=>s+x*q[j],0)/Math.sqrt(q.length)));
  return values[0].map((_,j)=>weights.reduce((s,w,i)=>s+w*values[i][j],0));
}
function validateSequence(Q,K,V){
 if(![Q,K,V].every(Array.isArray)||!Q.length||Q.length!==K.length||K.length!==V.length)throw Error('Invalid sequence shape');
 if(!Q[0]?.length||!V[0]?.length||!Q.every(q=>Array.isArray(q)&&q.length===Q[0].length)||!K.every(k=>Array.isArray(k)&&k.length===Q[0].length)||!V.every(v=>Array.isArray(v)&&v.length===V[0].length))throw Error('Shape mismatch');
 if(![...Q,...K,...V].flat().every(Number.isFinite))throw Error('Non-finite input');
}
export function full(Q,K,V) {
  validateSequence(Q,K,V);
  if (Q.length!==K.length || K.length!==V.length || !Q.length) throw Error('Invalid sequence shape');
  return Q.map((q,t)=>{
    const scores=K.map((k,j)=>j>t?-Infinity:k.reduce((s,x,i)=>s+x*q[i],0)/Math.sqrt(q.length));
    const weights=softmax(scores);
    return V[0].map((_,i)=>weights.reduce((sum,w,j)=>sum+w*V[j][i],0));
  });
}
export function cached(Q,K,V) {
  validateSequence(Q,K,V);
  const keys=[], values=[];
  return Q.map((q,t)=>{keys.push(K[t]);values.push(V[t]);return attend(q,keys,values)});
}
export function run() {
  const Q=[[1,0],[0,1],[1,1],[.5,-1]], K=[[1,0],[0,1],[1,1],[-1,1]], V=[[1],[3],[5],[7]];
  const baseline=full(Q,K,V), incremental=cached(Q,K,V);
  const changedK=K.map(x=>[...x]), changedV=V.map(x=>[...x]);changedK[3]=[100,-100];changedV[3]=[999];
  const perturbed=full(Q,changedK,changedV);
  const cacheError=Math.max(...baseline.flat().map((x,i)=>Math.abs(x-incremental.flat()[i])));
  const causalError=Math.max(...baseline.slice(0,3).flat().map((x,i)=>Math.abs(x-perturbed.slice(0,3).flat()[i])));
  const stable=softmax([1000,1001,999]);
  return {baseline,incremental,cacheError,causalError,stableSoftmaxSum:stable.reduce((a,b)=>a+b,0),
    passes:cacheError<1e-12&&causalError<1e-12&&stable.every(Number.isFinite)};
}

import assert from 'node:assert/strict';
import {run as attention,softmax,attend,full,cached} from './attention.mjs';
import {run as retrieval,retrieve} from './retrieval.mjs';
import {run as control,simulate} from './control.mjs';
assert.ok(attention().passes);assert.throws(()=>softmax([-Infinity,-Infinity]));assert.throws(()=>attend([1,2],[[1]],[[1]]));
assert.equal(retrieval().baseline.accuracy,.4);assert.equal(retrieval().treatment.accuracy,1);assert.equal(retrieve('unseen vocabulary'),null);
const result=control();assert.ok(result.baseline.maxForce<=5);assert.ok(result.delayed.maxForce<=5);assert.equal(simulate({steps:1}).trace[0].x,0);assert.equal(simulate({steps:1}).trace[0].v,.05);assert.throws(()=>simulate({mass:0}));

// Full and cached paths must reject inconsistent shapes, not just the single-query helper.
assert.throws(()=>full([[1,2]],[[1]],[[1]]),/Shape mismatch/);
assert.throws(()=>cached([[1,2]],[],[]),/sequence shape/);

console.log('PASS: 12 capstone checks. Toy arithmetic and invariants only; no GPU or hardware benchmark.');

Transfer task & ablation

Add padding masks to the full score matrix. Compare outputs with the incremental prefix implementation, including all-masked rows.

Rubric

Explain Q/K dimensions and √d scaling; implement without exponent overflow; pass prefix/cache invariance; explain why weights are not causal attribution.

Score each of the four criteria 0 (missing), 1 (with help), 2 (independent evidence). Passing requires 2 on every criterion. A self-check is not a certification.