Attention Heatmap

A grid heatmap for visualising attention weights between query tokens (rows) and key tokens (columns). Cells are color-tinted by weight; hovering one highlights its full row and column so the "which query attends to which key" relationship reads at a glance.

Installation

npx shadcn@latest add https://craftbits.dev/r/attention-heatmap.json

Usage

import { AttentionHeatmap } from "@craft-bits/core";
 
const weights = [
  [0.62, 0.18, 0.12, 0.08],
  [0.22, 0.51, 0.17, 0.10],
  [0.10, 0.15, 0.60, 0.15],
  [0.04, 0.09, 0.21, 0.66],
];
 
<AttentionHeatmap
  weights={weights}
  queries={["the", "cat", "sat", "down"]}
  keys={["the", "cat", "sat", "down"]}
/>

Drive the hover state from outside the component to sync the cross-highlight with a sidebar inspector or external tooltip:

const [hovered, setHovered] = useState<[number, number] | null>(null);
 
<AttentionHeatmap
  weights={weights}
  hoveredCell={hovered}
  onHoveredCellChange={setHovered}
/>

Understanding the component

1. Grid layout, real elements. The matrix is a CSS grid with one extra row + column for axis labels. Each cell is a motion.div — there is no <svg>, so cells get focus rings, real text rendering, and full hit areas for free.
2. Color encodes weight. Each cell's background-color is the active color scale tinted by 0.05 + weight * 0.9 alpha. Even tiny weights show a hint of fill; full-attention cells render at near-solid. The label color flips to --cb-accent-fg past ~0.55 weight so contrast stays AA on both themes.
3. Three color scales. accent (default) tints var(--cb-accent); mono tints var(--cb-fg) for a high-contrast monochrome ramp; spectrum walks an oklch hue ramp for viridis-like depth.
4. Cross-highlight on hover. Hovering a cell sets the hovered [row, col]. The row + column labels flip to --cb-accent, and every cell outside the hovered row or column dims to 0.35 opacity. The cross stays at full intensity, making "this query attends here" pop.
5. Controlled + uncontrolled. Pass hoveredCell + onHoveredCellChange to drive hover from a parent; omit them to let the component own its own state.
6. Reduced motion. When prefers-reduced-motion: reduce is set, the dim transition collapses to duration: 0 — the highlight snaps instead of springs.

Props

Accessibility

• The outer element is role="figure" with a visually hidden caption announcing matrix dimensions.
• When interactive, every cell is role="button" with an aria-label like "Attention from query 1 to key 3: 0.51" and a visible focus-visible ring keyed to --cb-accent.
• Keyboard focus drives the same cross-highlight as mouse hover — the relationship is reachable without a pointer.
• The dim transition respects prefers-reduced-motion and collapses to an instant swap.
• Color is never the only signal — every cell renders its numeric weight, so the matrix stays legible for colorblind users and at every color-scale setting.

Credits

• Extracted from: craftingattention (app/src/lessons/primitives/viz/AttentionHeatmap.tsx). The original was an interactive lesson widget bundled with embeddings, softmax, and four challenge modes; the library extract is the heatmap primitive — the lesson scaffolding lives in the source project.