feat: add formla of Q, K matrix calculation

app.py

@@ -77,6 +77,47 @@ Shaw relative attention (learned bias `b_{m-n}` on scores) is a **different**
 mechanism and is not computed here.
 """
 
+BULK_FORMULAS_MD = r"""
+### What each Q/K matrix entry means
+
+The selected heatmap entry is `X[h, k, j]`, where `X` is either **Q** or **K**, `h` is
+the attention-head index, `k` is the token position, and `j` is the dimension within
+that head.
+
+For a real model, the first-layer projections calculate each entry as:
+
+$$
+Q^{before}_{k,h,j} = \left(H_k W_Q + b_Q\right)_{h d + j},\qquad
+K^{before}_{k,h,j} = \left(H_k W_K + b_K\right)_{h d + j}
+$$
+
+Here `H_k` is the token's hidden vector, `d` is the dimension per head, and `W_Q`,
+`W_K` are the model's query/key projection weights. In random-matrix mode, the initial
+entries are sampled directly: `Q before ~ N(0, 1)` and `K before ~ N(0, 1)` using
+independent seeds.
+
+RoPE then transforms either matrix using
+
+$$
+\theta(k,i) = k\,base^{-2i/d}
+$$
+
+For the real-model Llama layout, the paired dimensions are `(i, i+d/2)`:
+
+$$
+X^{after}_{h,k,i} = X^{before}_{h,k,i}\cos\theta -
+X^{before}_{h,k,i+d/2}\sin\theta
+$$
+$$
+X^{after}_{h,k,i+d/2} = X^{before}_{h,k,i}\sin\theta +
+X^{before}_{h,k,i+d/2}\cos\theta
+$$
+
+where `X` means either **Q** or **K**. Random-matrix mode uses adjacent pairs
+`(2i, 2i+1)` instead. Therefore, **Before** contains projected or sampled values,
+**After** contains rotated values, and **Delta = After − Before**.
+"""
+
 PLACEHOLDER = go.Figure().update_layout(
     title="Run **Compute** on the Setup tab first",
     template="plotly_white",
@@ -381,6 +422,7 @@ with gr.Blocks(title="RoPE Explorer") as demo:
             head = gr.Slider(minimum=0, maximum=2, step=1, value=0, label="Head index (real models)")
 
         with gr.Tab("Bulk changes"):
+            gr.Markdown(BULK_FORMULAS_MD)
             which = gr.Radio(["Q", "K"], value="Q", label="Tensor")
             mod_2pi = gr.Checkbox(False, label="θ heatmap: wrap mod 2π")
             bulk_main = gr.Plot(label="Before / after / delta")

src/plots.py

@@ -59,13 +59,16 @@ def bulk_before_after_delta(before: np.ndarray, after: np.ndarray, tokens=None)
             go.Heatmap(
                 z=mat,
                 showscale=(i == 3),
-                hovertemplate="token=%{y}<br>dim=%{x}<br>z=%{z:.4f}<extra></extra>",
+                colorbar=dict(title="matrix value") if i == 3 else None,
+                hovertemplate="token position k=%{y}<br>dimension index j=%{x}<br>value=%{z:.4f}<extra></extra>",
             ),
             row=1,
             col=i,
         )
+        fig.update_xaxes(title_text="dimension index j (within selected head)", row=1, col=i)
+        fig.update_yaxes(title_text="token position k", row=1, col=i)
         fig.update_yaxes(autorange="reversed", row=1, col=i)
-    fig.update_layout(template="plotly_white", height=400, margin=dict(t=60))
+    fig.update_layout(template="plotly_white", height=430, margin=dict(t=60, b=65))
     return fig
 
 
@@ -78,10 +81,10 @@ def norms_and_cosine(before: np.ndarray, after: np.ndarray, tokens=None) -> go.F
     fig.add_trace(go.Scatter(x=xs, y=nb, name="before", mode="lines+markers"), row=1, col=1)
     fig.add_trace(go.Scatter(x=xs, y=na, name="after", mode="lines+markers"), row=1, col=1)
     fig.add_trace(go.Bar(x=xs, y=cos, name="cosine", showlegend=False), row=1, col=2)
-    fig.update_xaxes(title_text="token", row=1, col=1)
-    fig.update_xaxes(title_text="token", row=1, col=2)
-    fig.update_yaxes(title_text="L2", row=1, col=1)
-    fig.update_yaxes(title_text="cosine", range=[min(0.0, float(np.min(cos)) - 0.05), 1.02], row=1, col=2)
+    fig.update_xaxes(title_text="token position k", row=1, col=1)
+    fig.update_xaxes(title_text="token position k", row=1, col=2)
+    fig.update_yaxes(title_text="L2 norm of Q/K vector", row=1, col=1)
+    fig.update_yaxes(title_text="cosine similarity (before, after)", range=[min(0.0, float(np.min(cos)) - 0.05), 1.02], row=1, col=2)
     fig.update_layout(template="plotly_white", height=380, barmode="group")
     return fig
 
@@ -97,13 +100,13 @@ def theta_heatmap(seq_len: int, dim: int, base: float, mod_2pi: bool = False) ->
     fig = go.Figure(
         data=go.Heatmap(
             z=z,
-            colorbar=dict(title="radians"),
+            colorbar=dict(title="rotation angle θ (radians)"),
             hovertemplate="token k=%{y}<br>pair i=%{x}<br>θ=%{z:.4f}<extra></extra>",
         )
     )
     fig.update_layout(
         title=title,
-        xaxis_title="pair index i",
+        xaxis_title="pair index i (dimension pair)",
         yaxis_title="token position k",
         yaxis=dict(autorange="reversed"),
         template="plotly_white",
@@ -117,8 +120,8 @@ def frequency_strip(dim: int, base: float) -> go.Figure:
     fig = go.Figure(data=go.Bar(x=list(range(len(omega))), y=omega, name="ω_i"))
     fig.update_layout(
         title="Pair frequencies ω_i (pair 0 is fastest)",
-        xaxis_title="pair index i",
-        yaxis_title="ω",
+        xaxis_title="pair index i (dimension pair)",
+        yaxis_title="frequency ω_i (radians per position)",
         yaxis_type="log",
         template="plotly_white",
         height=280,