How does a tree of summaries search a context? How ZHA works

You find a passage in a thick book by checking the contents, choosing a chapter, choosing a section, and only then reading. You do not read each sentence to learn which sentence to read. Zyphra Hierarchical Attention (ZHA), published in November 2026, applies exactly that idea to the KV cache: build a table of contents with several levels, then walk from the top.

Try it: a question walks down the tree#

Each row is one level of the tree. A dark cell is a node that was opened, a pale cell is one scored and skipped, and empty space is nodes that were never touched. The last strip is the true attention, so you can see whether the tree went to the right place.

Try this order:

1. Keep Ask about one place with p at 90%. The levels score then open 64 → 4, 16 → 9, 36 → 25. In the end the tree reads only 100 keys and spends 564 comparisons instead of 4,096, keeping about 73% of the attention.
2. Pick Ask about eight places. Only the question changed, yet the tree opens wider by itself: it reads 1,084 keys and spends 3,500 comparisons.
3. Switch to Top-K with 6 branches. All three questions read exactly 24 keys. That sounds thrifty, but the eight-place question keeps only about 10% of the attention.
4. Go back to top-p, pick Random, no structure and Ask about one place. The tree spends 2,888 comparisons and keeps only 0.2% of the attention. The summaries no longer say anything.

How is the tree built?#

Each group of neighbouring keys is summarized into a node, and those nodes are summarized again above them, as the text grows. With chunks of 4, one level summarizes 4 keys, two levels 16 keys, and six levels 4,096 keys.

Two things matter here:

• The original keys are not deleted. The tree is only for searching. When the search ends, ordinary attention runs on exactly the original keys it chose.
• It is very cheap to keep up. While writing, a node is created each time a chunk fills. The paper puts this at around 1% of the time or less in its performance model.

What does a node summarize?#

A node does not summarize content. It estimates how much attention a question would give this region. Averaging the keys would lose information, because different questions care about different aspects of the same region.

ZHA uses a few reference queries, sample questions learned in advance. For each one, a node stores a compressed key and a correction number called a bias:

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 };
}

When a real question arrives, the node's score is dot product + bias, taking the highest across the reference queries. In maths terms, each reference query gives one tangent of the function "total attention of this region". That function is convex, so a tangent matches exactly at its reference and sits below the true value everywhere else.

We checked these properties with tests on the code the demo runs: the estimate is exact at a reference, never exceeds the true value, and merging child nodes gives the same result as summarizing the original keys directly. That last one is also what the paper states in a footnote to its architecture section.

How does top-p differ from top-K?#

Top-p opens the smallest group holding enough of the estimated attention. Top-K always opens exactly K nodes.

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;
}

The difference matters because real tasks vary a lot. The paper measures this at about 64K tokens: 90% of attention sits in a median of 2.3% of 16-key chunks on RULER, 11.7% on SWE-bench Verified and 15.0% on long fiction. A K large enough for fiction wastes about six times as much on RULER.

Does the model have to be retrained?#

No. The original model is frozen, and only the reference queries are learned, by comparing how the tree splits attention against how the model does (KL distillation). For Qwen3.6-35B-A3B the addition is 81,920 parameters, about 0.00024% of the model. They report needing only 1 billion training tokens at a length of 32K.

Note that the paper also says small models under 1 billion parameters often did poorly on retrieval when converted to ZHA, because their original attention scores were already noisy. Results improved clearly from around 4 billion parameters.

Run a small tree yourself#

Below is a complete tree with no outside libraries, using chunks of 4 and 4 reference queries. Two things are simplified: the keys are made-up vectors with a nested structure, and the reference queries are fixed rather than learned. It builds up to 64 nodes at the top level and measures for real.

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

// Seeded random numbers, so you get the numbers in this lesson.
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]!));
};

// Text-like keys: neighbours are about the same thing, and groups nest inside groups.
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));
};

// A node holds one tangent per reference query: a compressed key and a 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;
}

// Walk down: score the candidates, open the smallest group holding p of their mass.
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`,
  );
}

On 3 November 2026, on Node 26.11.1, the result was:

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

A context 64 times longer costs less than 2 times more per question. But the last column is falling too. The next lesson explains why.

When can a tree not help?#

When the data has no structure. The paper says plainly that sequences of random tokens, where any token may have to be retrieved, would need full quadratic attention. The compressibility does not come from the architecture, it comes from the data.

The evidence is in their measurements: on the Qwen model, the top-p set grows as a sublinear power law of the text's length. So as a text gets longer, the part that has to be read grows, but more slowly than the text. On Nemotron, which uses no positional encoding (NoPE), that law is less clear, and retrieval scores also dropped further after conversion to ZHA.

Is anyone else going this way?#

Yes, and at almost the same time. Latent Sparse Attention (LSA) groups keys into hierarchical clusters in a learned space, then scores clusters rather than each key. It has two variants: D-LSA distils into an existing model, and N-LSA trains natively. HiLS learns chunk selection from the language model's own loss.

We only read the abstracts of these two papers and did not compare them with ZHA in detail, so we draw no conclusion about which is better. Their numbers are also measured on different models.

Common misunderstandings#

• "ZHA summarizes, then reads the summary." No. The summaries only find the way. The final attention runs on the original keys and values.
• "It follows one branch, like an ordinary search tree." One question can open many branches at once, far apart in the context.
• "One walk down the tree serves the whole model." Each attention head in each layer walks separately for each query.
• "p of 90% means keeping 90% of the attention." That is the threshold at each level, and the levels multiply together. That is the subject of the last lesson.