Charts

Attention Heads

A small-multiples grid of causal attention maps — one miniature per (layer, head), sortable by row entropy or previous-token mass, with the picked head enlarged beside it under token labels on both axes and its mass split into four shares that add to exactly 100.

Preview in your theme

Loading preview…

"use client"

import * as React from "react"

import { cn } from "@/lib/utils"
import type {
  ChartAttentionHeadsData,
  ChartAttentionHeadsMap,
  ChartAttentionHeadsSelection,
} from "./chart-attention-heads.contract"

/** What the head grid is ordered by. `index` is the checkpoint's own numbering. */
export type AttentionSortKey = "index" | "entropy" | "previous"

Installation

npx shadcn@latest add https://ui.zyeon.ai/r/chart-attention-heads.json

Prompt

Build a React + TypeScript + Tailwind "ChartAttentionHeads" head-catalogue card
in plain SVG with zod. Not recharts and not a generic heatmap component: the
mark is a triangular softmax row, the grid is one miniature per (layer, head),
and every number on the card is derived from the weights rather than sent.

Contract
- One zod schema is the source of truth:
  { status: "loading" | "empty" | "error" | "ready";
    title?: string; caption?: string;
    model: { name: string; layers: int > 0; headsPerLayer: int > 0;
             contextTokens?: int > 0 };
    tokens: string[];      // the prompt, tokenizer strings, index IS position
    maps: { layer: int >= 0; head: int >= 0;
            rows?: number[][];              // dense causal: rows[q] covers 0..q
            cells?: { q: int; k: int; w: number }[] }[];   // sparse
    selected?: { layer, head } | null }.
- Exactly one of rows / cells per head — superRefine rejects both or neither,
  because merging them needs a rule nobody agreed on. A padded square IS
  accepted as rows, as long as everything above the diagonal is zero (that is
  what a masked tensor dumps); a padded square with mass up there is rejected
  as transposed or unmasked, and any cell with k > q likewise.
- Repeated (q, k) pairs are SUMMED so a payload assembled from shards can be
  concatenated; an absent cell is a measured zero, not a gap, because a softmax
  row is a partition of one unit and nothing can be missing from it.
- model.layers x model.headsPerLayer is the DENOMINATOR ("36 of 1,920 heads"),
  not the number plotted: nobody renders two thousand miniatures, so the header
  states the slice against the whole.
- Props = z.infer plus defaultSort ("index" | "entropy" | "previous"), onSelect,
  onRetry, className and the div's native props; forwardRef to the card.
- Export the maths beside the component so a test can print the numbers the
  picture is made of: attentionAlpha(), apportion(), rowEntropy(),
  inductionTargets(), archetypeOf(), buildAttentionHead(), buildAttentionGrid(),
  orderHeads(), staircasePath().

Behavior
- NORMALISE, THEN MEASURE. Fold both payload forms into dense causal rows,
  counting every cell that could not be placed. Renormalise any row carrying
  mass — a sparse payload legitimately drops its tail, and a row summing to 0.97
  did not attend 3% to nothing — and count the rows that were more than 2% off
  before rescaling, then say so in view. A head with no mass at all is dropped
  from the grid and reported, never drawn as an empty square.
- ROW 0 IS EXCLUDED FROM EVERY MEAN. A first token can only attend to itself, so
  its entropy is 0 by construction; averaging it in would score every head lower
  purely as a function of how short the prompt is.
- PER-HEAD MATHS, all over query rows 1 and up that carry mass:
    entropy   = mean of -sum(p ln p), in nats;
    uniform   = mean of ln(q + 1), what a uniform row would have scored;
    span      = exp(entropy), the tokens the average row spreads over.
- FOUR DISJOINT BUCKETS, one stated precedence: self at q, then previous at
  q - 1, then the first token, then anything earlier. The row-1 cell that is
  BOTH the first token and the previous one is therefore counted once, as
  previous — the most local relationship wins, and the rule is printed on the
  card. Because every scored row sums to 1 and every cell lands in exactly one
  bucket, the four shares partition the head's whole mass; apportion them by
  LARGEST REMAINDER so they add to exactly 100, and let the bar, the legend list
  and the "previous token" tile all read those same four integers. Four numbers
  describing one softmax row that print 101 read as a data error.
- INDUCTION SCORE, measured separately and allowed to overlap the buckets: the
  mass on the token that FOLLOWED an earlier copy of the query token, averaged
  over the rows that have such a copy, with the row count printed as the
  denominator. Null when no token repeats, and the card says so rather than
  printing a zero it did not measure.
- ARCHETYPE by the first rule that fires, thresholds printed: induction, then
  previous, then first token, then self at 45%; otherwise diffuse when the
  effective span reaches 85% of a uniform row's; otherwise mixed. Induction is
  tested FIRST on purpose — an induction head sinks onto the first token on
  every row that has nothing to match, so testing "sink" first would hide the
  pattern the card exists to surface.
- SORT INSIDE A LAYER, NEVER ACROSS. Rows are layers, ascending, always: depth
  is structural, and re-ranking the rows would destroy the only axis that
  carries meaning on its own. Within a row the reader picks head index (the
  checkpoint's own order), entropy ascending (sharpest first) or previous-token
  mass descending — so the left column becomes "the most previous-token-ish head
  in each layer" and the induction circuit's first hop surfaces by itself. Every
  miniature carries its own H-number underneath, so a sorted grid is still
  navigable.
- SELECTION, NOT HOVER. The grid is a single-select role="listbox" of
  role="option" miniatures with aria-activedescendant and one tab stop;
  selection follows focus, arrow keys walk 2-D (left/right inside a layer,
  up/down between layers) plus Home/End/PageUp/PageDown, and a delegated click
  handler resolves the option via closest("[data-head]"). Hover never moves the
  enlarged map: a detail panel that flickers under the pointer cannot be read.
- CONTROLLED WITHOUT TWO SOURCES OF TRUTH: the payload's `selected` seeds the
  pick, and when it CHANGES the component drops the reader's pick and follows it
  again (adjust-state-during-render, so there is no effect and nothing to clean
  up).
- Four first-class branches of one card: loading (grid-of-squares skeleton in
  the ready silhouette, aria-hidden, plus an sr-only role=status line), empty
  (the causal staircase, waiting, worded so it cannot be mistaken for a failed
  fetch), error (role=alert, "Try again" only when onRetry was passed), ready.
  status "ready" with nothing drawable, or fewer than two tokens, renders empty —
  and the empty card names WHICH of the two it was: a prompt too short to label
  the axes is not "maps that could not be placed", so the lost-value count is
  suppressed there rather than blaming the maps for a missing prompt.
  data-status reports that derived branch, never the raw prop.
- CLEANUP: there is nothing to clean up, and that is the design — no timers, no
  rAF, no ResizeObserver, no window listeners. Layout is CSS grid auto-fit and
  the SVGs scale by viewBox, so the card reflows without measuring.

Rendering & styling
- Semantic tokens only: bg-card / text-card-foreground panel, border, bg-muted
  for the causal canvas and the skeleton, text-muted-foreground for axis and
  small print, ring-ring plus bg-muted for the selected option, stroke-foreground
  at 60% for the row-maximum outline, bg-primary / text-primary-foreground for
  the archetype badge, destructive for the error heading. Weight uses ONE chart
  token, var(--chart-1), whose OPACITY climbs; the four buckets use
  var(--chart-2..5) so colour never means two things at once. Never a hex.
- THE RAMP IS A FUNCTION, NOT A TABLE: alpha = 0.22 + 0.78 * min(w,1)^0.7, and
  nothing under w = 0.02 paints at all (it reads as bare canvas — on an
  eleven-token prompt a uniform row is 9% a cell, so a fifth of that is softmax
  floor, not a link). The miniatures, the enlarged map and the legend swatches
  all call that one function AND paint it over the same bg-muted canvas, so a
  legend can never describe a ramp the map does not use — a cell composited over
  the card in one place and over the canvas in another is one weight with two
  colours. The 0.22 floor is deliberate: a tenth of a chart token on a
  near-black canvas lands within a few sRGB steps of the canvas itself.
- Each miniature is an svg with viewBox "0 0 n n" and shapeRendering=crispEdges:
  one path for the lower-triangular canvas (a staircase, cheaper than a
  background rect per cell), then one 1x1 rect per above-floor weight at exactly
  (k, q) — the same pair that sets its opacity and names it in the table.
- The enlarged map draws every causal cell at 20 user units, query tokens
  right-aligned down the left gutter with their index, key tokens rotated -90
  above their column over a row of key indices, and outlines each row's
  strongest key at the index the maximum was found at.
- ACCESSIBILITY: the listbox is NAMED by the card's heading (aria-labelledby)
  and DESCRIBED by the model line plus the sr-only summary paragraph — a name is
  re-announced on every arrow press, so the keyboard prose belongs in the
  description, stated once, not duplicated as name and sr-only text. Each
  option's aria-label is the same sentence its heading and its table row print;
  the enlarged map is role="img" (children presentational) with the sr-only
  tables carrying the numbers. TWO sr-only tables in a wrapper div (never
  sr-only on a table itself — CSS width is only a lower bound for a table box):
  every head with entropy, span and its four shares, and every weight of the
  selected map with the strongest key and row entropy per row.
- Motion: the only animation is the skeleton pulse, plus colour transitions on
  the options and the sort buttons; all three carry motion-reduce.

Customization levers
- Grid density: the miniature cap (max-w-12) and the grid gap decide how many
  heads fit before the card gets tall; drop the H-number caption for a denser
  contact sheet, or raise the cap for a six-head slice you want to actually read.
- The slice itself: send fewer layers or fewer heads rather than filtering here —
  the header's "N of M" is derived from what you send against model.layers x
  model.headsPerLayer, so cropping upstream keeps the denominator honest.
- The ramp: ALPHA_GAMMA under 1 lifts the faint texture, 1 is linear, above 1
  hides everything but the peaks; WEIGHT_FLOOR decides what counts as noise.
  Raise both for a poster, lower both for a debugging session.
- The bar for a pattern: ARCHETYPE_BAR (45%) and DIFFUSE_BAR (85%) are printed
  on the card, so move them together with the sentence. Add an archetype by
  inserting one test into archetypeOf, before the tests it should outrank.
- Sort keys: orderHeads is a switch over a literal union — add "self", "sink" or
  "induction" descending in three lines, and the radiogroup picks it up.
- Units: swap Math.log for Math.log2 in rowEntropy to report bits and 2^H spans;
  the card quotes nats because pretraining loss does.
- Layout: the grid and the detail panel are one auto-fit CSS grid, so they sit
  side by side on a wide card and stack on a narrow one. Pin one column for a
  report page, or drop the detail panel entirely and emit onSelect into a
  page-level inspector.
- Interaction: for a static figure, drop the listbox wiring and render the grid
  as role="img" with a fixed `selected`; for a shared cursor across a page of
  models, keep `selected` driven from above — it already outranks the internal
  pick whenever it changes.

Concepts

  • A softmax row is a partition, so an absent cell is a zero — the sparse form drops the tail on purpose, and the renderer treats "not sent" exactly like "sent as 0". That is the opposite of a generic heatmap, where a missing cell means nobody measured; here the row already told you it spent its whole unit of attention somewhere, so nothing can be missing from it. The ink follows the same rule: one weight-to-opacity function paints the miniatures, the enlarged map and the legend swatches over the same muted canvas, and anything under 2% of a row is softmax floor rather than a link, so it stays bare canvas too.
  • Row 0 is excluded from every mean — a first token can only attend to itself, so its entropy is zero by construction. Averaging it in would make every head on a short prompt look sharper than the same head on a long one, which turns a property of the prompt into a claim about the model.
  • Precedence makes the shares disjoint — self, then previous, then the first token, then anything earlier. The row-1 cell is both "the first token" and "the previous token", and counting it twice would let a head's mass add to more than itself; the most local relationship wins, the rule is printed, and largest-remainder apportionment keeps the four integers at exactly 100 for every head you click.
  • Induction is tested before sink, on purpose — an induction head parks on the first token on every row that has nothing to match, so a sink test placed first would relabel the very pattern the card exists to surface. The score has its own denominator (the rows that repeat an earlier token) and is printed with it, because a mean over three rows is not a mean over ten.
  • Sorting inside a layer, never across — depth is structural: layer 2 is not "better" than layer 27, and a global ranking would throw away the one axis that means something on its own. Re-ranking within each row instead turns the left column into "the sharpest head in this layer" or "the most previous-token-ish head in this layer", which is how the first hop of an induction circuit surfaces without anyone hunting for it.
  • Selection, not hover — the grid is a single-select listbox and the enlarged map is its detail view, so the panel only moves when the reader asks. A scan-on-hover readout works for one curve and fails for thirty-six maps: the panel would flicker under every pointer move, and nothing that flickers can be compared to anything.

On This Page