Update app.py

app.py

@@ -466,7 +466,272 @@ class VIndex:
             "mode":        "activation-guided" if use_act else "embed-based"
         }
 
-    # ── base ops (unchanged) ───────────────────────────────────
+    # ── Phase 1+: mechanistic attribution ─────────────────────
+
+    def gradient_slot_scores(self, prompt: str, target: str) -> Dict:
+        """One backward pass: grad norm of ∂(-log p(target))/∂W_down[:,slot] per KB layer.
+        Identifies which slots causally contributed to this prediction via gradient signal."""
+        target_id = self.token_id(target)
+
+        # Temporarily enable grad on down-proj weights
+        down_params: List[Tuple[int, torch.nn.Parameter]] = []
+        for li in range(self.arch.n_layers):
+            layer = self.arch._layer(li)
+            p = layer.mlp.c_proj.weight if self.arch.style == "gpt2" \
+                else layer.mlp.down_proj.weight
+            p.requires_grad_(True)
+            down_params.append((li, p))
+
+        self.model.zero_grad()
+        inputs = self.tok(prompt, return_tensors="pt").to(self.device)
+        out    = self.model(**inputs)
+        loss   = -F.log_softmax(out.logits[0, -1], dim=-1)[target_id]
+        loss.backward()
+
+        layer_scores = []
+        for li, p in down_params:
+            grad = p.grad
+            p.requires_grad_(False)
+            if grad is None:
+                layer_scores.append({"layer": li, "max_grad": 0.0, "top_slots": []})
+                continue
+            # gpt2: c_proj.weight [ffn_dim, hidden] → rows = slots
+            # gated: down_proj.weight [hidden, ffn_dim] → cols = slots
+            slot_norms = grad.norm(dim=1) if self.arch.style == "gpt2" \
+                         else grad.norm(dim=0)  # [ffn_dim]
+            k = min(20, slot_norms.shape[0])
+            vals, idxs = slot_norms.topk(k)
+            layer_scores.append({
+                "layer":     li,
+                "max_grad":  round(float(vals[0].item()), 6),
+                "top_slots": [{"slot": int(idx.item()),
+                               "grad_norm": round(float(v.item()), 6)}
+                              for idx, v in zip(idxs, vals)]
+            })
+
+        self.model.zero_grad()
+        return {"layer_scores": layer_scores}
+
+    def causal_patch_trace(self, prompt: str, subject: str, target: str,
+                           noise_std: float = 0.1) -> Dict:
+        """ROME-style causal tracing.
+        Corrupts subject embeddings, then for each KB layer measures how much
+        patching that layer's hidden state (at subject position) restores p(target).
+        Expensive: O(n_layers) forward passes."""
+        target_id = self.token_id(target)
+        W_u   = self.arch.get_unembedding().to(self.device)
+        inputs = self.tok(prompt, return_tensors="pt").to(self.device)
+        ids    = inputs["input_ids"][0].tolist()
+
+        # Find subject token positions via subsequence match
+        subj_ids = self.tok.encode(subject, add_special_tokens=False)
+        subj_pos: List[int] = []
+        for start in range(len(ids) - len(subj_ids) + 1):
+            if ids[start:start+len(subj_ids)] == subj_ids:
+                subj_pos = list(range(start, start+len(subj_ids)))
+                break
+        if not subj_pos:
+            for si in subj_ids:
+                if si in ids:
+                    subj_pos = [ids.index(si)]
+                    break
+        if not subj_pos:
+            subj_pos = [0]
+
+        # ── Clean forward — capture every layer's hidden states ──
+        clean_hs: Dict[int, torch.Tensor] = {}
+        clean_handles = []
+        def _mk_clean(li):
+            def _h(m, inp, out):
+                h = out[0] if isinstance(out, tuple) else out
+                clean_hs[li] = h[0].detach().clone()   # [seq, hidden]
+            return _h
+        for li in range(self.arch.n_layers):
+            clean_handles.append(self.arch._layer(li).register_forward_hook(_mk_clean(li)))
+        with torch.no_grad():
+            clean_out = self.model(**inputs)
+        for h in clean_handles: h.remove()
+        clean_prob = float(torch.softmax(clean_out.logits[0,-1], dim=-1)[target_id].item())
+
+        # ── Corrupted embeddings ──
+        E = self.arch.get_embedding().to(self.device)
+        emb = E[inputs["input_ids"][0]].unsqueeze(0).clone()     # [1, seq, hidden]
+        noise_scale = emb.std().item() * noise_std
+        for pos in subj_pos:
+            emb[0, pos] += torch.randn_like(emb[0, pos]) * noise_scale
+
+        with torch.no_grad():
+            corr_out = self.model(inputs_embeds=emb)
+        corr_prob = float(torch.softmax(corr_out.logits[0,-1], dim=-1)[target_id].item())
+
+        # ── Causal patch sweep ──
+        results = []
+        for li in range(self.kb_start, self.kb_end):
+            def _mk_patch(target_li):
+                def _h(m, inp, out):
+                    if target_li not in clean_hs:
+                        return out
+                    is_tuple = isinstance(out, tuple)
+                    h = list(out) if is_tuple else [out]
+                    clean = clean_hs[target_li]
+                    for pos in subj_pos:
+                        if pos < clean.shape[0]:
+                            h[0][0, pos] = clean[pos].to(h[0].device)
+                    return tuple(h) if is_tuple else h[0]
+                return _h
+            ph = self.arch._layer(li).register_forward_hook(_mk_patch(li))
+            with torch.no_grad():
+                patch_out = self.model(inputs_embeds=emb.clone())
+            ph.remove()
+            patch_prob = float(torch.softmax(patch_out.logits[0,-1], dim=-1)[target_id].item())
+            ie = patch_prob - corr_prob
+            results.append({
+                "layer":           li,
+                "patch_prob":      round(patch_prob, 6),
+                "indirect_effect": round(ie, 6),
+            })
+
+        return {
+            "clean_prob":      round(clean_prob, 6),
+            "corrupt_prob":    round(corr_prob, 6),
+            "subject_pos":     subj_pos,
+            "results":         results,
+        }
+
+    def smart_locate(self, prompt: str, subject: str, target: str,
+                     alpha: float = 0.4, beta: float = 0.3, gamma: float = 0.3,
+                     noise_std: float = 0.1) -> Dict:
+        """Combined gate_sim + grad_norm + causal_effect → precise layer/slot ranking.
+        alpha = weight for gate cosine sim
+        beta  = weight for gradient norm
+        gamma = weight for causal indirect effect"""
+        gate_data   = self.locate(prompt, subject, target)
+        grad_data   = self.gradient_slot_scores(prompt, target)
+        causal_data = self.causal_patch_trace(prompt, subject, target, noise_std=noise_std)
+
+        gate_map   = {ls["layer"]: ls["max_sim"]  for ls in gate_data["layer_scores"]}
+        grad_map   = {ls["layer"]: ls["max_grad"]  for ls in grad_data["layer_scores"]}
+        causal_map = {r["layer"]:  max(0.0, r["indirect_effect"])
+                      for r in causal_data["results"]}
+        grad_slots = {ls["layer"]: ls["top_slots"] for ls in grad_data["layer_scores"]}
+
+        layers = sorted(set(gate_map) | set(grad_map) | set(causal_map))
+
+        def _norm(vals: List[float]) -> List[float]:
+            m = max(vals) if vals else 1.0
+            return [v/m if m > 0 else 0.0 for v in vals]
+
+        gv = [gate_map.get(l, 0.0)   for l in layers]
+        dv = [grad_map.get(l, 0.0)   for l in layers]
+        cv = [causal_map.get(l, 0.0) for l in layers]
+        gn, dn, cn = _norm(gv), _norm(dv), _norm(cv)
+
+        ranked = []
+        for i, l in enumerate(layers):
+            score = alpha*gn[i] + beta*dn[i] + gamma*cn[i]
+            ranked.append({
+                "layer":          l,
+                "gate_sim":       round(gv[i], 4),
+                "grad_norm":      round(dv[i], 6),
+                "causal_effect":  round(cv[i], 6),
+                "gate_sim_n":     round(gn[i], 4),
+                "grad_norm_n":    round(dn[i], 4),
+                "causal_n":       round(cn[i], 4),
+                "combined":       round(score, 4),
+                "best_slots":     (grad_slots.get(l) or [])[:5],
+            })
+
+        ranked.sort(key=lambda x: -x["combined"])
+
+        return {
+            "ranked_layers":  ranked,
+            "phase_layer":    gate_data["phase_layer"],
+            "subject_pos":    gate_data["subject_pos"],
+            "clean_prob":     causal_data["clean_prob"],
+            "corrupt_prob":   causal_data["corrupt_prob"],
+            "recommendation": ranked[0] if ranked else None,
+            "weights":        {"alpha": alpha, "beta": beta, "gamma": gamma},
+        }
+
+    def smart_edit(self, prompt: str, subject: str, relation: str,
+                   old_target: str, new_target: str,
+                   top_layers: int = 3, slots_per_layer: int = 2,
+                   scale: float = 1.5, noise_std: float = 0.1,
+                   alpha: float = 0.4, beta: float = 0.4, gamma: float = 0.2,
+                   log: Optional[List[str]] = None) -> Dict:
+        """Auto edit: runs smart_locate on (prompt, subject, old_target) to find
+        the exact layer+slot targets via gradient+causal+gate consensus, then
+        patches those W_down columns toward embed(new_target).
+
+        old_target  = what the model currently predicts (used to locate)
+        new_target  = what you want to inject
+        top_layers  = how many top-ranked layers to patch
+        slots_per_layer = gradient-identified slots to patch per layer
+        scale       = col_norm multiplier (1.5-3.0 recommended)
+        beta > alpha because grad_norm is more reliable than gate_sim for small models."""
+        if log is None: log = []
+        self._snapshot()
+
+        log.append(f"SMART_EDIT: '{subject}' [{relation}] {old_target!r} → {new_target!r}")
+        log.append(f"  Running smart_locate on prompt: {prompt!r}")
+        log.append(f"  Weights: ��={alpha} β={beta} γ={gamma}  noise_std={noise_std}")
+
+        sl = self.smart_locate(prompt, subject, old_target,
+                               alpha=alpha, beta=beta, gamma=gamma,
+                               noise_std=noise_std)
+
+        log.append(f"  clean_prob={sl['clean_prob']:.6f}  corrupt_prob={sl['corrupt_prob']:.6f}")
+        log.append(f"  Phase layer: L{sl['phase_layer']}  Subject pos: {sl['subject_pos']}")
+
+        if sl["clean_prob"] < 1e-5:
+            log.append("  ⚠ clean_prob near zero — model barely knows this fact.")
+            log.append("    Grad-norm signal still valid. Causal IE=0 is expected.")
+            log.append("    Recommend: gpt2-medium or Qwen2.5-1.5B for stronger facts.")
+
+        tv    = self.embed(new_target)
+        tv_n  = F.normalize(tv, dim=0)
+        ops   = []
+        used  = []
+
+        top_ranked = sl["ranked_layers"][:top_layers]
+        for lr in top_ranked:
+            li = lr["layer"]
+            # Use gradient-identified slots — far more precise than gate cosine
+            grad_slots = [s["slot"] for s in lr["best_slots"][:slots_per_layer]]
+            if not grad_slots:
+                log.append(f"  L{li}: no grad slots, skipping")
+                continue
+            _, Wd = self.arch.get_ffn_weights(li)
+            Wd = Wd.to(self.device)
+            for slot in grad_slots:
+                col_norm = Wd[:, slot].norm().item()
+                new_col  = (tv_n * col_norm * scale).cpu().tolist()
+                ops.append({"op":"update_down","layer":li,"slot":slot,"down_col":new_col})
+                log.append(f"  ✓ L{li} slot {slot}: combined={lr['combined']}  "
+                            f"grad_norm={lr['grad_norm']:.4f}  col_norm={col_norm:.4f}  "
+                            f"inject={col_norm*scale:.4f}")
+            used.append({"layer":li,"slots":grad_slots,"combined":lr["combined"]})
+
+        self.patches.append({
+            "type":         "SMART_UPDATE",
+            "entity":       subject,
+            "relation":     relation,
+            "new_target":   new_target,
+            "old_target":   old_target,
+            "smart_top":    top_ranked,
+            "ops":          ops,
+        })
+        self._apply_all_patches()
+        log.append(f"\n  ✓ {len(ops)} op(s) across {len(used)} layer(s), patch #{len(self.patches)}")
+
+        return {
+            "ops":          ops,
+            "used_layers":  used,
+            "smart_locate": sl,
+            "log":          log,
+        }
+
+
 
     def infer(self, prompt: str, top_k: int = 5):
         probs  = torch.softmax(self._forward(prompt), dim=-1)
@@ -771,10 +1036,12 @@ svg text { font-family: var(--font); fill: var(--text); }
   <button class="tab-btn" onclick="showTab('describe')">② Describe</button>
   <button class="tab-btn" onclick="showTab('trace')">③ Trace</button>
   <button class="tab-btn" onclick="showTab('locate')">④ Locate</button>
-  <button class="tab-btn" onclick="showTab('heatmap')">⑤ Heatmap</button>
-  <button class="tab-btn" onclick="showTab('edit')">⑥ Edit</button>
-  <button class="tab-btn" onclick="showTab('patches')">⑦ Patches</button>
-  <button class="tab-btn" onclick="showTab('load')" style="margin-left:auto">⚙ Load</button>
+  <button class="tab-btn" onclick="showTab('smartlocate')">⑤ Smart Locate</button>
+  <button class="tab-btn" onclick="showTab('heatmap')">⑥ Heatmap</button>
+  <button class="tab-btn" onclick="showTab('edit')">⑦ Edit</button>
+  <button class="tab-btn" onclick="showTab('patches')">⑧ Patches</button>
+  <button class="tab-btn" onclick="showTab('guide')" style="margin-left:auto;color:var(--green)">📖 Guide</button>
+  <button class="tab-btn" onclick="showTab('load')">⚙ Load</button>
 </div>
 
 <!-- TOOLTIP -->
@@ -925,6 +1192,86 @@ svg text { font-family: var(--font); fill: var(--text); }
   </div>
 </div>
 
+<!-- ══════════ SMART LOCATE PANEL ══════════ -->
+<div id="panel-smartlocate" class="panel">
+  <div class="card">
+    <h3>Smart Locate — gradient + causal + gate_sim combined</h3>
+    <div style="color:var(--muted);font-size:11px;margin-bottom:12px;line-height:1.7">
+      Three independent signals combined into one ranked layer list.<br>
+      <span style="color:var(--blue)">■ gate_sim</span> — static embedding cosine (fast, weak proxy) &nbsp;
+      <span style="color:var(--green)">■ grad_norm</span> — ∂loss/∂W_down per slot (one backward pass) &nbsp;
+      <span style="color:var(--yellow)">■ causal IE</span> — indirect effect via subject-corruption patching (N_layers passes, slow)
+    </div>
+    <div class="row">
+      <div class="col2">
+        <label>Prompt</label>
+        <input type="text" id="sl-prompt" value="The capital of France is">
+      </div>
+      <div class="col">
+        <label>Subject</label>
+        <input type="text" id="sl-subject" value="France">
+      </div>
+      <div class="col">
+        <label>Target</label>
+        <input type="text" id="sl-target" value="Paris">
+      </div>
+    </div>
+    <div class="row" style="margin-top:10px">
+      <div class="col">
+        <label>α gate_sim: <span id="sl-a-val">0.4</span></label>
+        <input type="range" id="sl-alpha" min="0" max="1" step="0.05" value="0.4"
+               oninput="document.getElementById('sl-a-val').textContent=this.value">
+      </div>
+      <div class="col">
+        <label>β grad_norm: <span id="sl-b-val">0.3</span></label>
+        <input type="range" id="sl-beta" min="0" max="1" step="0.05" value="0.3"
+               oninput="document.getElementById('sl-b-val').textContent=this.value">
+      </div>
+      <div class="col">
+        <label>γ causal: <span id="sl-g-val">0.3</span></label>
+        <input type="range" id="sl-gamma" min="0" max="1" step="0.05" value="0.3"
+               oninput="document.getElementById('sl-g-val').textContent=this.value">
+      </div>
+      <div class="col">
+        <label>Noise σ: <span id="sl-noise-val">0.1</span></label>
+        <input type="range" id="sl-noise" min="0.02" max="0.5" step="0.02" value="0.1"
+               oninput="document.getElementById('sl-noise-val').textContent=this.value">
+      </div>
+    </div>
+    <div style="display:flex;gap:8px;margin-top:12px;flex-wrap:wrap">
+      <button onclick="runSmartLocate()">⚡ Smart Locate (full)</button>
+      <button class="secondary" onclick="runGradientOnly()">▶ Gradient only (fast)</button>
+      <button class="secondary" onclick="runCausalOnly()">▶ Causal trace only</button>
+    </div>
+    <div id="sl-status" style="color:var(--muted);font-size:11px;margin-top:8px"></div>
+  </div>
+
+  <div class="row">
+    <div class="col2 card">
+      <h3>Layer Rankings — 3-signal stacked bars</h3>
+      <div class="chart-wrap" id="sl-chart" style="min-height:320px"></div>
+    </div>
+    <div class="col card">
+      <h3>Recommendation</h3>
+      <div id="sl-rec" class="log">Run Smart Locate to see the best edit target.</div>
+      <h3 style="margin-top:14px">Collateral Probe</h3>
+      <div class="row" style="margin-top:8px">
+        <input type="text" id="sl-coll-prompt" value="Biggest cities in France"
+               style="flex:2" placeholder="Collateral prompt…">
+        <button class="secondary" onclick="runCollateral()" style="flex:0">▶</button>
+      </div>
+      <div id="sl-coll-out" class="log" style="margin-top:8px">Probe a prompt to check collateral damage.</div>
+    </div>
+  </div>
+
+  <div class="card">
+    <h3>Per-Layer Detail</h3>
+    <div id="sl-table" style="overflow-x:auto">
+      <div style="color:var(--muted);font-size:11px">Run Smart Locate first.</div>
+    </div>
+  </div>
+</div>
+
 <!-- ══════════ HEATMAP PANEL ══════════ -->
 <div id="panel-heatmap" class="panel">
   <div class="card">
@@ -980,6 +1327,7 @@ svg text { font-family: var(--font); fill: var(--text); }
         <div class="radio-group" id="edit-mode-group">
           <label><input type="radio" name="edit-mode" value="UPDATE" checked> UPDATE</label>
           <label><input type="radio" name="edit-mode" value="PRECISE"> PRECISE</label>
+          <label><input type="radio" name="edit-mode" value="SMART"> ★ SMART</label>
           <label><input type="radio" name="edit-mode" value="INSERT"> INSERT</label>
           <label><input type="radio" name="edit-mode" value="SUPPRESS"> SUPPRESS</label>
           <label><input type="radio" name="edit-mode" value="AMPLIFY"> AMPLIFY</label>
@@ -990,6 +1338,28 @@ svg text { font-family: var(--font); fill: var(--text); }
           <label>Prompt (PRECISE mode)</label>
           <input type="text" id="edit-prompt" value="The capital of France is">
         </div>
+        <div id="smart-row" style="display:none;margin-top:8px;background:var(--bg);border:1px solid var(--border);border-radius:6px;padding:10px">
+          <div style="color:var(--blue);font-size:11px;font-weight:700;margin-bottom:6px">★ SMART AUTO MODE</div>
+          <label>Prompt (used for locate + after-check)</label>
+          <input type="text" id="smart-prompt" value="The capital of France is">
+          <label style="margin-top:6px">Old value (what model currently says)</label>
+          <input type="text" id="smart-old" value="Paris">
+          <div class="row" style="margin-top:6px">
+            <div class="col">
+              <label>Top layers: <span id="smart-layers-val">3</span></label>
+              <input type="range" id="smart-layers" min="1" max="8" value="3"
+                     oninput="document.getElementById('smart-layers-val').textContent=this.value">
+            </div>
+            <div class="col">
+              <label>Slots/layer: <span id="smart-slots-val">2</span></label>
+              <input type="range" id="smart-slots" min="1" max="5" value="2"
+                     oninput="document.getElementById('smart-slots-val').textContent=this.value">
+            </div>
+          </div>
+          <div style="color:var(--muted);font-size:10px;margin-top:6px">
+            Runs smart_locate internally → patches gradient-identified slots. No manual tuning needed.
+          </div>
+        </div>
         <div id="style-shift-row" style="display:none;margin-top:8px">
           <label>From concept</label>
           <input type="text" id="ss-from" value="formal">
@@ -1069,6 +1439,209 @@ svg text { font-family: var(--font); fill: var(--text); }
   </div>
 </div>
 
+<!-- ══════════ GUIDE PANEL ══════════ -->
+<div id="panel-guide" class="panel">
+<div class="card">
+<h3>What is VINDEX doing?</h3>
+<div style="line-height:1.9;color:var(--muted)">
+In a transformer, factual associations like <span style="color:var(--text)">"France → capital → Paris"</span>
+are stored as direction vectors in the <span style="color:var(--blue)">W_down columns</span> of FFN layers.
+The <span style="color:var(--blue)">W_gate rows</span> act as keys: when the residual stream resembles "France",
+the matching gate fires, the down column adds "Paris" direction to the stream, and the unembedding reads out "Paris".
+VINDEX surgically replaces those down columns without retraining.
+</div>
+</div>
+
+<div class="card">
+<h3>Quickstart — 5-step experiment</h3>
+<div style="line-height:2;font-size:12px">
+<div style="color:var(--yellow);margin-bottom:4px">Step 1 — Load a model that actually knows facts</div>
+<div style="color:var(--muted);margin-left:16px;margin-bottom:10px">
+  ⚙ Load tab → <span style="color:var(--blue)">gpt2-medium</span> (1.5 GB, knows capitals) or
+  <span style="color:var(--blue)">Qwen/Qwen2.5-1.5B-Instruct</span> (3 GB, strong).<br>
+  distilgpt2 has clean_prob≈0 for most facts → causal IE=0 everywhere → misleading results.
+</div>
+
+<div style="color:var(--yellow);margin-bottom:4px">Step 2 — Verify the model knows the fact</div>
+<div style="color:var(--muted);margin-left:16px;margin-bottom:10px">
+  ① Infer: prompt = <code>"The capital of France is"</code><br>
+  ✓ Good: "Paris" appears in top-3 with prob &gt; 0.05<br>
+  ✗ Bad: top tokens are "a", "the", "known" → model doesn't know it → skip to INSERT mode
+</div>
+
+<div style="color:var(--yellow);margin-bottom:4px">Step 3 — Find where the fact lives</div>
+<div style="color:var(--muted);margin-left:16px;margin-bottom:10px">
+  ③ Trace: prompt = <code>"The capital of France is"</code>, target = <code>"Paris"</code><br>
+  → Look for phase layer: where rank drops from ~30000 to &lt;100. That's where the fact materializes.<br>
+  ⑤ Smart Locate → Gradient only (fast, 1 backward pass):<br>
+  <span style="margin-left:16px">subject = <code>France</code>, target = <code>Paris</code></span><br>
+  → The layer with highest grad_norm bar = best edit target. Note the slot numbers.
+</div>
+
+<div style="color:var(--yellow);margin-bottom:4px">Step 4 — Edit with SMART mode</div>
+<div style="color:var(--muted);margin-left:16px;margin-bottom:10px">
+  ⑦ Edit tab → mode = <span style="color:var(--blue)">★ SMART</span><br>
+  Entity = <code>France</code> | Relation = <code>capital</code><br>
+  Old value = <code>Paris</code> (what model says now — used for locate)<br>
+  New value = <code>Lyon</code> (what you want)<br>
+  Prompt = <code>"The capital of France is"</code><br>
+  Scale = <code>2.0</code> (start here; increase to 3.0 if effect is weak)<br>
+  → Click <b>Apply Edit</b>. Smart locate runs internally, patches grad-identified slots.
+</div>
+
+<div style="color:var(--yellow);margin-bottom:4px">Step 5 — Check collateral damage</div>
+<div style="color:var(--muted);margin-left:16px;margin-bottom:10px">
+  ① Infer: <code>"The capital of France is"</code> → should now say Lyon<br>
+  ① Infer: <code>"Biggest cities in France"</code> → should be unchanged (different slots)<br>
+  ① Infer: <code>"Paris is a city in"</code> → should still say France<br>
+  ① Infer: <code>"Lyon is a city in"</code> → might now also say France (collateral)<br>
+  ⑤ Smart Locate collateral probe → run these prompts, compare slot lists in ⑧ Patches
+</div>
+</div>
+</div>
+
+<div class="card">
+<h3>Interpreting Smart Locate results</h3>
+<div style="font-size:11px;line-height:1.9">
+<div class="row" style="gap:20px">
+<div class="col">
+<div style="color:var(--blue);font-weight:700;margin-bottom:6px">■ gate_sim (blue)</div>
+<div style="color:var(--muted)">
+  Cosine between W_gate[slot] and embed(subject).<br>
+  Fast, cheap, but <b>weak proxy</b> — measures embedding-space similarity,<br>
+  not causal contribution. Useful for finding <i>related</i> slots.<br>
+  <b>High gate_sim + low grad_norm</b> = slot activates for this entity<br>
+  but doesn't contribute much to this specific prediction.
+</div>
+</div>
+<div class="col">
+<div style="color:var(--green);font-weight:700;margin-bottom:6px">■ grad_norm (green)</div>
+<div style="color:var(--muted)">
+  ‖∂(-log p(target))/∂W_down[:,slot]‖ — how much changing this slot<br>
+  would affect the loss for this (prompt, target) pair.<br>
+  <b>Most reliable signal</b>, works even when clean_prob is tiny.<br>
+  One backward pass. Use β &gt; α to weight this higher.<br>
+  <b>High grad_norm</b> = this slot is causally upstream of the prediction.
+</div>
+</div>
+<div class="col">
+<div style="color:var(--yellow);font-weight:700;margin-bottom:6px">■ causal IE (yellow)</div>
+<div style="color:var(--muted)">
+  Indirect effect via noise-corruption patching (ROME-style).<br>
+  Measures: if I corrupt subject embeddings, how much does patching<br>
+  layer L's hidden state at subject pos <i>restore</i> the prediction?<br>
+  <b>Most interpretable</b> — true causal measurement. But:<br>
+  If clean_prob ≈ 0, IE = 0 everywhere (nothing to restore).<br>
+  Needs a model that actually knows the fact.
+</div>
+</div>
+</div>
+<div style="margin-top:12px;padding:10px;background:var(--bg);border-radius:6px;border:1px solid var(--border)">
+  <span style="color:var(--yellow)">⚠ Your distilgpt2 result:</span>
+  <span style="color:var(--muted)"> clean_prob=0.000001 → causal IE=0 everywhere (expected, not a bug).
+  grad_norm on L9/slot515 IS real signal — that slot responds to France+capital context in the gradient sense.
+  But the probability mass is too diffuse to show causal separation.
+  Switch to gpt2-medium for textbook causal results.</span>
+</div>
+</div>
+</div>
+
+<div class="card">
+<h3>Edit modes — when to use which</h3>
+<div style="font-size:11px">
+<table style="width:100%;border-collapse:collapse">
+<thead><tr style="border-bottom:1px solid var(--border);color:var(--muted)">
+  <th style="padding:6px 8px;text-align:left">Mode</th>
+  <th style="padding:6px 8px;text-align:left">Slot selection</th>
+  <th style="padding:6px 8px;text-align:left">Best for</th>
+  <th style="padding:6px 8px;text-align:left">Knobs</th>
+</tr></thead>
+<tbody style="color:var(--muted)">
+<tr style="border-bottom:1px solid var(--border)">
+  <td style="padding:6px 8px;color:var(--blue)">UPDATE</td>
+  <td style="padding:6px 8px">gate cosine sim to embed(entity)</td>
+  <td style="padding:6px 8px">Quick experiment, model knows the fact well</td>
+  <td style="padding:6px 8px">Top-K=3-5, Scale=1.5-3</td>
+</tr>
+<tr style="border-bottom:1px solid var(--border)">
+  <td style="padding:6px 8px;color:var(--purple)">PRECISE</td>
+  <td style="padding:6px 8px">gate cosine sim to h_L[subject_pos]</td>
+  <td style="padding:6px 8px">In-context subject representation (3-5× better than UPDATE)</td>
+  <td style="padding:6px 8px">+ Prompt field</td>
+</tr>
+<tr style="border-bottom:1px solid var(--border)">
+  <td style="padding:6px 8px;color:var(--yellow)">★ SMART</td>
+  <td style="padding:6px 8px">gradient norm → exact slots, then patch</td>
+  <td style="padding:6px 8px"><b>Best overall.</b> Auto-locates, no manual tuning</td>
+  <td style="padding:6px 8px">Top layers=3, Slots/layer=2, Scale=1.5-2.5</td>
+</tr>
+<tr style="border-bottom:1px solid var(--border)">
+  <td style="padding:6px 8px;color:var(--green)">INSERT</td>
+  <td style="padding:6px 8px">weakest slot (norm-based)</td>
+  <td style="padding:6px 8px">Model has no knowledge of fact, build from scratch</td>
+  <td style="padding:6px 8px">Alpha=0.4-0.7, Spread=4-6</td>
+</tr>
+<tr style="border-bottom:1px solid var(--border)">
+  <td style="padding:6px 8px;color:var(--red)">SUPPRESS</td>
+  <td style="padding:6px 8px">gate cosine → scale W_down to 0</td>
+  <td style="padding:6px 8px">Make model forget an entity (factor=0) or weaken (0.5)</td>
+  <td style="padding:6px 8px">Factor: 0=forget, 0.5=weaken</td>
+</tr>
+<tr style="border-bottom:1px solid var(--border)">
+  <td style="padding:6px 8px;color:var(--cyan)">STYLE-SHIFT</td>
+  <td style="padding:6px 8px">gate cosine → add direction vector</td>
+  <td style="padding:6px 8px">Bias/tone shifts: CEO→less male-coded, Paris→darker</td>
+  <td style="padding:6px 8px">from/to concepts, strength=0.3-0.8</td>
+</tr>
+</tbody>
+</table>
+</div>
+</div>
+
+<div class="card">
+<h3>Experiments to run</h3>
+<div style="font-size:11px;line-height:1.9;color:var(--muted)">
+<div style="color:var(--text);margin-bottom:4px">Experiment A — Capital swap (classic ROME benchmark)</div>
+Model: gpt2-medium | Prompt: "The capital of France is" | Old: Paris | New: Lyon<br>
+Check: "France's capital city" | "Lyon is now" | "Paris is in" | "Eiffel Tower is in"<br>
+Insight: does it generalize (paraphrase) or is it prompt-specific?<br><br>
+
+<div style="color:var(--text);margin-bottom:4px">Experiment B — Slot overlap analysis (your collateral question)</div>
+1. SMART locate "The capital of France is" → note slot numbers in recommendation<br>
+2. SMART locate "The biggest city in France is" → compare slot lists<br>
+3. Overlap = slots that will be collaterally damaged<br>
+4. No overlap = clean surgery ✓<br><br>
+
+<div style="color:var(--text);margin-bottom:4px">Experiment C — Suppression then INSERT</div>
+SUPPRESS France → then INSERT France capital Lyon → Infer<br>
+vs just UPDATE. Which gives cleaner, more confident result?<br><br>
+
+<div style="color:var(--text);margin-bottom:4px">Experiment D — Style shift (no factual change)</div>
+STYLE-SHIFT: anchor=CEO, from="male", to="female", strength=0.3<br>
+Then Infer: "The CEO of the company is a" — does pronoun distribution shift?<br>
+Insight: this is mechanical debiasing without retraining.<br><br>
+
+<div style="color:var(--text);margin-bottom:4px">Experiment E — Compile and compare</div>
+Edit 5 facts. Compile → save as new model directory.<br>
+Load compiled model fresh → Infer same prompts → edits should persist in weights.<br>
+Then Trace on compiled model → phase layers should shift or sharpen.
+</div>
+</div>
+
+<div class="card">
+<h3>α β γ tuning guide</h3>
+<div style="font-size:11px;line-height:1.9;color:var(--muted)">
+<b style="color:var(--text)">Default (0.4 / 0.3 / 0.3)</b> — balanced, works for unknown model quality<br>
+<b style="color:var(--text)">Grad-heavy (0.1 / 0.7 / 0.2)</b> — clean_prob &gt; 0.01. Grad signal is sharp, trust it.<br>
+<b style="color:var(--text)">Gate+Grad (0.4 / 0.4 / 0.2)</b> — recommended for smart_edit when causal IE is weak<br>
+<b style="color:var(--text)">Causal-heavy (0.2 / 0.2 / 0.6)</b> — only when clean_prob &gt; 0.1. IE is the gold signal then.<br>
+<b style="color:var(--text)">Gate-only (1.0 / 0.0 / 0.0)</b> — equivalent to basic locate(), sanity check<br>
+<br>
+<b style="color:var(--yellow)">Your distilgpt2 setting:</b> use (0.3 / 0.7 / 0.0) — gate+grad, skip causal (it's 0 anyway).
+</div>
+</div>
+</div>
+
 </div><!-- /app -->
 
 <script>
@@ -1459,6 +2032,7 @@ function showHmSlotDetail(layer, slot) {
 document.querySelectorAll('input[name="edit-mode"]').forEach(r=>{
   r.addEventListener('change', ()=>{
     document.getElementById('precise-prompt-row').style.display = r.value==='PRECISE'?'block':'none';
+    document.getElementById('smart-row').style.display          = r.value==='SMART'?'block':'none';
     document.getElementById('style-shift-row').style.display    = r.value==='STYLE-SHIFT'?'block':'none';
     document.getElementById('multiedit-row').style.display      = r.value==='MULTI-EDIT'?'block':'none';
   });
@@ -1500,6 +2074,30 @@ async function runEdit() {
     to_concept:  document.getElementById('ss-to').value,
     strength:    +document.getElementById('ss-strength').value,
   };
+  if(mode==='SMART'){
+    try {
+      const r = await api('/api/smart_edit', {
+        prompt:          document.getElementById('smart-prompt').value,
+        subject:         document.getElementById('edit-entity').value,
+        relation:        document.getElementById('edit-relation').value,
+        old_target:      document.getElementById('smart-old').value,
+        new_target:      document.getElementById('edit-new').value,
+        top_layers:      +document.getElementById('smart-layers').value,
+        slots_per_layer: +document.getElementById('smart-slots').value,
+        scale:           +document.getElementById('edit-scale').value,
+        noise_std:       0.1, alpha: 0.4, beta: 0.4, gamma: 0.2,
+      });
+      drawBeforeAfterChart(r.before, r.after);
+      let log = r.debug_log.join('\n');
+      log += '\n\nUsed layers:\n';
+      r.used_layers.forEach(l=>{ log+=`  L${l.layer}  slots=[${l.slots.join(',')}]  combined=${l.combined}\n`; });
+      log += '\nDelta:\n';
+      r.delta.slice(0,8).forEach(d=>{ log+=`  ${d.token}: ${d.before.toFixed(4)} → ${d.after.toFixed(4)}  ${d.delta>0?'+':''}${d.delta.toFixed(4)}\n`; });
+      document.getElementById('edit-log').textContent = log;
+      updatePatchCount();
+    } catch(e) { alert(e.message); }
+    return;
+  }
   if(mode==='MULTI-EDIT'){
     try {
       body.facts = JSON.parse(document.getElementById('multi-json').value);
@@ -1614,6 +2212,221 @@ async function updatePatchCount() {
   } catch(e){}
 }
 
+// ═══════════════════════════════════════════════
+// SMART LOCATE
+// ═══════════════════════════════════════════════
+let _slData = null;
+
+async function runSmartLocate() {
+  const st = document.getElementById('sl-status');
+  st.textContent = '⏳ Running gradient pass + causal sweep (may take ~20s for large models)…';
+  try {
+    const data = await api('/api/smart_locate', {
+      prompt:    document.getElementById('sl-prompt').value,
+      subject:   document.getElementById('sl-subject').value,
+      target:    document.getElementById('sl-target').value,
+      alpha:     +document.getElementById('sl-alpha').value,
+      beta:      +document.getElementById('sl-beta').value,
+      gamma:     +document.getElementById('sl-gamma').value,
+      noise_std: +document.getElementById('sl-noise').value,
+    });
+    _slData = data;
+    st.textContent = `✓ Done. clean_prob=${data.clean_prob.toFixed(4)}  corrupt_prob=${data.corrupt_prob.toFixed(4)}`;
+    drawSmartLocateChart(data.ranked_layers);
+    showSmartRec(data);
+    buildSlTable(data.ranked_layers);
+  } catch(e) { st.textContent = '✗ '+e.message; }
+}
+
+async function runGradientOnly() {
+  const st = document.getElementById('sl-status');
+  st.textContent = '⏳ Running gradient pass…';
+  try {
+    const data = await api('/api/gradient_scores', {
+      prompt: document.getElementById('sl-prompt').value,
+      target: document.getElementById('sl-target').value,
+    });
+    st.textContent = `✓ Gradient done. ${data.layer_scores.length} KB layers.`;
+    // Draw gradient-only bars
+    drawGradOnlyChart(data.layer_scores);
+  } catch(e) { st.textContent = '✗ '+e.message; }
+}
+
+async function runCausalOnly() {
+  const st = document.getElementById('sl-status');
+  st.textContent = '⏳ Running causal patch trace…';
+  try {
+    const data = await api('/api/causal_trace', {
+      prompt:    document.getElementById('sl-prompt').value,
+      subject:   document.getElementById('sl-subject').value,
+      target:    document.getElementById('sl-target').value,
+      noise_std: +document.getElementById('sl-noise').value,
+    });
+    st.textContent = `✓ Causal done. clean=${data.clean_prob.toFixed(4)} corrupt=${data.corrupt_prob.toFixed(4)}`;
+    drawCausalOnlyChart(data.results);
+  } catch(e) { st.textContent = '✗ '+e.message; }
+}
+
+async function runCollateral() {
+  const prompt = document.getElementById('sl-coll-prompt').value;
+  try {
+    const data = await api('/api/infer', { prompt, top_k: 5 });
+    const el = document.getElementById('sl-coll-out');
+    el.textContent = `"${prompt}"\n` +
+      data.results.map(r=>`  ${r.token.padEnd(18)} ${r.prob.toFixed(4)}`).join('\n');
+  } catch(e) { document.getElementById('sl-coll-out').textContent = '✗ '+e.message; }
+}
+
+function drawSmartLocateChart(ranked) {
+  // Sort by layer for chart display
+  const byLayer = [...ranked].sort((a,b)=>a.layer-b.layer);
+  const el = clearChart('sl-chart');
+  const W = el.clientWidth || 700, H = 40 + byLayer.length * 34;
+  el.style.height = H+'px';
+  const svg = d3.select(el).append('svg').attr('width','100%').attr('height',H);
+  const m = {left:50,right:110,top:20,bottom:20};
+  const w = W-m.left-m.right;
+  const g = svg.append('g').attr('transform',`translate(${m.left},${m.top})`);
+
+  // Each bar = 3 stacked segments (normalized: gate_sim_n, grad_norm_n, causal_n)
+  // Each segment width = signal_n * (w/3) so max of each is w/3
+  const segW = w / 3;
+
+  byLayer.forEach((d,i)=>{
+    const y = i*34;
+    // Label
+    g.append('text').attr('x',-6).attr('y',y+17).attr('text-anchor','end')
+     .attr('fill', d.layer===(_slData?.recommendation?.layer) ? C.yellow : C.muted)
+     .attr('font-size',10).text('L'+d.layer);
+
+    // gate_sim segment
+    g.append('rect').attr('x',0).attr('y',y+4).attr('width',d.gate_sim_n*segW)
+     .attr('height',12).attr('rx',2).attr('fill',C.blue).attr('opacity',.8)
+     .on('mousemove',(ev)=>showTooltip(`L${d.layer} gate_sim: ${d.gate_sim}`,ev.pageX,ev.pageY))
+     .on('mouseleave',hideTooltip);
+    // grad_norm segment
+    g.append('rect').attr('x',segW).attr('y',y+4).attr('width',d.grad_norm_n*segW)
+     .attr('height',12).attr('rx',2).attr('fill',C.green).attr('opacity',.8)
+     .on('mousemove',(ev)=>showTooltip(`L${d.layer} grad_norm: ${d.grad_norm}`,ev.pageX,ev.pageY))
+     .on('mouseleave',hideTooltip);
+    // causal segment
+    g.append('rect').attr('x',segW*2).attr('y',y+4).attr('width',d.causal_n*segW)
+     .attr('height',12).attr('rx',2).attr('fill',C.yellow).attr('opacity',.8)
+     .on('mousemove',(ev)=>showTooltip(`L${d.layer} causal_IE: ${d.causal_effect}`,ev.pageX,ev.pageY))
+     .on('mouseleave',hideTooltip);
+    // combined score label
+    g.append('text').attr('x',w+6).attr('y',y+14)
+     .attr('fill', d.combined===Math.max(...ranked.map(r=>r.combined)) ? C.yellow : C.muted)
+     .attr('font-size',10).text(d.combined.toFixed(3));
+  });
+
+  // Axis labels
+  const ax = g.append('g').attr('transform',`translate(0,${byLayer.length*34})`);
+  ax.append('text').attr('x',segW/2).attr('y',14).attr('text-anchor','middle')
+    .attr('fill',C.blue).attr('font-size',9).text('gate_sim');
+  ax.append('text').attr('x',segW*1.5).attr('y',14).attr('text-anchor','middle')
+    .attr('fill',C.green).attr('font-size',9).text('grad_norm');
+  ax.append('text').attr('x',segW*2.5).attr('y',14).attr('text-anchor','middle')
+    .attr('fill',C.yellow).attr('font-size',9).text('causal IE');
+
+  // Section dividers
+  [segW,segW*2].forEach(x=>{
+    g.append('line').attr('x1',x).attr('x2',x).attr('y1',0).attr('y2',byLayer.length*34)
+     .attr('stroke',C.border).attr('stroke-width',1).attr('stroke-dasharray','3,2');
+  });
+}
+
+function drawGradOnlyChart(layerScores) {
+  const el = clearChart('sl-chart');
+  const W = el.clientWidth || 700, H = 40 + layerScores.length * 28;
+  el.style.height = H+'px';
+  const svg = d3.select(el).append('svg').attr('width','100%').attr('height',H);
+  const m = {left:50,right:80,top:20,bottom:10};
+  const w = W-m.left-m.right;
+  const maxG = d3.max(layerScores, d=>d.max_grad) || 1;
+  const x = d3.scaleLinear().domain([0,maxG]).range([0,w]);
+  const g = svg.append('g').attr('transform',`translate(${m.left},${m.top})`);
+  layerScores.forEach((d,i)=>{
+    const y=i*28;
+    g.append('text').attr('x',-6).attr('y',y+14).attr('text-anchor','end')
+     .attr('fill',C.muted).attr('font-size',10).text('L'+d.layer);
+    g.append('rect').attr('x',0).attr('y',y+2).attr('width',x(d.max_grad))
+     .attr('height',16).attr('rx',2).attr('fill',C.green).attr('opacity',.8)
+     .on('mousemove',(ev)=>showTooltip(`L${d.layer} max_grad: ${d.max_grad}`,ev.pageX,ev.pageY))
+     .on('mouseleave',hideTooltip);
+    g.append('text').attr('x',x(d.max_grad)+4).attr('y',y+14)
+     .attr('fill',C.green).attr('font-size',9).text(d.max_grad.toExponential(2));
+  });
+}
+
+function drawCausalOnlyChart(results) {
+  const el = clearChart('sl-chart');
+  const W = el.clientWidth || 700, H = 40 + results.length * 28;
+  el.style.height = H+'px';
+  const svg = d3.select(el).append('svg').attr('width','100%').attr('height',H);
+  const m = {left:50,right:80,top:20,bottom:10};
+  const w = W-m.left-m.right;
+  const maxIE = Math.max(d3.max(results, d=>d.indirect_effect), 0.001);
+  const x = d3.scaleLinear().domain([0,maxIE]).range([0,w]);
+  const g = svg.append('g').attr('transform',`translate(${m.left},${m.top})`);
+  results.forEach((d,i)=>{
+    const y=i*28; const ie=Math.max(0,d.indirect_effect);
+    g.append('text').attr('x',-6).attr('y',y+14).attr('text-anchor','end')
+     .attr('fill',C.muted).attr('font-size',10).text('L'+d.layer);
+    g.append('rect').attr('x',0).attr('y',y+2).attr('width',x(ie))
+     .attr('height',16).attr('rx',2).attr('fill',C.yellow).attr('opacity',.8)
+     .on('mousemove',(ev)=>showTooltip(`L${d.layer} IE: ${d.indirect_effect}  patch_p: ${d.patch_prob}`,ev.pageX,ev.pageY))
+     .on('mouseleave',hideTooltip);
+    g.append('text').attr('x',x(ie)+4).attr('y',y+14)
+     .attr('fill',C.yellow).attr('font-size',9).text(d.indirect_effect.toFixed(5));
+  });
+}
+
+function showSmartRec(data) {
+  const rec = data.recommendation;
+  if(!rec){ document.getElementById('sl-rec').textContent='No recommendation.'; return; }
+  let txt = `★ Best layer: L${rec.layer}  combined=${rec.combined}\n\n`;
+  txt += `  gate_sim:      ${rec.gate_sim}  (norm ${rec.gate_sim_n})\n`;
+  txt += `  grad_norm:     ${rec.grad_norm}  (norm ${rec.grad_norm_n})\n`;
+  txt += `  causal_effect: ${rec.causal_effect}  (norm ${rec.causal_n})\n`;
+  if(rec.best_slots.length){
+    txt += `\nTop gradient slots in L${rec.layer}:\n`;
+    rec.best_slots.forEach(s=>{ txt+=`  slot ${s.slot}  grad_norm=${s.grad_norm}\n`; });
+  }
+  txt += `\nPhase layer (trace): L${data.phase_layer}\n`;
+  txt += `Subject pos: ${data.subject_pos}\n`;
+  txt += `clean_prob: ${data.clean_prob}  corrupt_prob: ${data.corrupt_prob}`;
+  document.getElementById('sl-rec').textContent = txt;
+}
+
+function buildSlTable(ranked) {
+  const el = document.getElementById('sl-table');
+  const maxC = Math.max(...ranked.map(r=>r.combined));
+  let html = `<table style="width:100%;border-collapse:collapse;font-size:11px">
+    <thead><tr style="color:var(--muted);border-bottom:1px solid var(--border)">
+      <th style="padding:4px 8px;text-align:left">Layer</th>
+      <th style="padding:4px 8px;text-align:right;color:${C.blue}">gate_sim</th>
+      <th style="padding:4px 8px;text-align:right;color:${C.green}">grad_norm</th>
+      <th style="padding:4px 8px;text-align:right;color:${C.yellow}">causal IE</th>
+      <th style="padding:4px 8px;text-align:right">combined ★</th>
+      <th style="padding:4px 8px;text-align:left;color:var(--muted)">top grad slots</th>
+    </tr></thead><tbody>`;
+  ranked.forEach(r=>{
+    const hi = r.combined===maxC ? `background:rgba(210,153,34,0.08)` : '';
+    const slots = r.best_slots.slice(0,3).map(s=>s.slot).join(', ');
+    html+=`<tr style="${hi};border-bottom:1px solid var(--border)">
+      <td style="padding:4px 8px;color:${r.combined===maxC?C.yellow:C.text}">L${r.layer}</td>
+      <td style="padding:4px 8px;text-align:right;color:${C.blue}">${r.gate_sim}</td>
+      <td style="padding:4px 8px;text-align:right;color:${C.green}">${r.grad_norm.toExponential(2)}</td>
+      <td style="padding:4px 8px;text-align:right;color:${C.yellow}">${r.causal_effect.toFixed(5)}</td>
+      <td style="padding:4px 8px;text-align:right;font-weight:700">${r.combined}</td>
+      <td style="padding:4px 8px;color:var(--muted)">${slots}</td>
+    </tr>`;
+  });
+  html += '</tbody></table>';
+  el.innerHTML = html;
+}
+
 // ═══════════════════════════════════════════════
 // INIT
 // ═══════════════════════════════════════════════
@@ -1672,6 +2485,39 @@ class HeatmapReq(BaseModel):
     use_activation: bool = False
     prompt: Optional[str] = None
 
+class GradientReq(BaseModel):
+    prompt: str
+    target: str
+
+class CausalTraceReq(BaseModel):
+    prompt: str
+    subject: str
+    target: str
+    noise_std: float = 0.1
+
+class SmartLocateReq(BaseModel):
+    prompt: str
+    subject: str
+    target: str
+    alpha: float = 0.4
+    beta: float = 0.3
+    gamma: float = 0.3
+    noise_std: float = 0.1
+
+class SmartEditReq(BaseModel):
+    prompt: str
+    subject: str
+    relation: str = ""
+    old_target: str
+    new_target: str
+    top_layers: int = 3
+    slots_per_layer: int = 2
+    scale: float = 1.5
+    noise_std: float = 0.1
+    alpha: float = 0.4
+    beta: float = 0.4
+    gamma: float = 0.2
+
 class DryRunReq(BaseModel):
     entity: str
     new_target: str
@@ -1778,6 +2624,54 @@ async def api_locate(req: LocateReq):
     return vi.locate(req.prompt, req.subject, req.target)
 
 
+@app.post("/api/gradient_scores")
+async def api_gradient_scores(req: GradientReq):
+    vi = _require()
+    return vi.gradient_slot_scores(req.prompt, req.target)
+
+
+@app.post("/api/causal_trace")
+async def api_causal_trace(req: CausalTraceReq):
+    vi = _require()
+    return vi.causal_patch_trace(req.prompt, req.subject, req.target,
+                                  noise_std=req.noise_std)
+
+
+@app.post("/api/smart_locate")
+async def api_smart_locate(req: SmartLocateReq):
+    vi = _require()
+    return vi.smart_locate(req.prompt, req.subject, req.target,
+                            alpha=req.alpha, beta=req.beta, gamma=req.gamma,
+                            noise_std=req.noise_std)
+
+
+@app.post("/api/smart_edit")
+async def api_smart_edit(req: SmartEditReq):
+    vi = _require()
+    prompt_str = req.prompt or f"The {req.relation} of {req.subject} is"
+    before = vi.infer(prompt_str, top_k=5)
+    log: List[str] = []
+    try:
+        result = vi.smart_edit(
+            prompt_str, req.subject, req.relation, req.old_target, req.new_target,
+            top_layers=req.top_layers, slots_per_layer=req.slots_per_layer,
+            scale=req.scale, noise_std=req.noise_std,
+            alpha=req.alpha, beta=req.beta, gamma=req.gamma, log=log
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    after = vi.infer(prompt_str, top_k=5)
+    b_map = {d["token"]: d["prob"] for d in before}
+    a_map = {d["token"]: d["prob"] for d in after}
+    all_toks = set(b_map) | set(a_map)
+    delta = sorted([{"token":t,"before":b_map.get(t,0),"after":a_map.get(t,0),
+                     "delta":a_map.get(t,0)-b_map.get(t,0)} for t in all_toks],
+                   key=lambda x: -abs(x["delta"]))
+    return {"before": before, "after": after, "delta": delta,
+            "debug_log": log, "used_layers": result["used_layers"],
+            "smart_locate": result["smart_locate"]}
+
+
 @app.post("/api/gate_heatmap")
 async def api_gate_heatmap(req: HeatmapReq):
     vi = _require()