因果注意力与 KV cache 一致性

先修与预测

最后一个 value 从 7 改为 999，哪些更早输出可以变化？

完整实现与测试入口

将下面四个文件下载到同一目录，执行 node check.mjs。检查覆盖三个项目，无需安装依赖或 API Key。

// 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.');

迁移任务与消融

给完整分数矩阵增加 padding 掩码，与增量前缀实现比较，并测试全掩码行。

评分标准

解释 Q/K 维度与 √d 缩放；实现无指数溢出；通过前缀/缓存不变性；解释权重为何不是因果归因。

四项分别按 0（缺失）、1（需帮助）、2（有独立证据）评分。每项均为 2 才过关，自评不是认证。