Spaces:

firobeid
/

RoPEHeatmapMatrix

Running

File size: 20,887 Bytes

<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>RoPE — Full Rotation Matrix</title>
<style>
  :root {
    --color-background-primary: #ffffff;
    --color-background-secondary: #f5f4f0;
    --color-background-tertiary: #eeecea;
    --color-background-info: #e6f1fb;
    --color-text-primary: #1a1a18;
    --color-text-secondary: #73726c;
    --color-text-info: #185fa5;
    --color-border-tertiary: rgba(0,0,0,0.12);
    --color-border-secondary: rgba(0,0,0,0.22);
    --color-border-info: #185fa5;
    --font-sans: "Anthropic Sans", system-ui, sans-serif;
  }
  @media (prefers-color-scheme: dark) {
    :root {
      --color-background-primary: #1e1e1c;
      --color-background-secondary: #2a2a27;
      --color-background-tertiary: #222220;
      --color-background-info: #0c2a42;
      --color-text-primary: #e8e6de;
      --color-text-secondary: #9c9a92;
      --color-text-info: #85b7eb;
      --color-border-tertiary: rgba(255,255,255,0.12);
      --color-border-secondary: rgba(255,255,255,0.22);
      --color-border-info: #378add;
    }
  }
  * { box-sizing: border-box; margin: 0; padding: 0; }
  body { font-family: var(--font-sans); background: var(--color-background-primary); color: var(--color-text-primary); padding: 1.5rem; }
  .pe-wrap { padding: 0.5rem 0 1rem; }
  .pe-title { font-size: 15px; font-weight: 500; color: var(--color-text-primary); margin-bottom: 4px; }
  .pe-sub { font-size: 12px; color: var(--color-text-secondary); margin-bottom: 14px; line-height: 1.6; }
  .pe-axis-label { font-size: 12px; color: var(--color-text-secondary); text-align: center; }
  .pe-y-label { writing-mode: vertical-rl; transform: rotate(180deg); font-size: 12px; color: var(--color-text-secondary); display: flex; align-items: center; justify-content: center; min-width: 20px; }
  .pe-row { display: flex; align-items: stretch; gap: 8px; }
  .pe-legend { display: flex; align-items: center; gap: 8px; margin-top: 8px; font-size: 11px; color: var(--color-text-secondary); }
  .annotation-box { background: var(--color-background-secondary); border: 0.5px solid var(--color-border-tertiary); border-radius: 8px; padding: 10px 14px; margin-top: 12px; font-size: 12px; color: var(--color-text-secondary); line-height: 1.6; }
  .ann-row { display: flex; gap: 16px; }
  .ann-col { flex: 1; }
  .ann-head { font-weight: 500; font-size: 12px; color: var(--color-text-primary); margin-bottom: 3px; }
  .controls-row { display: flex; gap: 8px; flex-wrap: wrap; margin-bottom: 8px; align-items: center; }
  .pe-btn { background: var(--color-background-secondary); border: 0.5px solid var(--color-border-secondary); border-radius: 6px; padding: 4px 11px; font-size: 12px; cursor: pointer; color: var(--color-text-primary); }
  .pe-btn.active { background: var(--color-background-info); color: var(--color-text-info); border-color: var(--color-border-info); }
  .tooltip { position: absolute; background: var(--color-background-primary); border: 0.5px solid var(--color-border-secondary); border-radius: 6px; padding: 6px 10px; font-size: 11px; color: var(--color-text-primary); pointer-events: none; opacity: 0; transition: opacity 0.15s; white-space: nowrap; z-index: 10; line-height: 1.6; }
  .badge { display: inline-block; font-size: 10px; padding: 2px 6px; border-radius: 4px; background: var(--color-background-info); color: var(--color-text-info); border: 0.5px solid var(--color-border-info); margin-left: 6px; }
  .divider { border: none; border-top: 0.5px solid var(--color-border-tertiary); margin: 8px 0; }
  .dual-wrap { display: flex; gap: 6px; flex: 1; }
  .matrix-panel { flex: 1; position: relative; }
</style>
</head>
<body>
<div class="pe-wrap">
  <div class="pe-title">RoPE — Full Rotation Matrix <span class="badge">Rotary Position Encoding</span></div>
  <div class="pe-sub">
    RoPE rotates each dimension pair (q<sub>2i</sub>, q<sub>2i+1</sub>) by angle <strong>pos · θ<sub>i</sub></strong> using a 2×2 rotation matrix — requiring both <strong style="color:#0F6E56">cos</strong> and <strong style="color:#BA7517">sin</strong>. Toggle between three views of the same underlying data.
  </div>

  <div class="controls-row">
    <span style="font-size:12px;color:var(--color-text-secondary)">View:</span>
    <button class="pe-btn active" id="btn-view-sidebyside" onclick="setView('sidebyside')">cos | sin side by side</button>
    <button class="pe-btn" id="btn-view-interleaved" onclick="setView('interleaved')">interleaved (as in code)</button>
    <button class="pe-btn" id="btn-view-angle" onclick="setView('angle')">rotation angle φ = pos·θ</button>
  </div>
  <div class="controls-row" style="margin-bottom:10px">
    <span style="font-size:12px;color:var(--color-text-secondary)">Hover mode:</span>
    <button class="pe-btn active" id="btn-cell" onclick="setMode('cell')">Cell value</button>
    <button class="pe-btn" id="btn-col" onclick="setMode('col')">Column (dim pair)</button>
    <button class="pe-btn" id="btn-row" onclick="setMode('row')">Row (position)</button>
  </div>

  <!-- Side-by-side view -->
  <div id="view-sidebyside">
    <div style="display:flex;gap:6px;margin-left:28px;margin-bottom:3px">
      <div style="flex:1;text-align:center;font-size:11px;font-weight:500;color:#0F6E56;background:rgba(29,158,117,0.1);border-radius:4px;padding:2px 0">cos(pos · θ<sub>i</sub>)</div>
      <div style="flex:1;text-align:center;font-size:11px;font-weight:500;color:#854F0B;background:rgba(186,117,23,0.1);border-radius:4px;padding:2px 0">sin(pos · θ<sub>i</sub>)</div>
    </div>
    <div class="pe-row">
      <div class="pe-y-label">Sequence position <em>(pos ↓)</em></div>
      <div class="dual-wrap">
        <div class="matrix-panel">
          <canvas id="pe-canvas-cos" style="width:100%;display:block;cursor:crosshair;"></canvas>
          <div class="tooltip" id="tt-cos"></div>
        </div>
        <div style="width:1px;background:var(--color-border-tertiary);margin:0 1px;flex-shrink:0"></div>
        <div class="matrix-panel">
          <canvas id="pe-canvas-sin" style="width:100%;display:block;cursor:crosshair;"></canvas>
          <div class="tooltip" id="tt-sin"></div>
        </div>
      </div>
    </div>
    <div style="margin-left:28px">
      <div class="pe-axis-label">Dimension pair index <em>(i →)</em></div>
      <div style="display:flex;gap:6px;margin-top:2px">
        <div style="flex:1;display:flex;justify-content:space-between;font-size:10px;color:var(--color-text-secondary)"><span>0 (fast)</span><span>64</span><span>127 (slow)</span></div>
        <div style="width:3px;flex-shrink:0"></div>
        <div style="flex:1;display:flex;justify-content:space-between;font-size:10px;color:var(--color-text-secondary)"><span>0 (fast)</span><span>64</span><span>127 (slow)</span></div>
      </div>
    </div>
    <div class="pe-legend" style="margin-left:28px">
      <span>−1</span>
      <canvas id="legend-bar-sbs" height="12" style="flex:1;border-radius:3px;"></canvas>
      <span>+1</span>
    </div>
  </div>

  <!-- Interleaved view -->
  <div id="view-interleaved" style="display:none">
    <div style="margin-left:28px;margin-bottom:3px;font-size:11px;color:var(--color-text-secondary)">
      Columns alternate <strong style="color:#0F6E56">cos</strong>, <strong style="color:#854F0B">sin</strong>, <strong style="color:#0F6E56">cos</strong>, <strong style="color:#854F0B">sin</strong>… matching how <code>cos_pos</code> and <code>sin_pos</code> are broadcast in the code via <code>repeat(..., rep=2)</code>.
    </div>
    <div class="pe-row">
      <div class="pe-y-label">Sequence position <em>(pos ↓)</em></div>
      <div style="flex:1;position:relative;">
        <canvas id="pe-canvas-interleaved" style="width:100%;display:block;cursor:crosshair;"></canvas>
        <div class="tooltip" id="tt-il"></div>
      </div>
    </div>
    <div style="margin-left:28px">
      <div class="pe-axis-label">Embedding dimension d (pairs of columns = one rotation)</div>
      <div style="display:flex;justify-content:space-between;font-size:10px;color:var(--color-text-secondary);margin-top:2px"><span>d=0 cos</span><span>d=128</span><span>d=255 sin</span></div>
    </div>
    <div class="pe-legend" style="margin-left:28px">
      <span>−1</span>
      <canvas id="legend-bar-il" height="12" style="flex:1;border-radius:3px;"></canvas>
      <span>+1</span>
    </div>
  </div>

  <!-- Angle view -->
  <div id="view-angle" style="display:none">
    <div style="margin-left:28px;margin-bottom:3px;font-size:11px;color:var(--color-text-secondary)">
      The raw rotation angle φ = pos · θ<sub>i</sub> in radians. Both cos and sin are derived from this single scalar. Colour = angle magnitude (dark = 0, bright = max).
    </div>
    <div class="pe-row">
      <div class="pe-y-label">Sequence position <em>(pos ↓)</em></div>
      <div style="flex:1;position:relative;">
        <canvas id="pe-canvas-angle" style="width:100%;display:block;cursor:crosshair;"></canvas>
        <div class="tooltip" id="tt-ang"></div>
      </div>
    </div>
    <div style="margin-left:28px">
      <div class="pe-axis-label">Dimension pair index <em>(i →)</em></div>
      <div style="display:flex;justify-content:space-between;font-size:10px;color:var(--color-text-secondary);margin-top:2px"><span>0 (fast, large θ)</span><span>64</span><span>127 (slow, tiny θ)</span></div>
    </div>
    <div class="pe-legend" style="margin-left:28px">
      <span>0 rad</span>
      <canvas id="legend-bar-ang" height="12" style="flex:1;border-radius:3px;"></canvas>
      <span id="ang-max-label">max rad</span>
    </div>
  </div>

  <div class="annotation-box" id="ann-box"></div>
</div>

<script>
const ROWS = 50, COLS = 128, D_MODEL = 256;
const isDark = matchMedia('(prefers-color-scheme: dark)').matches;

function theta(i) { return Math.pow(10000, -(2 * i) / D_MODEL); }

const cosMatrix = [], sinMatrix = [], ilMatrix = [], angMatrix = [];
const maxAngle = (ROWS - 1) * theta(0);

for (let r = 0; r < ROWS; r++) {
  cosMatrix[r] = []; sinMatrix[r] = []; ilMatrix[r] = []; angMatrix[r] = [];
  for (let c = 0; c < COLS; c++) {
    const phi = r * theta(c);
    cosMatrix[r][c] = Math.cos(phi);
    sinMatrix[r][c] = Math.sin(phi);
    angMatrix[r][c] = phi;
  }
  for (let d = 0; d < 256; d++) {
    const pair = Math.floor(d / 2), phi = r * theta(pair);
    ilMatrix[r][d] = (d % 2 === 0) ? Math.cos(phi) : Math.sin(phi);
  }
}

// Diverging color for trig values [-1, 1]: green for positive, purple for negative
function trig2color(v, alpha) {
  const t = (v + 1) / 2;
  if (isDark) {
    if (t > 0.5) {
      const s = (t - 0.5) * 2;
      return `rgba(${Math.round(20+s*10)},${Math.round(60+s*150)},${Math.round(70+s*60)},${alpha})`;
    } else {
      const s = (0.5 - t) * 2;
      return `rgba(${Math.round(40+s*120)},${Math.round(20+s*20)},${Math.round(70+s*150)},${alpha})`;
    }
  } else {
    if (t > 0.5) {
      const s = (t - 0.5) * 2;
      return `rgba(${Math.round(245-s*200)},${Math.round(252-s*70)},${Math.round(240-s*170)},${alpha})`;
    } else {
      const s = (0.5 - t) * 2;
      return `rgba(${Math.round(235-s*50)},${Math.round(225-s*170)},${Math.round(252-s*30)},${alpha})`;
    }
  }
}

// Sequential amber for angle values [0, maxAngle]
function angle2color(v, alpha) {
  const t = Math.min(v / maxAngle, 1);
  if (isDark) {
    return `rgba(${Math.round(15+t*230)},${Math.round(15+t*150)},${Math.round(15+t*10)},${alpha})`;
  } else {
    return `rgba(${Math.round(252-t*30)},${Math.round(252-t*145)},${Math.round(252-t*235)},${alpha})`;
  }
}

function drawCanvas(canvas, matrix, colorFn, cols, hc, hr) {
  const w = canvas.offsetWidth;
  if (!w) return;
  const rows = matrix.length;
  const h = Math.round(w * (rows / cols) * 1.4);
  const dpr = window.devicePixelRatio || 1;
  canvas.width = w * dpr; canvas.height = h * dpr;
  canvas.style.height = h + 'px';
  const ctx = canvas.getContext('2d');
  ctx.scale(dpr, dpr);
  const cw = w / cols, ch = h / rows;
  const dim = hc >= 0 || hr >= 0;
  for (let r = 0; r < rows; r++) {
    for (let c = 0; c < cols; c++) {
      const isHL = (hc >= 0 && c === hc) || (hr >= 0 && r === hr);
      ctx.fillStyle = colorFn(matrix[r][c], dim && !isHL ? 0.28 : 1);
      ctx.fillRect(c * cw, r * ch, cw + 0.5, ch + 0.5);
    }
  }
  if (hc >= 0) { ctx.strokeStyle='#EF9F27'; ctx.lineWidth=2; ctx.strokeRect(hc*cw+1,1,cw-1,rows*ch-2); }
  if (hr >= 0) { ctx.strokeStyle='#1D9E75'; ctx.lineWidth=2; ctx.strokeRect(1,hr*ch+1,cols*cw-2,ch-1); }
}

function drawLegend(id, colorFn, minV, maxV) {
  const lb = document.getElementById(id); if (!lb) return;
  const lctx = lb.getContext('2d');
  const w = lb.offsetWidth || 300, dpr = window.devicePixelRatio || 1;
  lb.width = w * dpr; lb.style.width = '100%';
  for (let x = 0; x < w; x++) {
    lctx.fillStyle = colorFn(minV + (x / w) * (maxV - minV), 1);
    lctx.fillRect(x * dpr, 0, dpr, 12);
  }
}

let currentView = 'sidebyside', mode = 'cell';
let hlCol = -1, hlRow = -1;

function redraw(hc, hr) {
  if (hc !== undefined) { hlCol = hc; hlRow = hr; }
  const hc_ = hlCol, hr_ = hlRow;
  if (currentView === 'sidebyside') {
    drawCanvas(document.getElementById('pe-canvas-cos'), cosMatrix, trig2color, COLS, hc_, hr_);
    drawCanvas(document.getElementById('pe-canvas-sin'), sinMatrix, trig2color, COLS, hc_, hr_);
    drawLegend('legend-bar-sbs', trig2color, -1, 1);
  } else if (currentView === 'interleaved') {
    const ilHc = hc_ >= 0 ? hc_ * 2 : -1;
    drawCanvas(document.getElementById('pe-canvas-interleaved'), ilMatrix, trig2color, 256, ilHc, hr_);
    drawLegend('legend-bar-il', trig2color, -1, 1);
  } else {
    drawCanvas(document.getElementById('pe-canvas-angle'), angMatrix, angle2color, COLS, hc_, hr_);
    drawLegend('legend-bar-ang', angle2color, 0, maxAngle);
    document.getElementById('ang-max-label').textContent = maxAngle.toFixed(1) + ' rad';
  }
}

function clearHL() { hlCol = -1; hlRow = -1; redraw(); }

function setView(v) {
  currentView = v;
  ['sidebyside','interleaved','angle'].forEach(k => {
    document.getElementById('view-'+k).style.display = k===v ? '' : 'none';
    document.getElementById('btn-view-'+k).classList.toggle('active', k===v);
  });
  clearHL(); updateAnnotation();
}

function setMode(m) {
  mode = m;
  ['cell','col','row'].forEach(k => document.getElementById('btn-'+k).classList.toggle('active', k===m));
  clearHL();
}

function attachHover(canvasId, ttId, colsOverride, getInfo) {
  const canvas = document.getElementById(canvasId);
  const tt = document.getElementById(ttId);
  canvas.addEventListener('mousemove', function(e) {
    const rect = canvas.getBoundingClientRect();
    const x = e.clientX - rect.left, y = e.clientY - rect.top;
    const cols = colsOverride || COLS;
    const col = Math.floor(x / (rect.width / cols));
    const row = Math.floor(y / (rect.height / ROWS));
    if (col < 0 || col >= cols || row < 0 || row >= ROWS) return;
    const info = getInfo(row, col);
    if (mode === 'cell')      { redraw(-1, -1); tt.innerHTML = info.cell; }
    else if (mode === 'col')  { redraw(info.pairCol, -1); tt.innerHTML = info.colTip; }
    else                      { redraw(-1, row); tt.innerHTML = info.rowTip; }
    tt.style.left = Math.min(x+10, rect.width-250)+'px';
    tt.style.top = Math.max(y-44,0)+'px';
    tt.style.opacity = '1';
  });
  canvas.addEventListener('mouseleave', () => { tt.style.opacity='0'; clearHL(); });
}

function pairInfo(row, col) {
  const phi = (row * theta(col)).toFixed(4);
  const cv = Math.cos(row*theta(col)).toFixed(4);
  const sv = Math.sin(row*theta(col)).toFixed(4);
  const th = theta(col).toExponential(3);
  const speed = col < 10 ? 'fast' : col < 60 ? 'moderate' : 'slow';
  return { phi, cv, sv, th, speed };
}

attachHover('pe-canvas-cos', 'tt-cos', null, (row, col) => {
  const { phi, cv, sv, th, speed } = pairInfo(row, col);
  return {
    pairCol: col,
    cell:   `<b>pos=${row}, pair i=${col}</b><br>cos(φ) = ${cv}<br>sin(φ) = ${sv}<br>φ = ${phi} rad`,
    colTip: `<b>Pair i=${col}</b> (${speed}) · θ = ${th}<br>cos col highlighted left, sin col right`,
    rowTip: `<b>pos=${row}</b> · full rotation vector (all pairs)`
  };
});

attachHover('pe-canvas-sin', 'tt-sin', null, (row, col) => {
  const { phi, cv, sv, th, speed } = pairInfo(row, col);
  return {
    pairCol: col,
    cell:   `<b>pos=${row}, pair i=${col}</b><br>sin(φ) = ${sv}<br>cos(φ) = ${cv}<br>φ = ${phi} rad`,
    colTip: `<b>Pair i=${col}</b> (${speed}) · θ = ${th}`,
    rowTip: `<b>pos=${row}</b> · full rotation vector (all pairs)`
  };
});

attachHover('pe-canvas-interleaved', 'tt-il', 256, (row, d) => {
  const pair = Math.floor(d / 2), isCos = d % 2 === 0;
  const { phi, cv, sv, th, speed } = pairInfo(row, pair);
  return {
    pairCol: d,
    cell:   `<b>pos=${row}, d=${d}</b> (pair ${pair}, ${isCos?'cos':'sin'})<br>value = ${isCos?cv:sv}<br>φ = ${phi} rad`,
    colTip: `<b>d=${d}</b> · pair ${pair} · ${isCos?'<b style="color:#1D9E75">cos</b>':'<b style="color:#BA7517">sin</b>'} channel`,
    rowTip: `<b>pos=${row}</b> · interleaved rotation vector (256 dims)`
  };
});

attachHover('pe-canvas-angle', 'tt-ang', null, (row, col) => {
  const { phi, cv, sv, th, speed } = pairInfo(row, col);
  return {
    pairCol: col,
    cell:   `<b>pos=${row}, pair i=${col}</b><br>φ = ${phi} rad<br>→ cos = ${cv}, sin = ${sv}`,
    colTip: `<b>Pair i=${col}</b> (${speed}) · θ = ${th}<br>angle grows by θ per position step`,
    rowTip: `<b>pos=${row}</b> · angle = pos × θ<sub>i</sub> for each pair`
  };
});

function updateAnnotation() {
  const ann = document.getElementById('ann-box');
  if (currentView === 'sidebyside') {
    ann.innerHTML = `<div class="ann-row">
      <div class="ann-col">
        <div class="ann-head" style="color:#0F6E56">cos matrix (left) — the "stay" component</div>
        cos(pos·θ<sub>i</sub>) scales how much of the original vector survives. At φ=0 it is 1 (identity). At φ=π it is −1 (full flip). Fast pairs (left columns) cycle rapidly; slow pairs (right) stay near 1 across all 50 positions.
      </div>
      <div class="ann-col">
        <div class="ann-head" style="color:#854F0B">sin matrix (right) — the "spin" component</div>
        sin(pos·θ<sub>i</sub>) scales the 90°-rotated copy <code>qw2 = [−q<sub>odd</sub>, q<sub>even</sub>]</code>. At φ=0 it is 0 (no spin); at φ=π/2 it is 1 (maximum rotation contribution). Together cos² + sin² = 1, preserving vector length exactly.
      </div>
    </div>
    <hr class="divider">
    <div style="font-size:11px"><b>Full rotation formula:</b> q<sub>new</sub> = q · cos_pos + [−q<sub>odd</sub>, q<sub>even</sub>] · sin_pos</div>`;
  } else if (currentView === 'interleaved') {
    ann.innerHTML = `<div class="ann-row">
      <div class="ann-col">
        <div class="ann-head">How the code broadcasts</div>
        <code>cos_pos = repeat(sinusoidal[..., 1::2], rep=2)</code> duplicates each cos value twice, so both dims of a pair share the same scalar.<br>
        <code>sin_pos = repeat(sinusoidal[..., ::2], rep=2)</code> does the same for sin.<br><br>
        Each consecutive pair of columns (d, d+1) gets the same rotation angle — only the trig function differs.
      </div>
      <div class="ann-col">
        <div class="ann-head">What to notice</div>
        Adjacent columns are always identical in value (the repeat). Visually this makes each pair look like a 2-pixel-wide stripe. Fast pairs (left) stripe rapidly across positions; slow pairs (right) are nearly uniform. The pattern <em>is</em> the side-by-side view — just interleaved instead of separated.
      </div>
    </div>`;
  } else {
    ann.innerHTML = `<div class="ann-row">
      <div class="ann-col">
        <div class="ann-head">The underlying clock: φ = pos · θ<sub>i</sub></div>
        Both cos and sin are simply read off the unit circle at this angle. The left column (pair 0) reaches ${maxAngle.toFixed(1)} rad by pos=49 — many full revolutions. The right column barely ticks above 0. This triangular ramp is the speed spectrum that gives RoPE its multi-scale position sensitivity.
      </div>
      <div class="ann-col">
        <div class="ann-head">Why the gradient is triangular</div>
        φ = pos × θ<sub>i</sub>, where θ<sub>i</sub> = 10000<sup>−2i/d</sup> decreases geometrically rightward and pos increases linearly downward. The brightest cell is bottom-left (pos=49, pair 0). Moving right, the brightness drops exponentially; moving up, it drops linearly. This is the "clock" RoPE winds at different speeds per pair.
      </div>
    </div>`;
  }
}

setTimeout(() => { redraw(); updateAnnotation(); }, 60);
window.addEventListener('resize', () => redraw());
</script>
</body>
</html>