feat(paper): overhaul figures pipeline — PDF output, hparam table, Pro results, lookahead variants

- Switch all figure exports from PNG to PDF; update dashboard to render PDFs via <embed>
- Merge run_figures.py + src/figures.py into run_paper_figures.py; delete old files
- Add paper/style.py: centralised visual style (colors, hatch, rcParams, figure dimensions)
- Add paper/__init__.py to make paper/ a package
- Add build_hparam_table() reading from paper/hparams.json; populate Verified values from
  confirmed SLURM training logs, Laguna Pro per-probe HPs from independent sweeps
- Extend build_auc_table() and build_calibration_table() with SWE-bench Pro sections
- Add plot_generalization_barplot(): main 2×2 model×dataset AUC barplot
- Add lookahead variants (max15, max50, _after_edit) via --lookahead-variants flag
- Bold best layer per row in AUC table (was: global best per probe across all conditions)
- Add colored deltas to transfer table
- Update dashboard: load figures from paper/figures/, PDF embed support
- Update slurm/figures.sh and README

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: andre15silva

README.md

@@ -2,7 +2,7 @@
 
 Measures whether a language model's internal hidden states linearly predict properties of its own output before those properties are realised.
 
-Pipeline: inference + activation hooks → per-sample `.pt` files → per-(layer, probe) cache tensors → linear probe training with W&B sweep → accuracy heatmaps → static dashboard.
+Pipeline: inference + activation hooks → per-sample `.pt` files → per-(layer, probe) cache tensors → linear probe training with W&B sweep → paper figures → static dashboard.
 
 ## Pipeline overview
 
@@ -27,7 +27,8 @@ run_probe.py final       train linear probe on best hparams → W&B run with met
                           (or pass --from-sweep to auto-load best HPs from a sweep)
         │
         ▼
-run_figures.py           accuracy heatmaps per layer × relative-position bin
+run_paper_figures.py     AUC heatmaps, layer plots, lookahead, barplot, LaTeX tables
+                         → paper/figures/  +  paper/*.tex
 run_export_dashboard.py  → dashboard/index.html  (open in browser)
 ```
 
@@ -62,7 +63,7 @@ run_attach_labels_swebench.py  load activations + trajectory + _labels.json,
                                → outputs/swebench/<run_id>_labeled/<instance_id>.pt
         │
         ▼
-run_build_cache.py  →  run_probe.py  →  run_figures.py   (same as standard path)
+run_build_cache.py  →  run_probe.py  →  run_paper_figures.py   (same as standard path)
 ```
 
 Dynamic probes (`currently_correct`, `currently_compiles`, `currently_reduces_failing`, `currently_has_regressions`) only make sense in the agentic path because they require `EditEvent` history with per-edit `test_results`. The carry-forward step in `run_attach_labels_swebench.py` expands one label per edit into one label per stride step; steps before the first edit are masked (`None` → `-1` → excluded from training).
@@ -258,10 +259,24 @@ Queries `sweep.best_run()` via the W&B API to auto-populate `lr`, `weight_decay`
 ### 4. Figures
 
 ```bash
-uv run python run_figures.py \
-  --run-id my_run \
-  --probe will_be_correct \
-  --model-config configs/models/qwen3_8b.yaml
+uv run python run_paper_figures.py \
+  --results-dir results/swebench \
+  --probes currently_compiles currently_correct currently_reduces_failing currently_has_regressions \
+  --model-run-ids laguna_xs2_full qwen36_35b_a3b_full \
+  --shuffled-run-ids laguna_xs2_full_shuffled qwen36_35b_a3b_full_shuffled
+```
+
+This generates all publication figures under `paper/figures/` and LaTeX tables under `paper/`.
+It also writes `paper/figures/manifest.json` for the dashboard figures gallery.
+
+For lookahead horizon figures, pass the shifted-label run IDs:
+
+```bash
+uv run python run_paper_figures.py \
+  ... \
+  --lookahead-run-ids laguna_xs2_full_shift0_max50 laguna_xs2_full_shift1_max50 ... laguna_xs2_full_shift50_max50 \
+  --lookahead-k-values $(seq 0 50) \
+  --lookahead-filename-suffix max50
 ```
 
 ### 5. Export dashboard
@@ -311,8 +326,11 @@ sbatch slurm/probe_final.sh \
   --model-config configs/models/qwen3_8b.yaml \
   --from-sweep abc123xyz
 
-sbatch slurm/figures.sh --run-id my_run --probe will_be_correct \
-  --model-config configs/models/qwen3_8b.yaml
+sbatch slurm/figures.sh \
+  --results-dir results/swebench \
+  --probes currently_compiles currently_correct currently_reduces_failing currently_has_regressions \
+  --model-run-ids laguna_xs2_full qwen36_35b_a3b_full \
+  --shuffled-run-ids laguna_xs2_full_shuffled qwen36_35b_a3b_full_shuffled
 
 sbatch slurm/export_dashboard.sh \
   --run-id my_run --probe will_be_correct \

dashboard/app.js

@@ -1,6 +1,7 @@
 "use strict";
 
 const DATA_ROOT = "data";
+const FIGURES_ROOT = "../paper/figures";
 let state = {
   manifest: [],
   runId: null,
@@ -807,7 +808,7 @@ function switchTab(name) {
 
 async function loadFiguresManifest() {
   try {
-    return await fetchJSON(`${DATA_ROOT}/figures/manifest.json`);
+    return await fetchJSON(`${FIGURES_ROOT}/manifest.json`);
   } catch {
     return [];
   }
@@ -824,7 +825,7 @@ function renderGallery() {
     );
 
     if (filtered.length === 0) {
-      container.innerHTML = '<p style="color:#888;font-size:13px;">No figures found. Run run_figures.py and re-export the dashboard.</p>';
+      container.innerHTML = '<p style="color:#888;font-size:13px;">No figures found. Run run_paper_figures.py to generate figures and manifest.</p>';
       return;
     }
 
@@ -853,19 +854,21 @@ function renderGallery() {
         card.style.cssText = "cursor:pointer;border:1px solid #ddd;border-radius:6px;overflow:hidden;width:200px;background:#fafafa;";
         card.title = fig.title || fig.path;
 
-        const img = document.createElement("img");
         const cacheBust = `?v=${Date.now()}`;
-        img.src = `${DATA_ROOT}/figures/${fig.path}${cacheBust}`;
-        img.style.cssText = "width:200px;height:130px;object-fit:contain;background:#fff;";
-        img.loading = "lazy";
+        const figSrc = `${FIGURES_ROOT}/${fig.path}${cacheBust}`;
+        const embed = document.createElement("embed");
+        embed.src = figSrc;
+        embed.type = "application/pdf";
+        embed.style.cssText = "width:200px;height:130px;background:#fff;pointer-events:none;";
 
         const caption = document.createElement("div");
-        caption.textContent = fig.title || fig.path.split("/").pop().replace(".png", "");
+        const ext = fig.path.split(".").pop();
+        caption.textContent = fig.title || fig.path.split("/").pop().replace(`.${ext}`, "");
         caption.style.cssText = "font-size:10px;color:#555;padding:4px 6px;white-space:nowrap;overflow:hidden;text-overflow:ellipsis;";
 
-        card.appendChild(img);
+        card.appendChild(embed);
         card.appendChild(caption);
-        card.addEventListener("click", () => openLightbox(`${DATA_ROOT}/figures/${fig.path}${cacheBust}`, fig.title || fig.path));
+        card.addEventListener("click", () => openLightbox(figSrc, fig.title || fig.path));
         grid.appendChild(card);
       }
 

dashboard/index.html

@@ -72,7 +72,7 @@ <h2>Detail View</h2>
 
   <!-- Lightbox overlay -->
   <div id="lightbox" style="display:none; position:fixed; inset:0; background:rgba(0,0,0,0.85); z-index:1000; align-items:center; justify-content:center; flex-direction:column;">
-    <img id="lightbox-img" style="max-width:92vw; max-height:88vh; border-radius:4px; box-shadow:0 4px 24px rgba(0,0,0,0.5);">
+    <embed id="lightbox-img" type="application/pdf" style="max-width:92vw; width:900px; height:88vh; border-radius:4px; box-shadow:0 4px 24px rgba(0,0,0,0.5);">
     <div id="lightbox-caption" style="color:#eee; font-size:12px; margin-top:10px;"></div>
     <button onclick="closeLightbox()" style="position:absolute;top:16px;right:22px;font-size:24px;background:none;border:none;color:#fff;cursor:pointer;">✕</button>
   </div>

paper/__init__.py

@@ -0,0 +1,40 @@
+\begin{table*}[h]
+\centering
+\caption{Test AUC-ROC per probe, model, and benchmark across transformer layers. \textbf{Bold} marks the best layer per row. \emph{Shuffled} uses label-permuted data as a sanity baseline.}
+\label{tab:auc_roc}
+\begin{tabular}{lccccc}
+\toprule
+ & \multicolumn{4}{c}{AUC-ROC $\uparrow$} & Shuffled \\
+\cmidrule(lr){2-5}
+Probe & 10 & 20 & 30 & 39 & (best layer) \\
+\midrule
+\multicolumn{6}{l}{\textsc{SWE-bench Verified}} \\
+\addlinespace[2pt]
+\multicolumn{6}{l}{\quad\textit{Laguna-XS2}} \\
+Syntactic correctness & \textbf{0.698} & 0.669 & 0.679 & 0.626 & 0.502 \\
+Semantic correctness & 0.775 & \textbf{0.791} & 0.769 & 0.716 & 0.501 \\
+Reduces failures & 0.743 & \textbf{0.769} & 0.743 & 0.674 & 0.501 \\
+Has regressions & 0.688 & \textbf{0.697} & 0.675 & 0.612 & 0.501 \\
+\addlinespace[3pt]
+\multicolumn{6}{l}{\quad\textit{Qwen3.6-35B-A3B}} \\
+Syntactic correctness & \textbf{0.777} & 0.755 & 0.771 & 0.692 & 0.499 \\
+Semantic correctness & 0.816 & 0.836 & \textbf{0.839} & 0.786 & 0.503 \\
+Reduces failures & 0.818 & \textbf{0.835} & 0.832 & 0.776 & 0.506 \\
+Has regressions & 0.723 & 0.773 & \textbf{0.778} & 0.698 & 0.500 \\
+\midrule
+\multicolumn{6}{l}{\textsc{SWE-bench Pro}} \\
+\addlinespace[2pt]
+\multicolumn{6}{l}{\quad\textit{Laguna-XS2}} \\
+Syntactic correctness & \textbf{0.780} & 0.746 & 0.757 & 0.756 & — \\
+Semantic correctness & 0.681 & \textbf{0.726} & 0.714 & 0.670 & — \\
+Reduces failures & 0.681 & \textbf{0.718} & 0.697 & 0.657 & — \\
+Has regressions & 0.645 & 0.705 & \textbf{0.739} & 0.691 & — \\
+\addlinespace[3pt]
+\multicolumn{6}{l}{\quad\textit{Qwen3.6-35B-A3B}} \\
+Syntactic correctness & — & — & — & — & — \\
+Semantic correctness & — & — & — & — & — \\
+Reduces failures & — & — & — & — & — \\
+Has regressions & — & — & — & — & — \\
+\bottomrule
+\end{tabular}
+\end{table*}

paper/auc_table.tex

@@ -0,0 +1,40 @@
+\begin{table*}[h]
+\centering
+\caption{Calibration metrics (ECE and Brier score) per probe across layers. Lower is better for both metrics.}
+\label{tab:calibration}
+\begin{tabular}{lcccccccc}
+\toprule
+Probe & \multicolumn{2}{c}{Layer 10} & \multicolumn{2}{c}{Layer 20} & \multicolumn{2}{c}{Layer 30} & \multicolumn{2}{c}{Layer 39} \\
+\cmidrule(lr){2-9}
+& ECE $\downarrow$ & Brier $\downarrow$ & ECE $\downarrow$ & Brier $\downarrow$ & ECE $\downarrow$ & Brier $\downarrow$ & ECE $\downarrow$ & Brier $\downarrow$ \\
+\midrule
+\multicolumn{9}{l}{\textsc{SWE-bench Verified}} \\
+\addlinespace[2pt]
+\multicolumn{9}{l}{\quad\textit{Laguna-XS2}} \\
+Syntactic correctness & 0.043 & 0.089 & 0.045 & 0.090 & 0.048 & 0.089 & 0.060 & 0.093 \\
+Semantic correctness & 0.062 & 0.186 & 0.040 & 0.178 & 0.039 & 0.185 & 0.051 & 0.203 \\
+Reduces failures & 0.058 & 0.203 & 0.060 & 0.194 & 0.050 & 0.202 & 0.072 & 0.226 \\
+Has regressions & 0.029 & 0.096 & 0.035 & 0.096 & 0.044 & 0.098 & 0.056 & 0.103 \\
+\addlinespace[3pt]
+\multicolumn{9}{l}{\quad\textit{Qwen3.6-35B-A3B}} \\
+Syntactic correctness & 0.042 & 0.076 & 0.042 & 0.076 & 0.041 & 0.075 & 0.036 & 0.077 \\
+Semantic correctness & 0.050 & 0.152 & 0.035 & 0.144 & 0.035 & 0.142 & 0.070 & 0.166 \\
+Reduces failures & 0.062 & 0.161 & 0.050 & 0.153 & 0.043 & 0.153 & 0.072 & 0.178 \\
+Has regressions & 0.010 & 0.060 & 0.012 & 0.059 & 0.013 & 0.059 & 0.030 & 0.064 \\
+\midrule
+\multicolumn{9}{l}{\textsc{SWE-bench Pro}} \\
+\addlinespace[2pt]
+\multicolumn{9}{l}{\quad\textit{Laguna-XS2}} \\
+Syntactic correctness & 0.061 & 0.189 & 0.041 & 0.198 & 0.056 & 0.196 & 0.070 & 0.200 \\
+Semantic correctness & 0.029 & 0.090 & 0.025 & 0.087 & 0.027 & 0.088 & 0.045 & 0.091 \\
+Reduces failures & 0.048 & 0.138 & 0.045 & 0.133 & 0.054 & 0.136 & 0.069 & 0.142 \\
+Has regressions & 0.013 & 0.079 & 0.011 & 0.077 & 0.018 & 0.076 & 0.008 & 0.076 \\
+\addlinespace[3pt]
+\multicolumn{9}{l}{\quad\textit{Qwen3.6-35B-A3B}} \\
+Syntactic correctness & — & — & — & — & — & — & — & — \\
+Semantic correctness & — & — & — & — & — & — & — & — \\
+Reduces failures & — & — & — & — & — & — & — & — \\
+Has regressions & — & — & — & — & — & — & — & — \\
+\bottomrule
+\end{tabular}
+\end{table*}

paper/calibration_table.tex

@@ -0,0 +1,19 @@
+\begin{table}[h]
+\centering
+\caption{Chosen hyperparameters for each model and benchmark, selected by 20-trial random search maximising mean validation AUC on the middle transformer layer. Pro sweeps were run independently per probe.}
+\label{tab:probe-hparams}
+\begin{tabular}{lllrrrr}
+\toprule
+Dataset & Model & Probe & Learning rate & Weight decay & Batch size & Patience \\
+\midrule
+\multirow{2}{*}{\textsc{Verified}} & Laguna-XS2 & & $3.42 \times 10^{-4}$ & $1.47 \times 10^{-3}$ & 256 & 50 \\
+ & Qwen3.6-35B-A3B & & $1.17 \times 10^{-3}$ & $9.81 \times 10^{-2}$ & 1024 & 10 \\
+\midrule
+\multirow{5}{*}{\textsc{Pro}} & \multirow{4}{*}{Laguna-XS2} & Syntactic correctness & $2.53 \times 10^{-4}$ & $2.75 \times 10^{-4}$ & 256 & 50 \\
+ &  & Semantic correctness & $6.16 \times 10^{-4}$ & $5.18 \times 10^{-3}$ & 1024 & 25 \\
+ &  & Reduces failures & $2.31 \times 10^{-4}$ & $4.04 \times 10^{-3}$ & 256 & 10 \\
+ &  & Has regressions & $1.51 \times 10^{-4}$ & $3.26 \times 10^{-2}$ & 1024 & 10 \\
+ & Qwen3.6-35B-A3B & & — & — & — & — \\
+\bottomrule
+\end{tabular}
+\end{table}

paper/hparam_table.tex

@@ -0,0 +1,15 @@
+{
+  "Verified": {
+    "Laguna-XS2":      {"lr": 0.00034219718825575905, "weight_decay": 0.0014714964923484015, "batch_size": 256,  "patience": 50},
+    "Qwen3.6-35B-A3B": {"lr": 0.00117,               "weight_decay": 0.0981,                "batch_size": 1024, "patience": 10}
+  },
+  "Pro": {
+    "Laguna-XS2": {
+      "Syntactic correctness": {"lr": 0.0002528899667354841,  "weight_decay": 0.00027521054719050283, "batch_size": 256,  "patience": 50},
+      "Semantic correctness":  {"lr": 0.0006160451703789833,  "weight_decay": 0.00517666984834032,    "batch_size": 1024, "patience": 25},
+      "Reduces failures":      {"lr": 0.0002310173345942377,  "weight_decay": 0.0040400920896719915,  "batch_size": 256,  "patience": 10},
+      "Has regressions":       {"lr": 0.00015055365890504895, "weight_decay": 0.0326146475479764,     "batch_size": 1024, "patience": 10}
+    },
+    "Qwen3.6-35B-A3B": null
+  }
+}

paper/hparams.json

@@ -0,0 +1,63 @@
+"""Visual style constants for paper figures.
+
+Adjust FULL_WIDTH for the target venue:
+  NeurIPS : 5.50 in
+  ICML    : 6.75 in
+  ICLR    : 6.00 in (approx)
+"""
+
+import matplotlib.pyplot as plt
+
+# --- Figure dimensions ---
+FULL_WIDTH      = 6.75   # inches; change per venue
+HALF_WIDTH      = 3.3
+FIG_HEIGHT_BAR  = 3.0
+FIG_HEIGHT_LINE = 2.6
+FIG_HEIGHT_HEAT = 2.8
+
+# --- Color palette ---
+# Model identity = color family; dataset identity = shade + hatch
+# solid fill = Verified,  hatched (///) = Pro
+COLORS = {
+    "laguna_verified": "#1565C0",   # dark blue
+    "laguna_pro":      "#64B5F6",   # light blue
+    "qwen_verified":   "#BF360C",   # dark coral
+    "qwen_pro":        "#FF8A65",   # light coral
+    "shuffled":        "#9E9E9E",   # neutral gray (both models)
+    "baseline":        "#AAAAAA",   # random-baseline dashed line
+}
+
+HATCH = {
+    "verified": "",     # solid fill
+    "pro":      "///",  # diagonal hatch
+}
+
+# --- Heatmap ---
+HEATMAP_CMAP = "Blues"
+HEATMAP_VMIN = 0.55
+HEATMAP_VMAX = 0.85
+
+
+def apply() -> None:
+    """Apply paper-wide rcParams. Call once at the top of main()."""
+    plt.rcParams.update({
+        "font.family":        "sans-serif",
+        "font.sans-serif":    ["Helvetica", "Arial", "DejaVu Sans"],
+        "font.size":          9,
+        "axes.titlesize":     11,
+        "axes.titleweight":   "normal",
+        "axes.labelsize":     10,
+        "xtick.labelsize":    9,
+        "ytick.labelsize":    9,
+        "legend.fontsize":    8,
+        "legend.framealpha":  0.9,
+        "figure.dpi":         150,
+        "savefig.dpi":        300,
+        "savefig.bbox":       "tight",
+        "axes.spines.top":    False,
+        "axes.spines.right":  False,
+        "axes.grid":          True,
+        "grid.color":         "#E0E0E0",
+        "grid.linewidth":     0.6,
+        "axes.axisbelow":     True,
+    })