Cây tóm tắt tìm trong ngữ cảnh thế nào? Cách ZHA hoạt động

function summarize(children: { keys: Vec[]; biases: number[] }[]) {
  const keys: Vec[] = [];
  const biases: number[] = [];
  references.forEach((reference, r) => {
    const scores = children.map((c) => SHARPNESS * dot(c.keys[r]!, reference) + c.biases[r]!);
    const weights = softmax(scores);
    const key = new Float32Array(DIM);
    children.forEach((c, i) => c.keys[r]!.forEach((x, d) => (key[d]! += weights[i]! * x)));
    keys.push(key);
    biases.push(logSumExp(scores) - SHARPNESS * dot(key, reference));
  });
  return { keys, biases };
}

function topP(probabilities: number[], p: number): number[] {
  const order = probabilities.map((_, i) => i).sort((a, b) => probabilities[b]! - probabilities[a]!);
  const picked: number[] = [];
  let total = 0;
  for (const i of order) {
    picked.push(i);
    total += probabilities[i]!;
    if (total >= p) break;
  }
  return picked;
}

// tree.ts: node tree.ts
type Vec = Float32Array;
const DIM = 32;
const CHUNK = 4;
const REFERENCES = 4;
const SHARPNESS = 20;

// Số ngẫu nhiên có seed, để bạn ra đúng các con số trong bài.
let seed = 1;
function random() {
  seed = (seed + 0x6d2b79f5) >>> 0;
  let t = seed;
  t = Math.imul(t ^ (t >>> 15), t | 1);
  t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
  return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
}
const dot = (a: Vec, b: Vec) => a.reduce((sum, x, i) => sum + x * b[i]!, 0);
function unit(v: Vec): Vec {
  const len = Math.sqrt(dot(v, v)) || 1;
  return v.map((x) => x / len);
}
const direction = () =>
  unit(Float32Array.from({ length: DIM }, () => Math.sqrt(-2 * Math.log(1 - random())) * Math.cos(2 * Math.PI * random())));
const nudge = (v: Vec, amount: number) => {
  const noise = direction();
  return unit(v.map((x, i) => x + amount * noise[i]!));
};

// Khoá giống văn bản: các khoá gần nhau nói về cùng một chuyện, và nhóm nằm lồng trong nhóm.
function textKeys(n: number): Vec[] {
  const sizes = [4096, 1024, 256, 64, 16];
  const stack = sizes.map(() => direction());
  const keys: Vec[] = [];
  for (let i = 0; i < n; i++) {
    sizes.forEach((size, level) => {
      if (i % size === 0) stack[level] = nudge(level === 0 ? direction() : stack[level - 1]!, 0.5);
    });
    keys.push(nudge(stack.at(-1)!, 0.35));
  }
  return keys;
}

function softmax(scores: number[]): number[] {
  const max = Math.max(...scores);
  const raw = scores.map((s) => Math.exp(s - max));
  const total = raw.reduce((a, b) => a + b, 0);
  return raw.map((x) => x / total);
}
const logSumExp = (scores: number[]) => {
  const max = Math.max(...scores);
  return max + Math.log(scores.reduce((sum, s) => sum + Math.exp(s - max), 0));
};

// Mỗi nút giữ một tiếp tuyến cho mỗi câu hỏi mẫu: một khoá nén và một bias.
type Node = { keys: Vec[]; biases: number[]; children: Node[]; leaves?: number[] };

const references: Vec[] = Array.from({ length: REFERENCES }, (_, r) =>
  r === 0 ? new Float32Array(DIM) : direction().map((x) => x * 0.5),
);

function summarize(children: { keys: Vec[]; biases: number[] }[]) {
  const keys: Vec[] = [];
  const biases: number[] = [];
  references.forEach((reference, r) => {
    const scores = children.map((c) => SHARPNESS * dot(c.keys[r]!, reference) + c.biases[r]!);
    const weights = softmax(scores);
    const key = new Float32Array(DIM);
    children.forEach((c, i) => c.keys[r]!.forEach((x, d) => (key[d]! += weights[i]! * x)));
    keys.push(key);
    biases.push(logSumExp(scores) - SHARPNESS * dot(key, reference));
  });
  return { keys, biases };
}

function buildTree(allKeys: Vec[]): Node[] {
  const zero = references.map(() => 0);
  let level: Node[] = [];
  for (let i = 0; i + CHUNK <= allKeys.length; i += CHUNK) {
    const leaves = Array.from({ length: CHUNK }, (_, c) => i + c);
    const children = leaves.map((l) => ({ keys: references.map(() => allKeys[l]!), biases: zero }));
    level.push({ ...summarize(children), children: [], leaves });
  }
  while (level.length > 64) {
    const next: Node[] = [];
    for (let i = 0; i + CHUNK <= level.length; i += CHUNK) {
      const children = level.slice(i, i + CHUNK);
      next.push({ ...summarize(children), children });
    }
    level = next;
  }
  return level;
}

// Đi xuống: chấm các ứng viên, mở nhóm nhỏ nhất giữ đủ p phần attention của chúng.
function route(top: Node[], query: Vec, p: number) {
  let candidates = top;
  let work = 0;
  const read: number[] = [];
  while (candidates.length > 0) {
    work += candidates.length * REFERENCES;
    const scores = candidates.map((node) =>
      Math.max(...node.keys.map((key, r) => SHARPNESS * dot(key, query) + node.biases[r]!)),
    );
    const mass = softmax(scores);
    const order = candidates.map((_, i) => i).sort((a, b) => mass[b]! - mass[a]!);
    const opened: Node[] = [];
    let total = 0;
    for (const i of order) {
      opened.push(candidates[i]!);
      total += mass[i]!;
      if (total >= p) break;
    }
    for (const node of opened) if (node.leaves) read.push(...node.leaves);
    candidates = opened.flatMap((node) => node.children);
  }
  return { work: work + read.length, read };
}

for (const n of [1024, 4096, 16384, 65536]) {
  const keys = textKeys(n);
  const top = buildTree(keys);
  let work = 0;
  let kept = 0;
  const queries = 12;
  for (let q = 0; q < queries; q++) {
    const query = nudge(keys[Math.floor(random() * n)]!, 0.2);
    const walk = route(top, query, 0.9);
    const truth = softmax(keys.map((key) => SHARPNESS * dot(key, query)));
    work += walk.work;
    kept += walk.read.reduce((sum, i) => sum + truth[i]!, 0);
  }
  console.log(
    `${n} keys: ${Math.round(work / queries)} comparisons per query instead of ${n}, ` +
      `keeps ${Math.round((kept / queries) * 100)}% of the attention`,
  );
}

1024 keys: 365 comparisons per query instead of 1024, keeps 85% of the attention
4096 keys: 411 comparisons per query instead of 4096, keeps 83% of the attention
16384 keys: 422 comparisons per query instead of 16384, keeps 81% of the attention
65536 keys: 572 comparisons per query instead of 65536, keeps 78% of the attention