Add semantic representation probes and QK similarity analysis

- Add semantic_probe.py for probing internal representations via cosine similarity to concept labels
- Add QK probe to measure query-key alignment across expert cross-attention layers
- Integrate representation probes as Phase 5/9 in the diagnostic pipeline
- Update DiagnosticPanel frontend to render semantic and QK probe results
- Enhance spatial object diagnosis with probe-informed evidence
- Update report rendering with representation probe sections
- Fix gradient attribution to support probe-compatible activation extraction

smolvla_inspect/diagnostic/agent.py

@@ -21,6 +21,14 @@ from .spatial_object import (
     choose_target_object,
     summarize_spatial_object_diagnosis,
 )
+from .semantic_probe import (
+    build_qk_probe_report,
+    build_semantic_probe_report,
+    probe_semantic_frame,
+    semantic_candidate_labels,
+    summarize_qk_probe,
+    summarize_semantic_probe,
+)
 from .prompts import (
     build_hypothesis_prompt, build_synthesis_prompt, build_triage_selection_prompt,
     format_diversity_summary, format_anomalies_json, format_hypotheses_with_results,
@@ -227,15 +235,16 @@ class DiagnosticAgent:
         """Run the full diagnostic pipeline."""
 
         _PHASE_LABELS = {
-            "scene_understanding": ("1/8", "Scene Understanding"),
-            "dataset_diversity":   ("2/8", "Dataset Diversity"),
-            "triage":              ("3/8", "Signal Triage"),
-            "matrix":              ("4/8", "Diagnostic Matrix"),
-            "hypothesize":         ("5/8", "Hypothesis Formation"),
-            "counterfactuals":     ("6/8", "Counterfactual Tests"),
-            "iteration":           ("6/8", "Follow-up Iteration"),
-            "disambiguation":      ("7/8", "Spatial vs Object Diagnosis"),
-            "synthesis":           ("8/8", "Report Synthesis"),
+            "scene_understanding": ("1/9", "Scene Understanding"),
+            "dataset_diversity":   ("2/9", "Dataset Diversity"),
+            "triage":              ("3/9", "Signal Triage"),
+            "matrix":              ("4/9", "Diagnostic Matrix"),
+            "representation":      ("5/9", "Representation Probes"),
+            "hypothesize":         ("6/9", "Hypothesis Formation"),
+            "counterfactuals":     ("7/9", "Counterfactual Tests"),
+            "iteration":           ("7/9", "Follow-up Iteration"),
+            "disambiguation":      ("8/9", "Spatial vs Object Diagnosis"),
+            "synthesis":           ("9/9", "Report Synthesis"),
             "complete":            ("OK", "Complete"),
         }
         _run_start = time.time()
@@ -390,11 +399,40 @@ class DiagnosticAgent:
             data={"anomaly_types": [a.type for a in anomalies]},
         ))
 
+        primary_semantic_frame = None
+        semantic_internal = None
+        semantic_probe = None
+        qk_probe = None
+        if self.policy is not None and target_object is not None:
+            _progress("representation", "Running semantic patch-to-text probe on the primary frame...")
+            try:
+                primary_semantic_frame, semantic_internal = probe_semantic_frame(
+                    self.policy,
+                    sample,
+                    self.dataset,
+                    self.image_key,
+                    self.device,
+                    target_object=target_object,
+                    candidate_labels=semantic_candidate_labels(scene, target_object),
+                    attention_heatmap=(signals.get("attention") or [None])[0],
+                    gradcam_map=(signals.get("gradcam_siglip") or [None])[0],
+                    image_map=self.image_map,
+                    frame_id="primary",
+                )
+                semantic_probe = build_semantic_probe_report(target_object, primary_semantic_frame, [])
+                if semantic_probe is not None:
+                    _progress("representation", summarize_semantic_probe(semantic_probe))
+            except Exception as e:
+                _progress("representation", f"Semantic probe skipped: {e}")
+        else:
+            _progress("representation", "Skipping semantic probe (no model or target object)")
+
         preliminary_spatial_object = build_spatial_object_diagnosis(
             matrix,
             target_object=target_object,
             cf_results=[],
             dataset_diversity=diversity,
+            semantic_probe=semantic_probe,
         )
 
         # ── Phase 5: LLM Hypothesis Formation ──────────────────
@@ -404,7 +442,7 @@ class DiagnosticAgent:
         _progress("hypothesize", f"Forming hypotheses via {llm_label}...")
         hypotheses = await self._form_hypotheses(
             task_string, scene, diversity, matrix, anomalies,
-            preliminary_spatial_object)
+            preliminary_spatial_object, semantic_probe, None)
         _progress("hypothesize", f"Formed {len(hypotheses)} hypotheses:")
         for h in hypotheses:
             test_label = f" -> test: {h.test_type}" if h.test_type != "none" else " (no test needed)"
@@ -472,7 +510,11 @@ class DiagnosticAgent:
             if surprising:
                 _progress("iteration", f"Found {len(surprising)} surprising results, forming follow-up hypotheses...")
                 followup_hypotheses = await self._form_hypotheses(
-                    task_string, scene, diversity, matrix, anomalies)
+                    task_string, scene, diversity, matrix, anomalies,
+                    pre_qk_spatial_object if 'pre_qk_spatial_object' in locals() else preliminary_spatial_object,
+                    semantic_probe,
+                    qk_probe,
+                )
                 # Filter to only genuinely new hypotheses
                 existing_ids = {h.id for h in hypotheses}
                 existing_tests = {(h.test_type, str(h.test_params)) for h in hypotheses}
@@ -552,24 +594,78 @@ class DiagnosticAgent:
                     except Exception as e:
                         _progress("counterfactuals", f"  -> FAILED: {e}")
 
-        # ── Phase 7: Spatial vs Object Diagnosis ─────────────────
+        semantic_probe = build_semantic_probe_report(target_object, primary_semantic_frame, cf_results)
+
+        # ── Phase 8: Spatial vs Object Diagnosis ─────────────────
         _progress("disambiguation", "Summarizing spatial priors vs object grounding...")
+        pre_qk_spatial_object = build_spatial_object_diagnosis(
+            matrix,
+            target_object=target_object,
+            cf_results=cf_results,
+            dataset_diversity=diversity,
+            semantic_probe=semantic_probe,
+        )
+        semantic_ambiguous = (
+            primary_semantic_frame is not None
+            and (
+                primary_semantic_frame.target_margin_over_best_non_target is None
+                or abs(primary_semantic_frame.target_margin_over_best_non_target) < 0.08
+            )
+        )
+        if (
+            self.policy is not None
+            and target_object is not None
+            and semantic_internal is not None
+            and (pre_qk_spatial_object.verdict in {"mixed", "inconclusive"} or semantic_ambiguous)
+        ):
+            _progress("disambiguation", "Running last-layer SigLIP QK decomposition...")
+            try:
+                relocation_result = next(
+                    (
+                        result for result in cf_results
+                        if result.test_type == "object_relocation"
+                        and (result.metrics or {}).get("target_object") == target_object
+                    ),
+                    None,
+                )
+                qk_probe = build_qk_probe_report(
+                    self.policy,
+                    sample,
+                    self.dataset,
+                    self.image_key,
+                    self.device,
+                    target_object=target_object,
+                    scene=scene,
+                    target_semantic_map=semantic_internal["target_map"],
+                    positional_baseline=signals.get("positional_baseline"),
+                    relocation_result=relocation_result,
+                    image_map=self.image_map,
+                )
+                if qk_probe is not None:
+                    _progress("disambiguation", summarize_qk_probe(qk_probe))
+            except Exception as e:
+                _progress("disambiguation", f"QK probe skipped: {e}")
+
         spatial_object_diagnosis = build_spatial_object_diagnosis(
             matrix,
             target_object=target_object,
             cf_results=cf_results,
             dataset_diversity=diversity,
+            semantic_probe=semantic_probe,
+            qk_probe=qk_probe,
         )
         _progress(
             "disambiguation",
             summarize_spatial_object_diagnosis(spatial_object_diagnosis),
         )
 
-        # ── Phase 8: LLM Synthesis ─────────────────────────────
+        # ── Phase 9: LLM Synthesis ─────────────────────────────
         _progress("synthesis", f"Synthesizing report via {llm_label}...")
         findings, narrative = await self._synthesize_report(
             task_string, scene, diversity, matrix, anomalies,
             hypotheses, cf_results, spatial_object_diagnosis,
+            semantic_probe,
+            qk_probe,
             cf_skip_reason=cf_skip_reason)
         _progress("synthesis", f"Generated {len(findings)} findings:")
         for f in findings:
@@ -608,6 +704,8 @@ class DiagnosticAgent:
             scene=scene,
             dataset_diversity=diversity,
             matrix=matrix,
+            semantic_probe=semantic_probe,
+            qk_probe=qk_probe,
             spatial_object_diagnosis=spatial_object_diagnosis,
             anomalies=anomalies,
             hypotheses=hypotheses,
@@ -1121,7 +1219,9 @@ class DiagnosticAgent:
     async def _form_hypotheses(self, task_string: str, scene: SceneSegmentation,
                                 diversity, matrix: DiagnosticMatrix,
                                 anomalies: list[Anomaly],
-                                spatial_object_diagnosis=None) -> list[Hypothesis]:
+                                spatial_object_diagnosis=None,
+                                semantic_probe=None,
+                                qk_probe=None) -> list[Hypothesis]:
         """Use LLM to form hypotheses from the diagnostic data."""
         prompt = build_hypothesis_prompt(
             task_string=task_string,
@@ -1130,6 +1230,8 @@ class DiagnosticAgent:
             matrix_markdown=matrix.to_markdown(),
             anomalies_json=format_anomalies_json(anomalies),
             spatial_object_summary=summarize_spatial_object_diagnosis(spatial_object_diagnosis),
+            semantic_probe_summary=summarize_semantic_probe(semantic_probe),
+            qk_probe_summary=summarize_qk_probe(qk_probe),
         )
 
         response = await self._call_llm(prompt)
@@ -1293,6 +1395,8 @@ class DiagnosticAgent:
                                   hypotheses: list[Hypothesis],
                                   cf_results: list[CounterfactualResult],
                                   spatial_object_diagnosis=None,
+                                  semantic_probe=None,
+                                  qk_probe=None,
                                   cf_skip_reason: str = "",
                                   ) -> tuple[list[Finding], str]:
         """Use LLM to synthesize findings and narrative from all evidence."""
@@ -1301,6 +1405,8 @@ class DiagnosticAgent:
             detected_objects=scene.region_names(),
             matrix_markdown=matrix.to_markdown(),
             spatial_object_summary=summarize_spatial_object_diagnosis(spatial_object_diagnosis),
+            semantic_probe_summary=summarize_semantic_probe(semantic_probe),
+            qk_probe_summary=summarize_qk_probe(qk_probe),
             anomalies_summary=format_anomalies_json(anomalies),
             diversity_summary=format_diversity_summary(diversity),
             hypotheses_with_results=format_hypotheses_with_results(

smolvla_inspect/diagnostic/counterfactual.py

@@ -22,6 +22,7 @@ from ..data import build_policy_batch_from_sample
 from ..gradient import _patch_eager_attention_bool_mask
 from .models import CounterfactualResult, SceneSegmentation
 from .registry import register_primitive
+from .semantic_probe import compute_patch_text_similarity
 
 
 # ---------------------------------------------------------------------------
@@ -255,6 +256,38 @@ def _mask_share(heatmap: np.ndarray | None, mask: np.ndarray) -> float:
     return float(np.sum(hm * mask.astype(np.float32)) / total)
 
 
+def _normalize_similarity_map(similarity_map: np.ndarray, valid_mask: np.ndarray) -> np.ndarray:
+    out = np.zeros_like(similarity_map, dtype=np.float32)
+    if not valid_mask.any():
+        return out
+    valid = similarity_map[valid_mask].astype(np.float32)
+    vmin = float(valid.min())
+    vmax = float(valid.max())
+    if vmax > vmin:
+        out[valid_mask] = (valid - vmin) / (vmax - vmin)
+    return out
+
+
+def _semantic_region_mean(similarity_map: np.ndarray | None, region_mask: np.ndarray,
+                          grid_shape: tuple[int, int], valid_mask: np.ndarray) -> float | None:
+    if similarity_map is None:
+        return None
+    region_grid = _resize_mask(region_mask, grid_shape) & valid_mask
+    if not region_grid.any():
+        return None
+    return float(similarity_map[region_grid].mean())
+
+
+def _semantic_region_peak(similarity_map: np.ndarray | None, region_mask: np.ndarray,
+                          grid_shape: tuple[int, int], valid_mask: np.ndarray) -> float | None:
+    if similarity_map is None:
+        return None
+    region_grid = _resize_mask(region_mask, grid_shape) & valid_mask
+    if not region_grid.any():
+        return None
+    return float(similarity_map[region_grid].max())
+
+
 def _shift_mask(mask: np.ndarray, shift_pixels: tuple[int, int]) -> np.ndarray:
     """Translate a binary mask by ``(dx, dy)`` with clipping."""
     dx, dy = shift_pixels
@@ -588,6 +621,40 @@ def object_relocation(
         "focus_shift_gap": moved_target_share - old_anchor_share,
     }
 
+    original_semantics = compute_patch_text_similarity(
+        policy, sample, dataset, image_key, device, [target_object], image_map=image_map,
+    )
+    modified_semantics = compute_patch_text_similarity(
+        policy, new_sample, dataset, image_key, device, [target_object], image_map=image_map,
+    )
+    if original_semantics is not None and modified_semantics is not None:
+        orig_map = original_semantics["similarity_maps"].get(target_object)
+        mod_map = modified_semantics["similarity_maps"].get(target_object)
+        if orig_map is not None and mod_map is not None:
+            orig_valid = original_semantics["valid_patch_mask"]
+            mod_valid = modified_semantics["valid_patch_mask"]
+            orig_grid = original_semantics["grid_size"]
+            mod_grid = modified_semantics["grid_size"]
+            orig_peak = _semantic_region_peak(orig_map, obj_mask, orig_grid, orig_valid)
+            old_anchor_sem = _semantic_region_mean(mod_map, old_anchor_mask, mod_grid, mod_valid)
+            moved_target_sem = _semantic_region_mean(mod_map, moved_target_mask, mod_grid, mod_valid)
+            mod_norm = _normalize_similarity_map(mod_map, mod_valid)
+            old_anchor_norm = _semantic_region_mean(mod_norm, old_anchor_mask, mod_grid, mod_valid) or 0.0
+            moved_target_norm = _semantic_region_mean(mod_norm, moved_target_mask, mod_grid, mod_valid) or 0.0
+            semantic_total = old_anchor_norm + moved_target_norm + 1e-8
+            metrics.update({
+                "original_target_semantic_peak": orig_peak if orig_peak is not None else 0.0,
+                "modified_old_anchor_target_semantic": (
+                    old_anchor_sem if old_anchor_sem is not None else 0.0
+                ),
+                "modified_new_object_target_semantic": (
+                    moved_target_sem if moved_target_sem is not None else 0.0
+                ),
+                "semantic_follow_ratio": moved_target_norm / semantic_total,
+                "semantic_anchor_ratio": old_anchor_norm / semantic_total,
+                "semantic_shift_gap": moved_target_norm - old_anchor_norm,
+            })
+
     return _compute_result(
         baseline_actions, modified_actions,
         img_hwc, modified_hwc,
@@ -1144,11 +1211,35 @@ def occlusion_targeted(
     policy.reset()
     modified_actions = _get_actions(policy, new_sample, dataset, image_key, device, image_map)
 
+    metrics = {"target_object": target_object, "fill": fill}
+    original_semantics = compute_patch_text_similarity(
+        policy, sample, dataset, image_key, device, [target_object], image_map=image_map,
+    )
+    occluded_semantics = compute_patch_text_similarity(
+        policy, new_sample, dataset, image_key, device, [target_object], image_map=image_map,
+    )
+    if original_semantics is not None and occluded_semantics is not None:
+        orig_map = original_semantics["similarity_maps"].get(target_object)
+        occ_map = occluded_semantics["similarity_maps"].get(target_object)
+        if orig_map is not None and occ_map is not None:
+            orig_valid = original_semantics["valid_patch_mask"]
+            occ_valid = occluded_semantics["valid_patch_mask"]
+            orig_grid = original_semantics["grid_size"]
+            occ_grid = occluded_semantics["grid_size"]
+            orig_peak = _semantic_region_peak(orig_map, obj_mask, orig_grid, orig_valid)
+            occ_peak = _semantic_region_peak(occ_map, obj_mask, occ_grid, occ_valid)
+            if orig_peak is not None and occ_peak is not None:
+                metrics.update({
+                    "original_target_semantic_peak": orig_peak,
+                    "occluded_target_semantic_peak": occ_peak,
+                    "occlusion_target_semantic_drop": orig_peak - occ_peak,
+                })
+
     return _compute_result(
         baseline_actions, modified_actions,
         img_hwc, modified_hwc,
         hypothesis_id=f"occlusion_targeted_{target_object}",
         test_type="occlusion_targeted",
         affected_mask=obj_mask,
-        metrics={"target_object": target_object, "fill": fill},
+        metrics=metrics,
     )

smolvla_inspect/diagnostic/models.py

@@ -329,6 +329,70 @@ class SpatialObjectDiagnosis:
     object_evidence: list[DiagnosticEvidence] = field(default_factory=list)
 
 
+@dataclass
+class SemanticFrameSummary:
+    """Semantic-probe summary for a single image/frame."""
+
+    frame_id: str
+    target_object: str | None
+    candidate_labels: list[str] = field(default_factory=list)
+    best_non_target_label: str | None = None
+    target_semantic_peak: float | None = None
+    target_semantic_mean_on_causal_patches: float | None = None
+    target_margin_over_best_non_target: float | None = None
+    causal_semantic_alignment: float | None = None
+    background_semantic_gap: float | None = None
+    summary: str = ""
+
+
+@dataclass
+class SemanticCounterfactualSummary:
+    """Semantic summary derived from a counterfactual probe."""
+
+    test_type: str
+    summary: str
+    metrics: dict = field(default_factory=dict)
+
+
+@dataclass
+class SemanticProbeReport:
+    """Top-level semantic object-vs-location probe report."""
+
+    target_object: str | None
+    summary: str
+    primary_frame: SemanticFrameSummary | None = None
+    frames: list[SemanticFrameSummary] = field(default_factory=list)
+    counterfactuals: dict[str, SemanticCounterfactualSummary] = field(default_factory=dict)
+
+
+@dataclass
+class QKHeadSummary:
+    """Per-head QK decomposition summary for the final SigLIP layer."""
+
+    head_index: int
+    head_type: str
+    score: float
+    semantic_map_correlation: float
+    positional_baseline_correlation: float
+    target_region_logit_mass: float
+    background_logit_mass: float
+    old_anchor_logit_mass: float | None = None
+    moved_object_logit_mass: float | None = None
+
+
+@dataclass
+class QKProbeReport:
+    """Compact summary of last-layer SigLIP QK behavior."""
+
+    layer: str
+    summary: str
+    dominant_head_type: str
+    semantic_head_fraction: float
+    positional_head_fraction: float
+    mixed_head_fraction: float
+    top_heads: list[QKHeadSummary] = field(default_factory=list)
+
+
 @dataclass
 class EvidenceEntry:
     """A single piece of evidence collected during a diagnostic phase."""
@@ -380,6 +444,8 @@ class DiagnosticReport:
     scene: SceneSegmentation
     dataset_diversity: DatasetDiversityReport | None
     matrix: DiagnosticMatrix
+    semantic_probe: SemanticProbeReport | None
+    qk_probe: QKProbeReport | None
     spatial_object_diagnosis: SpatialObjectDiagnosis | None
     anomalies: list[Anomaly]
     hypotheses: list[Hypothesis]
@@ -439,6 +505,8 @@ class DiagnosticReport:
             "scene": self._serialisable_scene(),
             "dataset_diversity": asdict(self.dataset_diversity) if self.dataset_diversity else None,
             "matrix": self.matrix.to_dict(),
+            "semantic_probe": asdict(self.semantic_probe) if self.semantic_probe else None,
+            "qk_probe": asdict(self.qk_probe) if self.qk_probe else None,
             "spatial_object_diagnosis": (
                 asdict(self.spatial_object_diagnosis)
                 if self.spatial_object_diagnosis
@@ -665,6 +733,73 @@ class DiagnosticReport:
                     sections.append(f"- {item.summary} (score={item.score:.2f}, source={item.source})")
                 sections.append("")
 
+        # ── Semantic Probe ───────────────────────────────────────
+        if self.semantic_probe is not None:
+            probe = self.semantic_probe
+            sections.append("## Semantic Feature Probe")
+            sections.append("")
+            sections.append(probe.summary)
+            sections.append("")
+
+            if probe.primary_frame is not None:
+                frame = probe.primary_frame
+                sections.append("### Primary Frame")
+                sections.append("")
+                sections.append("| Metric | Value |")
+                sections.append("|---|---|")
+                sections.append(f"| Target object | {frame.target_object or 'N/A'} |")
+                sections.append(f"| Candidate labels | {', '.join(frame.candidate_labels)} |")
+                sections.append(f"| Best non-target label | {frame.best_non_target_label or 'N/A'} |")
+                for key in (
+                    "target_semantic_peak",
+                    "target_semantic_mean_on_causal_patches",
+                    "target_margin_over_best_non_target",
+                    "causal_semantic_alignment",
+                    "background_semantic_gap",
+                ):
+                    value = getattr(frame, key)
+                    rendered = f"{value:.4f}" if isinstance(value, float) else "N/A"
+                    sections.append(f"| {key.replace('_', ' ').title()} | {rendered} |")
+                if frame.summary:
+                    sections.append("")
+                    sections.append(frame.summary)
+                sections.append("")
+
+            if probe.counterfactuals:
+                sections.append("### Counterfactual Semantic Follow-Through")
+                sections.append("")
+                for name, summary in probe.counterfactuals.items():
+                    sections.append(f"- **{name.replace('_', ' ')}**: {summary.summary}")
+                sections.append("")
+
+        # ── QK Probe ─────────────────────────────────────────────
+        if self.qk_probe is not None:
+            probe = self.qk_probe
+            sections.append("## QK Decomposition")
+            sections.append("")
+            sections.append(probe.summary)
+            sections.append("")
+            sections.append("| Metric | Value |")
+            sections.append("|---|---|")
+            sections.append(f"| Layer | {probe.layer} |")
+            sections.append(f"| Dominant head type | {probe.dominant_head_type} |")
+            sections.append(f"| Semantic head fraction | {probe.semantic_head_fraction:.4f} |")
+            sections.append(f"| Positional head fraction | {probe.positional_head_fraction:.4f} |")
+            sections.append(f"| Mixed head fraction | {probe.mixed_head_fraction:.4f} |")
+            sections.append("")
+            if probe.top_heads:
+                sections.append("### Top Heads")
+                sections.append("")
+                sections.append("| Head | Type | Score | Semantic Corr | Positional Corr | Target Mass | Background Mass |")
+                sections.append("|---|---|---|---|---|---|---|")
+                for head in probe.top_heads:
+                    sections.append(
+                        f"| {head.head_index} | {head.head_type} | {head.score:.4f} | "
+                        f"{head.semantic_map_correlation:.4f} | {head.positional_baseline_correlation:.4f} | "
+                        f"{head.target_region_logit_mass:.4f} | {head.background_logit_mass:.4f} |"
+                    )
+                sections.append("")
+
         # ── Anomalies ─────────────────────────────────────────────
         sections.append("## Detected Anomalies")
         sections.append("")

smolvla_inspect/diagnostic/prompts.py

@@ -31,6 +31,12 @@ HYPOTHESIS_PROMPT = """You are an expert robotics ML researcher diagnosing a vis
 **Spatial vs Object Diagnosis:**
 {spatial_object_summary}
 
+**Semantic Probe:**
+{semantic_probe_summary}
+
+**QK Decomposition:**
+{qk_probe_summary}
+
 **Detected Anomalies:**
 {anomalies_json}
 
@@ -91,6 +97,12 @@ SYNTHESIS_PROMPT = """You are an expert robotics ML researcher writing a diagnos
 **Spatial vs Object Diagnosis:**
 {spatial_object_summary}
 
+**Semantic Probe:**
+{semantic_probe_summary}
+
+**QK Decomposition:**
+{qk_probe_summary}
+
 **Anomalies detected:** {anomalies_summary}
 
 **Dataset diversity:** {diversity_summary}
@@ -136,6 +148,8 @@ def build_hypothesis_prompt(task_string: str, detected_objects: list[str],
                             diversity_summary: str, matrix_markdown: str,
                             anomalies_json: str,
                             spatial_object_summary: str = "Spatial-vs-object diagnosis unavailable.",
+                            semantic_probe_summary: str = "Semantic patch-to-text probe unavailable.",
+                            qk_probe_summary: str = "QK decomposition unavailable.",
                             ) -> str:
     """Build the hypothesis formation prompt with all context filled in."""
     return HYPOTHESIS_PROMPT.format(
@@ -145,6 +159,8 @@ def build_hypothesis_prompt(task_string: str, detected_objects: list[str],
         diversity_summary=diversity_summary,
         matrix_markdown=matrix_markdown,
         spatial_object_summary=spatial_object_summary,
+        semantic_probe_summary=semantic_probe_summary,
+        qk_probe_summary=qk_probe_summary,
         anomalies_json=anomalies_json,
     )
 
@@ -154,6 +170,8 @@ def build_synthesis_prompt(task_string: str, detected_objects: list[str],
                            diversity_summary: str,
                            hypotheses_with_results: str,
                            spatial_object_summary: str = "Spatial-vs-object diagnosis unavailable.",
+                           semantic_probe_summary: str = "Semantic patch-to-text probe unavailable.",
+                           qk_probe_summary: str = "QK decomposition unavailable.",
                            ) -> str:
     """Build the synthesis/report prompt with all context filled in."""
     return SYNTHESIS_PROMPT.format(
@@ -162,6 +180,8 @@ def build_synthesis_prompt(task_string: str, detected_objects: list[str],
         detected_objects=", ".join(detected_objects),
         matrix_markdown=matrix_markdown,
         spatial_object_summary=spatial_object_summary,
+        semantic_probe_summary=semantic_probe_summary,
+        qk_probe_summary=qk_probe_summary,
         anomalies_summary=anomalies_summary,
         diversity_summary=diversity_summary,
         hypotheses_with_results=hypotheses_with_results,

smolvla_inspect/diagnostic/report.py

@@ -104,6 +104,23 @@ def save_diagnostic_report(report: DiagnosticReport, run_dir: str) -> str:
             "metrics": cf.metrics,
         })
 
+    if report.semantic_probe is not None:
+        evidence_data.append({
+            "phase": "representation",
+            "type": "semantic_probe",
+            "target_object": report.semantic_probe.target_object,
+            "summary": report.semantic_probe.summary,
+        })
+
+    if report.qk_probe is not None:
+        evidence_data.append({
+            "phase": "representation",
+            "type": "qk_probe",
+            "layer": report.qk_probe.layer,
+            "summary": report.qk_probe.summary,
+            "dominant_head_type": report.qk_probe.dominant_head_type,
+        })
+
     if report.spatial_object_diagnosis is not None:
         diag = report.spatial_object_diagnosis
         evidence_data.append({

smolvla_inspect/diagnostic/semantic_probe.py

@@ -0,0 +1,685 @@
+"""Semantic and QK probes for spatial-vs-object diagnostics."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from ..gradient import compute_gradcam_map, extract_siglip_last_layer_features
+from .models import (
+    CounterfactualResult,
+    QKHeadSummary,
+    QKProbeReport,
+    SceneSegmentation,
+    SemanticCounterfactualSummary,
+    SemanticFrameSummary,
+    SemanticProbeReport,
+)
+
+_DEFAULT_SIGLIP_MODEL = "google/siglip-base-patch16-512"
+_SIGLIP_TEXT_CACHE: dict[tuple[str, str], tuple[object, object]] = {}
+
+
+def semantic_candidate_labels(
+    scene: SceneSegmentation,
+    target_object: str | None = None,
+) -> list[str]:
+    """Collect unique non-gripper labels, keeping *target_object* first."""
+    labels = []
+    if target_object is not None:
+        labels.append(target_object)
+    for obj in scene.objects:
+        label = obj.label
+        if "gripper" in label.lower():
+            continue
+        if label not in labels:
+            labels.append(label)
+    return labels
+
+
+def _infer_siglip_model_id(policy) -> str:
+    """Best-effort guess of the matching SigLIP text tower."""
+    try:
+        from ..data import find_vision_encoder
+
+        vision_encoder = find_vision_encoder(policy)
+    except Exception:
+        vision_encoder = None
+
+    if vision_encoder is None:
+        return _DEFAULT_SIGLIP_MODEL
+
+    cfg = getattr(vision_encoder, "config", None)
+    image_size = getattr(cfg, "image_size", 512)
+    patch_size = getattr(cfg, "patch_size", 16)
+    hidden_size = getattr(cfg, "hidden_size", 768)
+    if hidden_size >= 1024:
+        return f"google/siglip-large-patch{patch_size}-{image_size}"
+    return f"google/siglip-base-patch{patch_size}-{image_size}"
+
+
+def _text_device(device: str | torch.device) -> str:
+    device_str = str(device)
+    if device_str.startswith(("cuda", "cpu")):
+        return device_str
+    return "cpu"
+
+
+def _get_siglip_text_encoder(model_id: str, device: str):
+    key = (model_id, device)
+    if key in _SIGLIP_TEXT_CACHE:
+        return _SIGLIP_TEXT_CACHE[key]
+
+    from transformers import AutoModel, AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModel.from_pretrained(model_id).to(device).eval()
+    _SIGLIP_TEXT_CACHE[key] = (tokenizer, model)
+    return tokenizer, model
+
+
+def _unique_labels(labels: list[str]) -> list[str]:
+    seen: set[str] = set()
+    result = []
+    for label in labels:
+        if not label or label in seen:
+            continue
+        seen.add(label)
+        result.append(label)
+    return result
+
+
+def _resize_map(arr: np.ndarray, target_shape: tuple[int, int]) -> np.ndarray:
+    if arr.shape == target_shape:
+        return arr.astype(np.float32, copy=False)
+    tensor = torch.from_numpy(arr.astype(np.float32)).unsqueeze(0).unsqueeze(0)
+    resized = torch.nn.functional.interpolate(
+        tensor, size=target_shape, mode="bilinear", align_corners=False,
+    )
+    return resized.squeeze(0).squeeze(0).cpu().numpy()
+
+
+def _resize_mask(mask: np.ndarray, target_shape: tuple[int, int]) -> np.ndarray:
+    if mask.shape == target_shape:
+        return mask.astype(bool, copy=False)
+    tensor = torch.from_numpy(mask.astype(np.float32)).unsqueeze(0).unsqueeze(0)
+    resized = torch.nn.functional.interpolate(
+        tensor, size=target_shape, mode="nearest",
+    )
+    return resized.squeeze(0).squeeze(0).cpu().numpy() > 0.5
+
+
+def _normalize_map(values: np.ndarray, valid_mask: np.ndarray) -> np.ndarray:
+    out = np.zeros_like(values, dtype=np.float32)
+    if not valid_mask.any():
+        return out
+    valid = values[valid_mask].astype(np.float32)
+    vmin = float(valid.min())
+    vmax = float(valid.max())
+    if vmax > vmin:
+        out[valid_mask] = (valid - vmin) / (vmax - vmin)
+    return out
+
+
+def _topk_mask(values: np.ndarray, valid_mask: np.ndarray, pct: float = 0.05) -> np.ndarray:
+    mask = np.zeros_like(values, dtype=bool)
+    if not valid_mask.any():
+        return mask
+    flat = values[valid_mask]
+    k = max(4, int(np.ceil(flat.size * pct)))
+    k = min(k, flat.size)
+    threshold = np.partition(flat, flat.size - k)[flat.size - k]
+    mask[valid_mask] = values[valid_mask] >= threshold
+    return mask
+
+
+def _masked_mean(values: np.ndarray, mask: np.ndarray) -> float | None:
+    region = mask.astype(bool)
+    if not region.any():
+        return None
+    return float(values[region].mean())
+
+
+def _masked_max(values: np.ndarray, mask: np.ndarray) -> float | None:
+    region = mask.astype(bool)
+    if not region.any():
+        return None
+    return float(values[region].max())
+
+
+def _weighted_mean(values: np.ndarray, weights: np.ndarray, valid_mask: np.ndarray) -> float | None:
+    weights = np.where(valid_mask, np.maximum(weights, 0.0), 0.0).astype(np.float32)
+    total = float(weights.sum())
+    if total <= 1e-8:
+        return None
+    return float((values * weights).sum() / total)
+
+
+def _pearson_corr(a: np.ndarray, b: np.ndarray, valid_mask: np.ndarray) -> float:
+    mask = valid_mask.astype(bool)
+    if mask.sum() < 4:
+        return 0.0
+    av = a[mask].astype(np.float32)
+    bv = b[mask].astype(np.float32)
+    av = av - av.mean()
+    bv = bv - bv.mean()
+    denom = float(np.linalg.norm(av) * np.linalg.norm(bv))
+    if denom <= 1e-8:
+        return 0.0
+    return float(np.dot(av, bv) / denom)
+
+
+def _map_mass(values: np.ndarray, region_mask: np.ndarray, valid_mask: np.ndarray) -> float:
+    if not valid_mask.any():
+        return 0.0
+    norm = _normalize_map(values, valid_mask)
+    total = float(norm[valid_mask].sum())
+    if total <= 1e-8:
+        return 0.0
+    return float(norm[region_mask & valid_mask].sum() / total)
+
+
+def _clone_sample(sample: dict, image_key: str) -> dict:
+    cloned = dict(sample)
+    cloned[image_key] = sample[image_key].clone()
+    return cloned
+
+
+def _tensor_to_hwc(tensor: torch.Tensor) -> np.ndarray:
+    arr = tensor.detach().cpu().numpy()
+    if arr.ndim == 3 and arr.shape[0] in (1, 3):
+        arr = np.transpose(arr, (1, 2, 0))
+    return arr
+
+
+def _hwc_to_tensor(arr: np.ndarray, device: str = "cpu") -> torch.Tensor:
+    if arr.ndim == 3 and arr.shape[2] in (1, 3):
+        arr = np.transpose(arr, (2, 0, 1))
+    return torch.from_numpy(arr.astype(np.float32)).to(device)
+
+
+def _shift_mask(mask: np.ndarray, shift_pixels: tuple[int, int]) -> np.ndarray:
+    dx, dy = shift_pixels
+    shifted = np.zeros_like(mask, dtype=bool)
+    ys, xs = np.where(mask)
+    if len(ys) == 0:
+        return shifted
+    ny = ys + dy
+    nx = xs + dx
+    valid = (ny >= 0) & (ny < mask.shape[0]) & (nx >= 0) & (nx < mask.shape[1])
+    shifted[ny[valid], nx[valid]] = True
+    return shifted
+
+
+def _make_relocated_sample(sample: dict, image_key: str, segmentation: SceneSegmentation,
+                           target_object: str, shift_pixels: tuple[int, int]) -> tuple[dict | None, dict]:
+    img_tensor = sample[image_key]
+    img_hwc = _tensor_to_hwc(img_tensor)
+    h, w = img_hwc.shape[:2]
+    obj_mask = segmentation.get_mask(target_object)
+    if obj_mask is None:
+        return None, {}
+    obj_mask = _resize_mask(obj_mask, (h, w))
+    dx, dy = shift_pixels
+    modified_hwc = img_hwc.copy()
+    ys, xs = np.where(obj_mask)
+    if len(ys) == 0:
+        return None, {}
+
+    fill_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
+    modified_hwc[obj_mask] = fill_color
+    for oy, ox in zip(ys, xs):
+        ny = oy + dy
+        nx = ox + dx
+        if 0 <= ny < h and 0 <= nx < w:
+            modified_hwc[ny, nx] = img_hwc[oy, ox]
+    modified_hwc = np.clip(modified_hwc, 0.0, 1.0)
+
+    moved_mask = _shift_mask(obj_mask, shift_pixels)
+    old_anchor_mask = obj_mask & ~moved_mask
+    moved_target_mask = moved_mask & ~obj_mask
+    if not old_anchor_mask.any():
+        old_anchor_mask = obj_mask
+    if not moved_target_mask.any():
+        moved_target_mask = moved_mask
+
+    new_sample = _clone_sample(sample, image_key)
+    new_sample[image_key] = _hwc_to_tensor(modified_hwc, device=str(img_tensor.device))
+    return new_sample, {
+        "original_mask": obj_mask,
+        "old_anchor_mask": old_anchor_mask,
+        "moved_target_mask": moved_target_mask,
+    }
+
+
+def compute_patch_text_similarity(
+    policy,
+    sample,
+    dataset,
+    image_key,
+    device,
+    labels: list[str],
+    *,
+    image_map=None,
+    capture_qk: bool = False,
+):
+    """Extract final-layer patch features and cosine maps for *labels*."""
+    labels = _unique_labels(labels)
+    if not labels:
+        return None
+
+    features = extract_siglip_last_layer_features(
+        policy, sample, dataset, image_key, device,
+        image_map=image_map, capture_qk=capture_qk,
+    )
+    if features is None:
+        return None
+
+    model_id = _infer_siglip_model_id(policy)
+    device_str = _text_device(device)
+    try:
+        tokenizer, text_model = _get_siglip_text_encoder(model_id, device_str)
+    except Exception as exc:
+        if model_id != _DEFAULT_SIGLIP_MODEL:
+            try:
+                model_id = _DEFAULT_SIGLIP_MODEL
+                tokenizer, text_model = _get_siglip_text_encoder(model_id, device_str)
+            except Exception:
+                print(f"    WARNING: SigLIP text encoder unavailable for semantic probe ({exc})")
+                return None
+        else:
+            print(f"    WARNING: SigLIP text encoder unavailable for semantic probe ({exc})")
+            return None
+
+    with torch.no_grad():
+        inputs = tokenizer(labels, padding=True, truncation=True, return_tensors="pt")
+        inputs = {key: value.to(device_str) for key, value in inputs.items()}
+        text_embeddings = text_model.get_text_features(**inputs).float()
+        text_embeddings = torch.nn.functional.normalize(text_embeddings, dim=-1)
+
+    patch_features = torch.from_numpy(features["patch_features"]).float()
+    if patch_features.shape[-1] != text_embeddings.shape[-1]:
+        print(
+            "    WARNING: SigLIP patch/text dimensions do not match for semantic probe "
+            f"({patch_features.shape[-1]} vs {text_embeddings.shape[-1]})"
+        )
+        return None
+    similarity_maps: dict[str, np.ndarray] = {}
+    for idx, label in enumerate(labels):
+        sim = torch.tensordot(patch_features, text_embeddings[idx].cpu(), dims=([-1], [0]))
+        similarity_maps[label] = sim.cpu().numpy()
+
+    result = dict(features)
+    result["labels"] = labels
+    result["model_id"] = model_id
+    result["similarity_maps"] = similarity_maps
+    return result
+
+
+def probe_semantic_frame(
+    policy,
+    sample,
+    dataset,
+    image_key,
+    device,
+    *,
+    target_object: str | None,
+    candidate_labels: list[str],
+    attention_heatmap: np.ndarray | None = None,
+    gradcam_map: np.ndarray | None = None,
+    image_map=None,
+    frame_id: str = "primary",
+):
+    """Build a semantic summary for one frame and return internal artifacts."""
+    if target_object is None:
+        return None, None
+
+    artifact = compute_patch_text_similarity(
+        policy, sample, dataset, image_key, device, candidate_labels,
+        image_map=image_map, capture_qk=False,
+    )
+    if artifact is None:
+        return None, None
+
+    similarity_maps = artifact["similarity_maps"]
+    valid_mask = artifact["valid_patch_mask"]
+    grid_shape = artifact["grid_size"]
+    target_map = similarity_maps.get(target_object)
+    if target_map is None:
+        return None, None
+
+    if gradcam_map is None:
+        try:
+            gradcam_map = compute_gradcam_map(
+                policy, sample, dataset, image_key, device, image_map=image_map,
+            )
+        except Exception:
+            gradcam_map = None
+
+    gradcam_grid = _resize_map(gradcam_map, grid_shape) if gradcam_map is not None else None
+    attention_grid = _resize_map(attention_heatmap, grid_shape) if attention_heatmap is not None else None
+
+    if gradcam_grid is not None:
+        gradcam_top = _topk_mask(gradcam_grid, valid_mask)
+    else:
+        gradcam_top = np.zeros(grid_shape, dtype=bool)
+    if attention_grid is not None:
+        attention_top = _topk_mask(attention_grid, valid_mask)
+    else:
+        attention_top = np.zeros(grid_shape, dtype=bool)
+
+    causal_mask = np.zeros(grid_shape, dtype=bool)
+    if gradcam_top.any() and attention_top.any():
+        causal_mask = gradcam_top & attention_top
+        if not causal_mask.any():
+            causal_mask = gradcam_top
+    elif gradcam_top.any():
+        causal_mask = gradcam_top
+    elif attention_top.any():
+        causal_mask = attention_top
+    else:
+        causal_mask = valid_mask.copy()
+
+    if gradcam_grid is not None:
+        causal_weights = _normalize_map(gradcam_grid, valid_mask)
+    elif attention_grid is not None:
+        causal_weights = _normalize_map(attention_grid, valid_mask)
+    else:
+        causal_weights = valid_mask.astype(np.float32)
+    causal_weights = np.where(causal_mask, causal_weights, 0.0)
+
+    target_peak = _masked_max(target_map, valid_mask)
+    target_mean_causal = _masked_mean(target_map, causal_mask & valid_mask)
+    causal_alignment = _weighted_mean(target_map, causal_weights, valid_mask)
+    background_mask = valid_mask & ~causal_mask
+    background_mean = _masked_mean(target_map, background_mask)
+    background_gap = None
+    if target_mean_causal is not None and background_mean is not None:
+        background_gap = float(target_mean_causal - background_mean)
+
+    non_target_scores = []
+    for label, sim_map in similarity_maps.items():
+        if label == target_object:
+            continue
+        score = _masked_mean(sim_map, causal_mask & valid_mask)
+        if score is not None:
+            non_target_scores.append((label, score))
+
+    best_non_target_label = None
+    target_margin = None
+    if target_mean_causal is not None and non_target_scores:
+        best_non_target_label, best_non_target_score = max(non_target_scores, key=lambda item: item[1])
+        target_margin = float(target_mean_causal - best_non_target_score)
+
+    summary_bits = []
+    if target_margin is not None:
+        if target_margin > 0.10:
+            summary_bits.append(
+                f"Causal patches align more strongly with '{target_object}' than other detected objects "
+                f"(margin={target_margin:.3f})."
+            )
+        elif target_margin < 0.0:
+            summary_bits.append(
+                f"Causal patches are not semantically distinctive for '{target_object}' "
+                f"(margin={target_margin:.3f})."
+            )
+    if background_gap is not None:
+        if background_gap > 0.10:
+            summary_bits.append(
+                f"Target semantics concentrate on causal patches rather than the low-causal background "
+                f"(gap={background_gap:.3f})."
+            )
+        elif background_gap < 0.0:
+            summary_bits.append(
+                f"Target semantics are no stronger on causal patches than on the low-causal background "
+                f"(gap={background_gap:.3f})."
+            )
+
+    frame = SemanticFrameSummary(
+        frame_id=frame_id,
+        target_object=target_object,
+        candidate_labels=candidate_labels,
+        best_non_target_label=best_non_target_label,
+        target_semantic_peak=target_peak,
+        target_semantic_mean_on_causal_patches=target_mean_causal,
+        target_margin_over_best_non_target=target_margin,
+        causal_semantic_alignment=causal_alignment,
+        background_semantic_gap=background_gap,
+        summary=" ".join(summary_bits) if summary_bits else "Semantic evidence was weak or unavailable on the causal patches.",
+    )
+    internal = {
+        "target_map": target_map,
+        "valid_patch_mask": valid_mask,
+        "causal_mask": causal_mask,
+        "causal_weights": causal_weights,
+        "similarity_maps": similarity_maps,
+        "artifact": artifact,
+    }
+    return frame, internal
+
+
+def summarize_semantic_probe(report: SemanticProbeReport | None) -> str:
+    """Short summary string for prompts and logs."""
+    if report is None or report.primary_frame is None:
+        return "Semantic patch-to-text probe unavailable."
+    frame = report.primary_frame
+    bits = [f"Target object: {report.target_object or 'N/A'}."]
+    if frame.target_margin_over_best_non_target is not None:
+        bits.append(f"Target margin over best non-target: {frame.target_margin_over_best_non_target:.3f}.")
+    if frame.background_semantic_gap is not None:
+        bits.append(f"Background gap: {frame.background_semantic_gap:.3f}.")
+    bits.append(frame.summary)
+    return " ".join(bits)
+
+
+def build_semantic_probe_report(
+    target_object: str | None,
+    primary_frame: SemanticFrameSummary | None,
+    cf_results: list[CounterfactualResult],
+) -> SemanticProbeReport | None:
+    """Assemble the primary-frame and counterfactual semantic summaries."""
+    if target_object is None and primary_frame is None:
+        return None
+
+    counterfactuals: dict[str, SemanticCounterfactualSummary] = {}
+    for result in cf_results:
+        metrics = result.metrics or {}
+        if result.test_type == "object_relocation" and "semantic_follow_ratio" in metrics:
+            follow = float(metrics.get("semantic_follow_ratio", 0.0))
+            anchor = float(metrics.get("semantic_anchor_ratio", 0.0))
+            summary = (
+                f"After relocation, target semantics follow the moved object "
+                f"(follow={follow:.2f}, anchor={anchor:.2f})."
+                if follow >= anchor
+                else f"After relocation, target semantics stay closer to the old anchor "
+                f"(anchor={anchor:.2f}, follow={follow:.2f})."
+            )
+            counterfactuals[result.test_type] = SemanticCounterfactualSummary(
+                test_type=result.test_type,
+                summary=summary,
+                metrics=metrics,
+            )
+        elif result.test_type == "occlusion_targeted" and "occlusion_target_semantic_drop" in metrics:
+            drop = float(metrics.get("occlusion_target_semantic_drop", 0.0))
+            summary = (
+                f"Occluding the target reduces target-semantic evidence in the target region "
+                f"(drop={drop:.3f})."
+                if drop > 0
+                else f"Occluding the target does not reduce target-semantic evidence "
+                f"(drop={drop:.3f})."
+            )
+            counterfactuals[result.test_type] = SemanticCounterfactualSummary(
+                test_type=result.test_type,
+                summary=summary,
+                metrics=metrics,
+            )
+
+    if primary_frame is not None:
+        summary = primary_frame.summary
+    elif counterfactuals:
+        summary = "Semantic evidence is available only through counterfactual follow-through probes."
+    else:
+        summary = "Semantic patch-to-text probe unavailable."
+
+    frames = [primary_frame] if primary_frame is not None else []
+    return SemanticProbeReport(
+        target_object=target_object,
+        summary=summary,
+        primary_frame=primary_frame,
+        frames=frames,
+        counterfactuals=counterfactuals,
+    )
+
+
+def build_qk_probe_report(
+    policy,
+    sample,
+    dataset,
+    image_key,
+    device,
+    *,
+    target_object: str | None,
+    scene: SceneSegmentation,
+    target_semantic_map: np.ndarray | None,
+    positional_baseline: np.ndarray | None = None,
+    relocation_result: CounterfactualResult | None = None,
+    image_map=None,
+) -> QKProbeReport | None:
+    """Summarize last-layer SigLIP QK behavior with semantic-vs-positional evidence."""
+    if target_object is None or target_semantic_map is None:
+        return None
+
+    artifact = compute_patch_text_similarity(
+        policy, sample, dataset, image_key, device, [target_object],
+        image_map=image_map, capture_qk=True,
+    )
+    if artifact is None or "qk_key_maps" not in artifact:
+        return None
+
+    valid_mask = artifact["valid_patch_mask"]
+    grid_shape = artifact["grid_size"]
+    qk_maps = artifact["qk_key_maps"]
+    target_map = _resize_map(target_semantic_map, grid_shape)
+    positional_grid = (
+        _resize_map(positional_baseline, grid_shape)
+        if positional_baseline is not None
+        else np.zeros(grid_shape, dtype=np.float32)
+    )
+
+    target_mask = scene.get_mask(target_object)
+    target_grid_mask = (
+        _resize_mask(target_mask, grid_shape) & valid_mask
+        if target_mask is not None
+        else np.zeros(grid_shape, dtype=bool)
+    )
+    background_grid_mask = _resize_mask(scene.background_mask, grid_shape) & valid_mask
+
+    relocation_old_mask = None
+    relocation_new_mask = None
+    relocated_qk_maps = None
+    if relocation_result is not None:
+        metrics = relocation_result.metrics or {}
+        shift = metrics.get("shift_pixels")
+        if isinstance(shift, (list, tuple)) and len(shift) == 2:
+            relocated_sample, relocated_masks = _make_relocated_sample(
+                sample, image_key, scene, target_object, (int(shift[0]), int(shift[1])),
+            )
+            if relocated_sample is not None:
+                relocated_artifact = compute_patch_text_similarity(
+                    policy, relocated_sample, dataset, image_key, device, [target_object],
+                    image_map=image_map, capture_qk=True,
+                )
+                if relocated_artifact is not None and "qk_key_maps" in relocated_artifact:
+                    relocated_qk_maps = relocated_artifact["qk_key_maps"]
+                    relocation_old_mask = _resize_mask(relocated_masks["old_anchor_mask"], grid_shape) & valid_mask
+                    relocation_new_mask = _resize_mask(relocated_masks["moved_target_mask"], grid_shape) & valid_mask
+
+    head_summaries: list[QKHeadSummary] = []
+    for head_index, head_map in enumerate(qk_maps):
+        semantic_corr = _pearson_corr(head_map, target_map, valid_mask)
+        positional_corr = _pearson_corr(head_map, positional_grid, valid_mask)
+        target_mass = _map_mass(head_map, target_grid_mask, valid_mask)
+        background_mass = _map_mass(head_map, background_grid_mask, valid_mask)
+
+        old_anchor_mass = None
+        moved_object_mass = None
+        relocation_bias = 0.0
+        if relocated_qk_maps is not None and relocation_old_mask is not None and relocation_new_mask is not None:
+            old_anchor_mass = _map_mass(relocated_qk_maps[head_index], relocation_old_mask, valid_mask)
+            moved_object_mass = _map_mass(relocated_qk_maps[head_index], relocation_new_mask, valid_mask)
+            relocation_bias = moved_object_mass - old_anchor_mass
+
+        semantic_adv = semantic_corr - positional_corr
+        if relocation_bias > 0.05:
+            semantic_adv += 0.10
+        elif relocation_bias < -0.05:
+            semantic_adv -= 0.10
+
+        if semantic_adv > 0.08:
+            head_type = "semantic"
+        elif semantic_adv < -0.08:
+            head_type = "positional"
+        else:
+            head_type = "mixed"
+
+        score = max(abs(semantic_adv), abs(relocation_bias), abs(target_mass - background_mass))
+        head_summaries.append(QKHeadSummary(
+            head_index=head_index,
+            head_type=head_type,
+            score=score,
+            semantic_map_correlation=semantic_corr,
+            positional_baseline_correlation=positional_corr,
+            target_region_logit_mass=target_mass,
+            background_logit_mass=background_mass,
+            old_anchor_logit_mass=old_anchor_mass,
+            moved_object_logit_mass=moved_object_mass,
+        ))
+
+    if not head_summaries:
+        return None
+
+    semantic_count = sum(1 for head in head_summaries if head.head_type == "semantic")
+    positional_count = sum(1 for head in head_summaries if head.head_type == "positional")
+    mixed_count = sum(1 for head in head_summaries if head.head_type == "mixed")
+    total = len(head_summaries)
+
+    semantic_fraction = semantic_count / total
+    positional_fraction = positional_count / total
+    mixed_fraction = mixed_count / total
+
+    if semantic_fraction > positional_fraction + 0.15:
+        dominant = "semantic"
+        summary = (
+            f"Most last-layer SigLIP heads align more with target semantics than with the positional baseline "
+            f"({semantic_fraction:.0%} semantic vs {positional_fraction:.0%} positional)."
+        )
+    elif positional_fraction > semantic_fraction + 0.15:
+        dominant = "positional"
+        summary = (
+            f"Most last-layer SigLIP heads align more with the positional baseline than with target semantics "
+            f"({positional_fraction:.0%} positional vs {semantic_fraction:.0%} semantic)."
+        )
+    else:
+        dominant = "mixed"
+        summary = (
+            f"Last-layer SigLIP heads split between semantic and positional signals "
+            f"({semantic_fraction:.0%} semantic, {positional_fraction:.0%} positional, {mixed_fraction:.0%} mixed)."
+        )
+
+    top_heads = sorted(head_summaries, key=lambda head: head.score, reverse=True)[:3]
+    return QKProbeReport(
+        layer="siglip_last",
+        summary=summary,
+        dominant_head_type=dominant,
+        semantic_head_fraction=semantic_fraction,
+        positional_head_fraction=positional_fraction,
+        mixed_head_fraction=mixed_fraction,
+        top_heads=top_heads,
+    )
+
+
+def summarize_qk_probe(report: QKProbeReport | None) -> str:
+    """Short summary string for prompts and logs."""
+    if report is None:
+        return "QK decomposition unavailable."
+    return report.summary

smolvla_inspect/diagnostic/spatial_object.py

@@ -7,7 +7,9 @@ from .models import (
     DatasetDiversityReport,
     DiagnosticEvidence,
     DiagnosticMatrix,
+    QKProbeReport,
     SceneSegmentation,
+    SemanticProbeReport,
     SpatialObjectDiagnosis,
 )
 
@@ -100,6 +102,8 @@ def build_spatial_object_diagnosis(
     target_object: str | None,
     cf_results: list[CounterfactualResult],
     dataset_diversity: DatasetDiversityReport | None = None,
+    semantic_probe: SemanticProbeReport | None = None,
+    qk_probe: QKProbeReport | None = None,
 ) -> SpatialObjectDiagnosis:
     """Summarize whether the policy is reading a location or an object."""
     spatial_evidence: list[DiagnosticEvidence] = []
@@ -182,6 +186,72 @@ def build_spatial_object_diagnosis(
                 {"target_object": target_object, "target_attention_share": attention_share},
             )
 
+    if semantic_probe is not None and semantic_probe.primary_frame is not None:
+        frame = semantic_probe.primary_frame
+        if frame.target_semantic_peak is not None:
+            key_metrics["target_semantic_peak"] = float(frame.target_semantic_peak)
+        if frame.target_semantic_mean_on_causal_patches is not None:
+            key_metrics["target_semantic_mean_on_causal_patches"] = float(
+                frame.target_semantic_mean_on_causal_patches
+            )
+        if frame.target_margin_over_best_non_target is not None:
+            margin = float(frame.target_margin_over_best_non_target)
+            key_metrics["target_margin_over_best_non_target"] = margin
+            add_object(
+                "semantic_target_margin",
+                _scale(margin, 0.02, 0.18),
+                0.18,
+                f"Causal patches look more like '{target_object}' than other detected objects "
+                f"(margin={margin:.3f}).",
+                {
+                    "target_object": target_object,
+                    "best_non_target_label": frame.best_non_target_label,
+                    "target_margin_over_best_non_target": margin,
+                },
+            )
+            add_spatial(
+                "weak_semantic_target_margin",
+                _scale(0.05 - margin, 0.01, 0.10),
+                0.14,
+                f"Causal patches are not semantically distinctive for '{target_object}' "
+                f"(margin={margin:.3f}).",
+                {
+                    "target_object": target_object,
+                    "best_non_target_label": frame.best_non_target_label,
+                    "target_margin_over_best_non_target": margin,
+                },
+            )
+        if frame.causal_semantic_alignment is not None:
+            alignment = float(frame.causal_semantic_alignment)
+            key_metrics["causal_semantic_alignment"] = alignment
+            add_object(
+                "causal_semantic_alignment",
+                _scale(alignment, 0.15, 0.45),
+                0.10,
+                f"High-causal patches contain clear target semantics for '{target_object}' "
+                f"(alignment={alignment:.3f}).",
+                {"target_object": target_object, "causal_semantic_alignment": alignment},
+            )
+        if frame.background_semantic_gap is not None:
+            bg_gap = float(frame.background_semantic_gap)
+            key_metrics["background_semantic_gap"] = bg_gap
+            add_object(
+                "background_semantic_gap",
+                _scale(bg_gap, 0.03, 0.18),
+                0.10,
+                f"Target semantics are stronger on high-causal patches than on the low-causal background "
+                f"(gap={bg_gap:.3f}).",
+                {"target_object": target_object, "background_semantic_gap": bg_gap},
+            )
+            add_spatial(
+                "low_background_semantic_gap",
+                _scale(0.04 - bg_gap, 0.01, 0.08),
+                0.10,
+                f"Target semantics are no stronger on causal patches than on the surrounding background "
+                f"(gap={bg_gap:.3f}).",
+                {"target_object": target_object, "background_semantic_gap": bg_gap},
+            )
+
     for traj in matrix.temporal_trajectories or []:
         corr = float(traj.object_tracking_correlation)
         key_metrics["object_tracking_correlation"] = max(
@@ -228,10 +298,16 @@ def build_spatial_object_diagnosis(
         metrics = relocation.metrics or {}
         follow_ratio = metrics.get("focus_follow_ratio")
         anchor_ratio = metrics.get("anchor_retention_ratio")
+        semantic_follow_ratio = metrics.get("semantic_follow_ratio")
+        semantic_anchor_ratio = metrics.get("semantic_anchor_ratio")
         if follow_ratio is not None:
             key_metrics["relocation_follow_ratio"] = float(follow_ratio)
         if anchor_ratio is not None:
             key_metrics["relocation_anchor_ratio"] = float(anchor_ratio)
+        if semantic_follow_ratio is not None:
+            key_metrics["relocation_semantic_follow_ratio"] = float(semantic_follow_ratio)
+        if semantic_anchor_ratio is not None:
+            key_metrics["relocation_semantic_anchor_ratio"] = float(semantic_anchor_ratio)
 
         if follow_ratio is not None and anchor_ratio is not None:
             add_object(
@@ -248,6 +324,23 @@ def build_spatial_object_diagnosis(
                 f"After relocation, GradCAM stays anchored to the old location more than the moved object (anchor={float(anchor_ratio):.2f}, follow={float(follow_ratio):.2f}).",
                 metrics,
             )
+        if semantic_follow_ratio is not None and semantic_anchor_ratio is not None:
+            add_object(
+                "semantic_relocation_follow_probe",
+                _scale(float(semantic_follow_ratio) - float(semantic_anchor_ratio), 0.05, 0.35),
+                0.20,
+                f"After relocation, target semantics follow the moved object more than the old anchor "
+                f"(follow={float(semantic_follow_ratio):.2f}, anchor={float(semantic_anchor_ratio):.2f}).",
+                metrics,
+            )
+            add_spatial(
+                "semantic_relocation_anchor_probe",
+                _scale(float(semantic_anchor_ratio) - float(semantic_follow_ratio), 0.05, 0.35),
+                0.20,
+                f"After relocation, target semantics stay more anchored to the old location than to the moved object "
+                f"(anchor={float(semantic_anchor_ratio):.2f}, follow={float(semantic_follow_ratio):.2f}).",
+                metrics,
+            )
 
     occlusion = _pick_counterfactual(
         cf_results,
@@ -257,6 +350,7 @@ def build_spatial_object_diagnosis(
     if occlusion is not None:
         delta = float(occlusion.action_delta_l2)
         key_metrics["occlusion_action_delta"] = delta
+        semantic_drop = occlusion.metrics.get("occlusion_target_semantic_drop") if occlusion.metrics else None
         add_object(
             "target_occlusion_response",
             _scale(delta, 0.02, 0.10),
@@ -271,6 +365,25 @@ def build_spatial_object_diagnosis(
             f"Occluding '{target_object}' barely changes the predicted action (delta={delta:.4f}).",
             {"target_object": target_object, "action_delta_l2": delta},
         )
+        if semantic_drop is not None:
+            semantic_drop = float(semantic_drop)
+            key_metrics["occlusion_target_semantic_drop"] = semantic_drop
+            add_object(
+                "target_occlusion_semantic_drop",
+                _scale(semantic_drop, 0.03, 0.20),
+                0.12,
+                f"Occluding '{target_object}' removes target-semantic evidence from the target region "
+                f"(drop={semantic_drop:.3f}).",
+                {"target_object": target_object, "occlusion_target_semantic_drop": semantic_drop},
+            )
+            add_spatial(
+                "target_occlusion_semantic_insensitivity",
+                _scale(0.02 - semantic_drop, 0.005, 0.02),
+                0.08,
+                f"Occluding '{target_object}' does not materially reduce target-semantic evidence "
+                f"(drop={semantic_drop:.3f}).",
+                {"target_object": target_object, "occlusion_target_semantic_drop": semantic_drop},
+            )
 
     background = _pick_counterfactual(cf_results, "background_substitution")
     if background is not None:
@@ -300,6 +413,34 @@ def build_spatial_object_diagnosis(
             {"target_object": target_object, "action_delta_l2": delta},
         )
 
+    if qk_probe is not None:
+        key_metrics["qk_semantic_head_fraction"] = qk_probe.semantic_head_fraction
+        key_metrics["qk_positional_head_fraction"] = qk_probe.positional_head_fraction
+        if qk_probe.dominant_head_type == "semantic":
+            add_object(
+                "qk_semantic_heads",
+                _scale(
+                    qk_probe.semantic_head_fraction - qk_probe.positional_head_fraction,
+                    0.05,
+                    0.40,
+                ),
+                0.14,
+                qk_probe.summary,
+                {"dominant_head_type": qk_probe.dominant_head_type},
+            )
+        elif qk_probe.dominant_head_type == "positional":
+            add_spatial(
+                "qk_positional_heads",
+                _scale(
+                    qk_probe.positional_head_fraction - qk_probe.semantic_head_fraction,
+                    0.05,
+                    0.40,
+                ),
+                0.14,
+                qk_probe.summary,
+                {"dominant_head_type": qk_probe.dominant_head_type},
+            )
+
     spatial_score = spatial_sum / spatial_weight if spatial_weight > 1e-8 else 0.0
     object_score = object_sum / object_weight if object_weight > 1e-8 else 0.0
 
@@ -340,7 +481,7 @@ def build_spatial_object_diagnosis(
         summary = (
             f"The available evidence is not strong enough to cleanly separate spatial-prior use from "
             f"object grounding for '{target_object}'. The model may be weakly grounded, weakly shortcut-driven, "
-            "or the available probes may be too noisy."
+                "or the available probes may be too noisy."
         )
 
     return SpatialObjectDiagnosis(

smolvla_inspect/gradient.py

@@ -11,6 +11,7 @@ import warnings
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 
 try:
     from tqdm import tqdm
@@ -25,6 +26,8 @@ from .data import (
     _resolve_task_string,
     find_vision_encoder,
 )
+from ._compat import resize_with_pad
+from .heatmap import compute_padding_patches
 
 
 # ---------------------------------------------------------------------------
@@ -234,6 +237,196 @@ def _gradcam_from_captured(activations, gradients, spatial_shape=None):
     return cam.detach().float().cpu().numpy()
 
 
+def _prepare_siglip_vision_inputs(sample, image_key, vision_encoder, device):
+    """Normalize a sample image for direct SigLIP vision-encoder forward passes."""
+    img = sample[image_key].unsqueeze(0).to(device)
+    if img.max() > 1.0:
+        img = img.float() / 255.0
+
+    original_hw = (int(img.shape[-2]), int(img.shape[-1]))
+    target_size = getattr(
+        getattr(vision_encoder, "config", None), "image_size", None,
+    ) or getattr(vision_encoder, "image_size", 384)
+    if img.shape[-2] != target_size or img.shape[-1] != target_size:
+        img_resized = resize_with_pad(img, target_size, target_size, pad_value=0)
+    else:
+        img_resized = img
+    img_resized = img_resized * 2.0 - 1.0
+
+    try:
+        enc_dtype = next(vision_encoder.parameters()).dtype
+        img_resized = img_resized.to(enc_dtype)
+    except StopIteration:
+        pass
+
+    patch_size = getattr(vision_encoder, "patch_size", None) or getattr(
+        getattr(vision_encoder, "config", None), "patch_size", 14,
+    )
+    grid_h = int(img_resized.size(2) // patch_size)
+    grid_w = int(img_resized.size(3) // patch_size)
+    patch_mask = torch.ones(1, grid_h, grid_w, dtype=torch.bool, device=device)
+
+    crop_h, crop_w = compute_padding_patches(original_hw, target_size, patch_size)
+    valid_patch_mask = np.ones((grid_h, grid_w), dtype=bool)
+    if crop_h > 0:
+        valid_patch_mask[:crop_h, :] = False
+    if crop_w > 0:
+        valid_patch_mask[:, :crop_w] = False
+
+    return img_resized, patch_mask, {
+        "grid_size": (grid_h, grid_w),
+        "image_shape": original_hw,
+        "patch_size": patch_size,
+        "content_crop": (crop_h, crop_w),
+        "valid_patch_mask": valid_patch_mask,
+    }
+
+
+def _forward_siglip_vision_encoder(vision_encoder, img_resized, patch_mask):
+    """Run a direct forward pass through the SigLIP vision encoder."""
+    if hasattr(vision_encoder, "embeddings") and hasattr(vision_encoder, "encoder"):
+        embeddings = vision_encoder.embeddings(img_resized, patch_mask)
+        return vision_encoder.encoder(embeddings)
+    if hasattr(vision_encoder, "forward"):
+        return vision_encoder(img_resized)
+    return vision_encoder(pixel_values=img_resized, patch_attention_mask=patch_mask)
+
+
+def extract_siglip_last_layer_features(
+    policy,
+    sample,
+    dataset,
+    image_key,
+    device,
+    *,
+    image_map=None,
+    capture_qk: bool = False,
+):
+    """Extract final-layer SigLIP patch features, with optional QK key-logit maps."""
+    vision_encoder = find_vision_encoder(policy)
+    if vision_encoder is None:
+        print("    WARNING: Could not find vision encoder for semantic probe")
+        return None
+
+    try:
+        last_layer = vision_encoder.encoder.layers[-1]
+    except (AttributeError, IndexError):
+        print("    WARNING: Could not access last SigLIP encoder layer")
+        return None
+
+    last_attn = getattr(last_layer, "self_attn", None) or getattr(last_layer, "attention", None)
+    if capture_qk and last_attn is None:
+        print("    WARNING: Could not access last-layer attention module for QK probe")
+        capture_qk = False
+
+    img_resized, patch_mask, meta = _prepare_siglip_vision_inputs(
+        sample, image_key, vision_encoder, device,
+    )
+
+    activations: dict[str, torch.Tensor] = {}
+    q_capture: dict[str, torch.Tensor] = {}
+    k_capture: dict[str, torch.Tensor] = {}
+    attn_input: dict[str, torch.Tensor] = {}
+    handles = []
+
+    def _layer_hook(module, input_args, output):
+        out = output[0] if isinstance(output, tuple) else output
+        activations["value"] = out.detach()
+        return output
+
+    handles.append(last_layer.register_forward_hook(_layer_hook))
+
+    if capture_qk:
+        if hasattr(last_attn, "q_proj") and hasattr(last_attn, "k_proj"):
+            def _q_hook(module, input_args, output):
+                q_capture.setdefault("value", output.detach())
+                return output
+
+            def _k_hook(module, input_args, output):
+                k_capture.setdefault("value", output.detach())
+                return output
+
+            handles.append(last_attn.q_proj.register_forward_hook(_q_hook))
+            handles.append(last_attn.k_proj.register_forward_hook(_k_hook))
+        elif isinstance(last_attn, torch.nn.MultiheadAttention):
+            def _pre_hook(module, input_args):
+                attn_input.setdefault("value", input_args[0].detach())
+                return None
+
+            handles.append(last_attn.register_forward_pre_hook(_pre_hook))
+
+    try:
+        with torch.no_grad():
+            try:
+                _forward_siglip_vision_encoder(vision_encoder, img_resized, patch_mask)
+            except Exception as direct_error:
+                batch, _ = build_policy_batch_from_sample(
+                    sample, policy, device, batch_size=1,
+                    image_key_for_grad=None, dataset=dataset,
+                    image_map=image_map,
+                )
+                try:
+                    policy.reset()
+                    policy.select_action(batch)
+                except Exception as fallback_error:
+                    print(
+                        "    WARNING: SigLIP feature extraction failed "
+                        f"(direct={direct_error}, fallback={fallback_error})"
+                    )
+                    return None
+
+        if "value" not in activations:
+            print("    WARNING: Last-layer feature hook did not fire")
+            return None
+
+        patch_features = activations["value"]
+        if patch_features.dim() == 3:
+            patch_features = patch_features[0]
+        patch_features = F.normalize(patch_features.float(), dim=-1)
+
+        grid_h, grid_w = meta["grid_size"]
+        patch_grid = patch_features.reshape(grid_h, grid_w, -1).detach().cpu().numpy()
+        result = {
+            "patch_features": patch_grid,
+            "grid_size": meta["grid_size"],
+            "image_shape": meta["image_shape"],
+            "patch_size": meta["patch_size"],
+            "content_crop": meta["content_crop"],
+            "valid_patch_mask": meta["valid_patch_mask"],
+        }
+
+        if capture_qk:
+            q_tensor = q_capture.get("value")
+            k_tensor = k_capture.get("value")
+            if q_tensor is None or k_tensor is None:
+                hidden = attn_input.get("value")
+                if hidden is not None and isinstance(last_attn, torch.nn.MultiheadAttention):
+                    qkv = F.linear(hidden, last_attn.in_proj_weight, last_attn.in_proj_bias)
+                    q_tensor, k_tensor, _ = qkv.chunk(3, dim=-1)
+
+            if q_tensor is not None and k_tensor is not None:
+                if q_tensor.dim() == 3:
+                    q_tensor = q_tensor[0]
+                if k_tensor.dim() == 3:
+                    k_tensor = k_tensor[0]
+                num_heads = getattr(last_attn, "num_heads", None) or getattr(
+                    last_attn, "num_attention_heads", None,
+                )
+                if num_heads is None or num_heads <= 0:
+                    num_heads = 1
+                head_dim = int(q_tensor.shape[-1] // num_heads)
+                q = q_tensor.reshape(-1, num_heads, head_dim).permute(1, 0, 2).float()
+                k = k_tensor.reshape(-1, num_heads, head_dim).permute(1, 0, 2).float()
+                scores = torch.matmul(q, k.transpose(1, 2)) * (head_dim ** -0.5)
+                key_maps = scores.mean(dim=1).detach().cpu().numpy()
+                result["qk_key_maps"] = key_maps.reshape(num_heads, grid_h, grid_w)
+
+        return result
+    finally:
+        for handle in handles:
+            handle.remove()
+
+
 # ---------------------------------------------------------------------------
 # Saliency map (vanilla input-gradient)
 # ---------------------------------------------------------------------------

web/frontend/src/components/DiagnosticPanel.tsx

@@ -6,6 +6,46 @@ import CounterfactualComparison from "./CounterfactualComparison";
 
 interface DiagnosticReport {
   metadata: Record<string, unknown>;
+  semantic_probe?: {
+    target_object: string | null;
+    summary: string;
+    primary_frame?: {
+      frame_id: string;
+      target_object: string | null;
+      candidate_labels: string[];
+      best_non_target_label: string | null;
+      target_semantic_peak: number | null;
+      target_semantic_mean_on_causal_patches: number | null;
+      target_margin_over_best_non_target: number | null;
+      causal_semantic_alignment: number | null;
+      background_semantic_gap: number | null;
+      summary: string;
+    } | null;
+    counterfactuals?: Record<string, {
+      test_type: string;
+      summary: string;
+      metrics: Record<string, number | string | boolean | number[] | string[]>;
+    }>;
+  };
+  qk_probe?: {
+    layer: string;
+    summary: string;
+    dominant_head_type: string;
+    semantic_head_fraction: number;
+    positional_head_fraction: number;
+    mixed_head_fraction: number;
+    top_heads: Array<{
+      head_index: number;
+      head_type: string;
+      score: number;
+      semantic_map_correlation: number;
+      positional_baseline_correlation: number;
+      target_region_logit_mass: number;
+      background_logit_mass: number;
+      old_anchor_logit_mass?: number | null;
+      moved_object_logit_mass?: number | null;
+    }>;
+  };
   spatial_object_diagnosis?: {
     target_object: string | null;
     verdict: string;
@@ -290,7 +330,11 @@ export default function DiagnosticPanel() {
       </div>
 
       {report.spatial_object_diagnosis && (
-        <SpatialObjectDiagnosisCard diagnosis={report.spatial_object_diagnosis} />
+        <SpatialObjectDiagnosisCard
+          diagnosis={report.spatial_object_diagnosis}
+          semanticProbe={report.semantic_probe}
+          qkProbe={report.qk_probe}
+        />
       )}
 
       {/* Scene Understanding */}
@@ -426,8 +470,12 @@ export default function DiagnosticPanel() {
 
 function SpatialObjectDiagnosisCard({
   diagnosis,
+  semanticProbe,
+  qkProbe,
 }: {
   diagnosis: NonNullable<DiagnosticReport["spatial_object_diagnosis"]>;
+  semanticProbe?: DiagnosticReport["semantic_probe"];
+  qkProbe?: DiagnosticReport["qk_probe"];
 }) {
   const verdictLabel: Record<string, string> = {
     spatial_prior: "Spatial Prior",
@@ -480,6 +528,72 @@ function SpatialObjectDiagnosisCard({
         <EvidenceList title="Evidence For Spatial Priors" items={diagnosis.spatial_evidence} emptyText="No strong spatial-prior evidence." />
         <EvidenceList title="Evidence For Object Grounding" items={diagnosis.object_evidence} emptyText="No strong object-grounding evidence." />
       </div>
+
+      {(semanticProbe || qkProbe) && (
+        <div style={{ marginTop: 12, display: "grid", gridTemplateColumns: "repeat(auto-fit, minmax(260px, 1fr))", gap: 12 }}>
+          {semanticProbe && (
+            <div style={{ background: "#FCFDFD", border: "1px solid #DBE4E8", borderRadius: 8, padding: 12 }}>
+              <div style={{ fontWeight: 600, color: "var(--text-heading)", fontSize: 12, marginBottom: 8 }}>
+                Semantic Probe
+              </div>
+              <div style={{ color: "var(--text-body)", fontSize: 13, lineHeight: 1.5, marginBottom: 8 }}>
+                {semanticProbe.summary}
+              </div>
+              {semanticProbe.primary_frame && (
+                <div style={{ display: "grid", gridTemplateColumns: "repeat(auto-fit, minmax(160px, 1fr))", gap: 8, marginBottom: 8 }}>
+                  <MetricPill label="Semantic peak" value={formatOptionalMetric(semanticProbe.primary_frame.target_semantic_peak)} />
+                  <MetricPill label="Causal mean" value={formatOptionalMetric(semanticProbe.primary_frame.target_semantic_mean_on_causal_patches)} />
+                  <MetricPill label="Target margin" value={formatOptionalMetric(semanticProbe.primary_frame.target_margin_over_best_non_target)} />
+                  <MetricPill label="Background gap" value={formatOptionalMetric(semanticProbe.primary_frame.background_semantic_gap)} />
+                </div>
+              )}
+              {semanticProbe.counterfactuals && Object.keys(semanticProbe.counterfactuals).length > 0 && (
+                <div style={{ display: "flex", flexDirection: "column", gap: 8 }}>
+                  {Object.entries(semanticProbe.counterfactuals).map(([key, value]) => (
+                    <div key={key} style={{ fontSize: 12 }}>
+                      <div style={{ color: "var(--text-heading)", fontWeight: 500 }}>
+                        {key.replace(/_/g, " ")}
+                      </div>
+                      <div style={{ color: "var(--text-body)" }}>{value.summary}</div>
+                    </div>
+                  ))}
+                </div>
+              )}
+            </div>
+          )}
+
+          {qkProbe && (
+            <div style={{ background: "#FCFDFD", border: "1px solid #DBE4E8", borderRadius: 8, padding: 12 }}>
+              <div style={{ fontWeight: 600, color: "var(--text-heading)", fontSize: 12, marginBottom: 8 }}>
+                QK Decomposition
+              </div>
+              <div style={{ color: "var(--text-body)", fontSize: 13, lineHeight: 1.5, marginBottom: 8 }}>
+                {qkProbe.summary}
+              </div>
+              <div style={{ display: "grid", gridTemplateColumns: "repeat(auto-fit, minmax(140px, 1fr))", gap: 8, marginBottom: 8 }}>
+                <MetricPill label="Layer" value={qkProbe.layer} />
+                <MetricPill label="Dominant" value={qkProbe.dominant_head_type} />
+                <MetricPill label="Semantic heads" value={qkProbe.semantic_head_fraction.toFixed(3)} />
+                <MetricPill label="Positional heads" value={qkProbe.positional_head_fraction.toFixed(3)} />
+              </div>
+              {qkProbe.top_heads.length > 0 && (
+                <div style={{ display: "flex", flexDirection: "column", gap: 8 }}>
+                  {qkProbe.top_heads.map((head) => (
+                    <div key={head.head_index} style={{ fontSize: 12 }}>
+                      <div style={{ color: "var(--text-heading)", fontWeight: 500 }}>
+                        Head {head.head_index} · {head.head_type}
+                      </div>
+                      <div style={{ color: "var(--text-body)" }}>
+                        semantic corr={head.semantic_map_correlation.toFixed(3)} · positional corr={head.positional_baseline_correlation.toFixed(3)} · score={head.score.toFixed(3)}
+                      </div>
+                    </div>
+                  ))}
+                </div>
+              )}
+            </div>
+          )}
+        </div>
+      )}
     </div>
   );
 }
@@ -536,6 +650,11 @@ function formatMetricValue(value: number | string) {
   return String(value);
 }
 
+function formatOptionalMetric(value: number | null | undefined) {
+  if (typeof value !== "number") return "N/A";
+  return Math.abs(value) >= 1 ? value.toFixed(3) : value.toFixed(4);
+}
+
 function SeverityBadge({ severity }: { severity: string }) {
   const styles: Record<string, { bg: string; color: string }> = {
     critical: { bg: "#dc3545", color: "#fff" },