fix: make KL optimization measured and reversible (#182)

README.md

@@ -116,7 +116,7 @@ OBLITERATUS implements several techniques that go beyond prior work:
 | **Parametric Kernel Optimization** | Bell-curve layer weighting with 7 global parameters via Optuna TPE search | Heretic-inspired |
 | **Refusal Direction Optimization (RDO)** | Gradient-based refinement of SVD-extracted directions using a linear refusal probe | Wollschlager et al., ICML 2025 |
 | **Float Direction Interpolation** | Continuous SVD direction index via Gaussian-shaped weighting for smoother refusal removal | Novel |
-| **KL-Divergence Co-Optimization** | Post-projection feedback loop that partially reverts over-projected layers if KL budget exceeded | Novel |
+| **KL-Divergence Co-Optimization** | Measures forward sequence-token KL against a bounded pristine prompt baseline and exactly restores damaging weak-signal layers until the configured budget is met | Novel |
 | **Component-Specific Scaling** | Separate attention vs MLP projection strengths (MLP layers are more sensitive) | Novel |
 | **LoRA-Based Reversible Ablation** | Rank-1 LoRA adapters instead of permanent weight surgery, enabling reversible ablation | Novel |
 | **Activation Winsorization** | Clamps activation vectors to percentile range before SVD to prevent outlier-dominated directions | Heretic-inspired |

app.py

@@ -1240,8 +1240,10 @@ def _format_obliteration_metrics(
         icon = "🟢" if ppl < 12 else "🟡" if ppl < 20 else "🔴"
         parts.append(f"| Perplexity | **{ppl:.2f}** | {icon} |")
     if kl is not None:
-        icon = "🟢" if kl < 0.05 else "🟡" if kl < 0.1 else "🔴"
-        parts.append(f"| KL Divergence | **{kl:.4f}** | {icon} |")
+        kl_budget = metrics.get("kl_budget", getattr(pipeline, "kl_budget", None))
+        icon = "🟢" if kl_budget is not None and kl <= kl_budget else "🔴"
+        budget_note = f" / {kl_budget:.4f}" if kl_budget is not None else ""
+        parts.append(f"| Token KL / Budget | **{kl:.4f}{budget_note}** | {icon} |")
     if n_layers > 0:
         parts.append(f"| Layers Modified | **{n_layers}** | |")
     if load_settings is not None:

ci/test-risk-map.json

@@ -75,6 +75,7 @@
       "required_tests": [
         "tests/test_abliterate.py",
         "tests/test_abliterate_extended.py",
+        "tests/test_kl_optimization_contracts.py",
         "tests/test_app_benchmark_lifecycle.py",
         "tests/test_qwen35_contracts.py",
         "tests/test_auto_obliterate.py",

obliteratus/abliterate.py

@@ -1150,10 +1150,15 @@ class AbliterationPipeline:
         self._expert_directions: dict[int, dict[int, torch.Tensor]] = {}
         # Layer-adaptive projection weights (layer → scale 0..1)
         self._layer_excise_weights: dict[int, float] = {}
+        self._refusal_strengths: dict[int, float] = {}
         # SAE-derived refusal directions (layer → tensor of shape (n_features, hidden))
         self._sae_directions: dict[int, torch.Tensor] = {}
         # Pre-EXCISE first-token logits for KL divergence in VERIFY
         self._baseline_first_token_logits: torch.Tensor | None = None
+        # Full prompt-token logits used by the KL optimizer and final verifier.
+        # Kept on CPU and deliberately bounded by _capture_baseline_kl_logits.
+        self._baseline_token_logits: list[torch.Tensor] = []
+        self._last_first_token_kl: float | None = None
         self._kl_eval_prompts: list[str] = []
         # Attention head refusal attribution (layer → list of (head_idx, score))
         self._refusal_heads: dict[int, list[tuple[int, float]]] = {}
@@ -2631,6 +2636,7 @@ class AbliterationPipeline:
         import math
         sorted_layers = [(idx, n) for idx, n in sorted_layers
                          if not (math.isnan(n) or math.isinf(n))]
+        self._refusal_strengths = dict(sorted_layers)
         max_norm = sorted_layers[0][1] if sorted_layers else 1.0
         if math.isnan(max_norm) or math.isinf(max_norm) or max_norm <= 0:
             max_norm = 1.0
@@ -3882,23 +3888,28 @@ class AbliterationPipeline:
     # ── Pre-EXCISE baseline capture for KL divergence ──────────────────
 
     def _capture_baseline_kl_logits(self):
-        """Capture first-token logits on harmless prompts before EXCISE.
+        """Capture prompt-token logits on harmless prompts before EXCISE.
 
-        These are compared against post-EXCISE logits in _verify() to compute
-        first-token KL divergence — the standard metric used by Heretic and
-        Young (2025) for measuring collateral damage from abliteration.
+        These are compared against post-EXCISE logits in optimization and
+        verification to compute forward sequence-token KL. The last input-token
+        slice is retained separately as a compatibility diagnostic.
 
-        Uses chat template (matching PROBE stage formatting) and padding-aware
-        indexing to extract logits at the last real token per sequence.
+        Uses the same chat-template and tokenization contract as PROBE.
         """
         model = self.handle.model
         tokenizer = self.handle.tokenizer
         device = self._get_model_device(model)
 
-        # Use a subset of harmless prompts (100 is the Heretic standard)
-        raw_prompts = self.harmless_prompts[:100]
+        # Full-vocabulary token logits are large for modern models. Sixteen
+        # deterministic prompts provide a bounded measurement set while still
+        # covering substantially more evidence than a first-token-only proxy.
+        raw_prompts = self.harmless_prompts[:16]
         if len(raw_prompts) < 10:
             self.log("Skipping baseline KL capture (too few harmless prompts)")
+            if self.use_kl_optimization:
+                raise RuntimeError(
+                    "KL optimization requires at least 10 harmless baseline prompts"
+                )
             return
 
         # Apply chat template for consistency with how the model was probed
@@ -3906,33 +3917,39 @@ class AbliterationPipeline:
 
         self.log(f"Capturing baseline logits on {len(self._kl_eval_prompts)} harmless prompts for KL...")
         all_first_logits = []
-        batch_size = 8
+        all_token_logits = []
 
         try:
-            for i in range(0, len(self._kl_eval_prompts), batch_size):
-                batch = self._kl_eval_prompts[i:i + batch_size]
+            for prompt in self._kl_eval_prompts:
                 inputs = tokenizer(
-                    batch, return_tensors="pt",
-                    padding=True, truncation=True, max_length=self.max_seq_length or 256,
+                    prompt, return_tensors="pt",
+                    truncation=True, max_length=self.max_seq_length or 256,
                 )
                 inputs = {k: v.to(device) for k, v in inputs.items()}
 
                 with torch.no_grad():
                     logits = model(**inputs).logits
-                    # Padding-aware: extract logits at last REAL token per sequence
-                    attn_mask = inputs["attention_mask"]
-                    last_idx = attn_mask.sum(dim=1) - 1  # (batch,)
-                    batch_range = torch.arange(logits.shape[0], device=device)
-                    first_logits = logits[batch_range, last_idx].cpu()
-                    all_first_logits.append(first_logits)
+                    prompt_logits = logits[0].detach().cpu()
+                    all_token_logits.append(prompt_logits)
+                    all_first_logits.append(prompt_logits[-1:].clone())
 
                 del inputs, logits
 
             self._baseline_first_token_logits = torch.cat(all_first_logits, dim=0)
+            self._baseline_token_logits = all_token_logits
             self.log(f"  Captured baseline logits: {self._baseline_first_token_logits.shape}")
+            self.log(
+                f"  Captured sequence-token KL baseline: "
+                f"{sum(item.shape[0] for item in all_token_logits)} tokens"
+            )
         except Exception as e:
-            self.log(f"  Baseline KL capture failed (non-fatal): {e}")
+            self.log(f"  Baseline KL capture failed: {e}")
             self._baseline_first_token_logits = None
+            self._baseline_token_logits = []
+            if self.use_kl_optimization:
+                raise RuntimeError(
+                    "KL optimization requires a finite, shape-compatible pristine baseline"
+                ) from e
 
         self._free_gpu_memory()
 
@@ -3962,6 +3979,30 @@ class AbliterationPipeline:
         arch = self.handle.architecture
         config = self.handle.config
 
+        # Exact KL rollback is intentionally bounded to the weakest third of
+        # selected refusal layers. Snapshot before any pass so restoration is
+        # byte-for-byte pristine rather than an approximate inverse projection.
+        kl_pristine_snapshots: dict[int, dict[str, torch.Tensor]] = {}
+        if self.use_kl_optimization and self._strong_layers:
+            candidate_count = max(1, len(self._strong_layers) // 3)
+            ranked = sorted(
+                self._strong_layers,
+                key=lambda idx: (
+                    self._layer_excise_weights.get(idx, 1.0),
+                    self._refusal_strengths.get(idx, float("inf")),
+                    idx,
+                ),
+            )
+            candidate_layers = ranked[:candidate_count]
+            self.log(
+                f"Capturing exact KL rollback state for {len(candidate_layers)} "
+                "weak-signal layers..."
+            )
+            kl_pristine_snapshots = {
+                idx: self._snapshot_layer_state(layers[idx])
+                for idx in candidate_layers
+            }
+
         text_cfg = getattr(config, "text_config", None)
         n_heads = (
             getattr(config, "num_attention_heads", None)
@@ -3974,11 +4015,26 @@ class AbliterationPipeline:
         grad_ctx = torch.no_grad()
         grad_ctx.__enter__()
         try:
-            self._excise_inner(layers, arch, config, n_heads, t0)
+            self._excise_inner(
+                layers,
+                arch,
+                config,
+                n_heads,
+                t0,
+                kl_pristine_snapshots,
+            )
         finally:
             grad_ctx.__exit__(None, None, None)
 
-    def _excise_inner(self, layers, arch, config, n_heads, t0):
+    def _excise_inner(
+        self,
+        layers,
+        arch,
+        config,
+        n_heads,
+        t0,
+        kl_pristine_snapshots,
+    ):
         """Inner excise logic, called within torch.no_grad() context."""
         total_modified = 0
         total_neurons_masked = 0
@@ -4600,19 +4656,22 @@ class AbliterationPipeline:
         # ── KL-divergence co-optimization ──────────────────────────────
         # Inspired by Heretic's Bayesian optimization approach, but
         # implemented as a post-projection feedback loop rather than a
-        # search-based method.  Measures KL divergence on harmless prompts
-        # after each refinement pass and compensates over-projected layers.
+        # search-based method. Measures sequence-token KL on harmless prompts
+        # after refinement and restores over-damaging weak-signal layers.
         #
         # Algorithm:
         # 1. Run a small forward pass on harmless reference prompts
-        # 2. Compute per-layer KL divergence contribution
-        # 3. If total KL exceeds budget, identify the worst layers and
-        #    partially revert their projection (additive correction)
+        # 2. Measure exact marginal KL reduction for rollback candidates
+        # 3. If total KL exceeds budget, restore pristine candidates greedily
         #
         # This is NOVEL: Heretic optimizes KL during ablation via search;
-        # we optimize via post-hoc correction with minimal compute overhead.
+        # we optimize via measured post-hoc correction.
         if self.use_kl_optimization and self.handle and self._strong_layers:
-            self._kl_optimize_corrections(layers, total_modified)
+            self._kl_optimize_corrections(
+                layers,
+                total_modified,
+                kl_pristine_snapshots,
+            )
 
         # ── lm_head projection ────────────────────────────────────────
         # The language model head converts hidden states to token logits.
@@ -4941,6 +5000,7 @@ class AbliterationPipeline:
                 norms[idx] = (S[:k] ** 2).sum().item()
 
         sorted_layers = sorted(norms.items(), key=lambda x: x[1], reverse=True)
+        self._refusal_strengths = dict(sorted_layers)
 
         # Respect configured layer_selection (matching _distill)
         selection_method = self.layer_selection
@@ -5024,219 +5084,156 @@ class AbliterationPipeline:
             self._expert_safety_scores.clear()
             self._identify_safety_experts()
 
-    def _kl_optimize_corrections(self, layers: nn.ModuleList, total_modified: int):
-        """KL-divergence co-optimization: measure and correct over-projection.
+    @staticmethod
+    def _snapshot_layer_state(layer: nn.Module) -> dict[str, torch.Tensor]:
+        """Return an exact CPU snapshot suitable for deterministic rollback."""
+        snapshot: dict[str, torch.Tensor] = {}
+        for name, value in layer.state_dict().items():
+            if value.device.type == "meta":
+                raise RuntimeError(f"cannot snapshot meta tensor {name!r} for KL rollback")
+            snapshot[name] = value.detach().cpu().clone()
+        if not snapshot:
+            raise RuntimeError("cannot optimize KL without a non-empty layer snapshot")
+        return snapshot
 
-        Measures per-layer KL divergence contribution on harmless reference
-        prompts and partially reverts projections that caused excessive KL.
+    @staticmethod
+    def _restore_layer_state(layer: nn.Module, snapshot: dict[str, torch.Tensor]) -> None:
+        """Restore a snapshot with strict key, shape, and dtype validation."""
+        current = layer.state_dict(keep_vars=True)
+        if current.keys() != snapshot.keys():
+            raise RuntimeError("KL rollback layer state keys changed after mutation")
+        with torch.no_grad():
+            for name, pristine in snapshot.items():
+                target = current[name]
+                if target.device.type == "meta":
+                    raise RuntimeError(f"cannot restore meta tensor {name!r} during KL rollback")
+                if target.shape != pristine.shape or target.dtype != pristine.dtype:
+                    raise RuntimeError(
+                        f"KL rollback tensor contract changed for {name!r}: "
+                        f"{tuple(target.shape)}/{target.dtype} != "
+                        f"{tuple(pristine.shape)}/{pristine.dtype}"
+                    )
+                target.copy_(pristine.to(device=target.device))
 
-        This bridges Heretic's Bayesian optimization approach (co-minimize
-        refusal + KL) with OBLITERATUS's analysis-based pipeline:
-        - Heretic: search-based (200 Optuna trials, expensive)
-        - OBLITERATUS: project first, measure KL, correct surgically (cheap)
+    def _measure_sequence_token_kl(self) -> float:
+        """Measure mean KL(P_pristine || Q_current) across prompt tokens."""
+        if not self._baseline_token_logits or not self._kl_eval_prompts:
+            raise RuntimeError("sequence-token KL baseline is unavailable")
+        if len(self._baseline_token_logits) != len(self._kl_eval_prompts):
+            raise RuntimeError("sequence-token KL baseline prompt count changed")
 
-        The correction is additive: for over-projected layers, we add back
-        a fraction of the removed refusal component to reduce KL divergence
-        while keeping most of the refusal removal intact.
-        """
         model = self.handle.model
         tokenizer = self.handle.tokenizer
         device = self._get_model_device(model)
-        arch = self.handle.architecture
-
-        # Reference prompts for KL measurement (harmless only)
-        kl_prompts = [
-            "The weather today is",
-            "In mathematics, the concept of",
-            "The history of ancient Rome",
-            "How to cook a simple pasta",
-            "The process of photosynthesis involves",
-        ]
-
-        self.log("KL-divergence co-optimization: measuring capability damage...")
-
-        # Collect pre-correction logits
-        all_logits = []
-        try:
-            for prompt in kl_prompts:
-                inputs = tokenizer(
-                    prompt, return_tensors="pt", truncation=True, max_length=self.max_seq_length or 64,
-                )
-                inputs = {k: v.to(device) for k, v in inputs.items()}
-                with torch.no_grad():
-                    outputs = model(**inputs)
-                    all_logits.append(outputs.logits[:, -1, :].detach().cpu().float())
-                del inputs, outputs
-        except Exception as e:
-            self.log(f"  KL optimization skipped (forward pass failed: {e})")
-            return
-
-        if not all_logits:
-            return
-
-        # Compute per-layer KL contribution by temporarily removing each
-        # layer's projection and measuring the change.  This is expensive
-        # with the full model, so we use an approximation: the projection
-        # magnitude as a proxy for KL contribution.
-        layer_kl_proxy: dict[int, float] = {}
-        for idx in self._strong_layers:
-            if idx not in self.refusal_directions:
-                continue
-            d = self.refusal_directions[idx]
-
-            # Proxy: mean absolute projection of refusal direction onto weight
-            # matrices at this layer.  Larger projection = more modification = more KL.
-            total_proj = 0.0
-            n_proj = 0
-            try:
-                attn = get_attention_module(layers[idx], arch)
-                for name in _ATTN_OUT_NAMES:
-                    W = getattr(attn, name, None)
-                    if W is not None and hasattr(W, "weight"):
-                        d_dev = d.to(device=W.weight.device, dtype=W.weight.dtype)
-                        if W.weight.shape[-1] == d_dev.shape[0]:
-                            proj_mag = (W.weight.data @ d_dev).abs().mean().item()
-                        elif W.weight.shape[0] == d_dev.shape[0]:
-                            proj_mag = (d_dev @ W.weight.data).abs().mean().item()
-                        else:
-                            continue
-                        total_proj += proj_mag
-                        n_proj += 1
-            except (AttributeError, RuntimeError):
-                pass
-            try:
-                ffn = get_ffn_module(layers[idx], arch)
-                for name in _FFN_OUT_NAMES:
-                    W = getattr(ffn, name, None)
-                    if W is not None and hasattr(W, "weight"):
-                        d_dev = d.to(device=W.weight.device, dtype=W.weight.dtype)
-                        if W.weight.shape[-1] == d_dev.shape[0]:
-                            proj_mag = (W.weight.data @ d_dev).abs().mean().item()
-                        elif W.weight.shape[0] == d_dev.shape[0]:
-                            proj_mag = (d_dev @ W.weight.data).abs().mean().item()
-                        else:
-                            continue
-                        total_proj += proj_mag
-                        n_proj += 1
-            except (AttributeError, RuntimeError):
-                pass
-
-            avg_proj = total_proj / max(n_proj, 1)
-            layer_kl_proxy[idx] = avg_proj
-            self._kl_contributions[idx] = avg_proj
-
-        if not layer_kl_proxy:
-            return
-
-        # Compute total loss (perplexity) as KL proxy
-        total_loss = 0.0
-        n_tokens = 0
-        try:
-            for prompt in kl_prompts[:3]:
-                inputs = tokenizer(
-                    prompt, return_tensors="pt", truncation=True, max_length=self.max_seq_length or 64,
-                )
-                inputs = {k: v.to(device) for k, v in inputs.items()}
-                with torch.no_grad():
-                    outputs = model(**inputs, labels=inputs["input_ids"])
-                    loss_val = outputs.loss.item()
-                    if not math.isnan(loss_val) and not math.isinf(loss_val):
-                        total_loss += loss_val * inputs["input_ids"].shape[1]
-                        n_tokens += inputs["input_ids"].shape[1]
-                del inputs, outputs
-        except Exception:
-            pass
-
-        if n_tokens > 0:
-            avg_loss = total_loss / n_tokens
-            try:
-                current_ppl = math.exp(min(avg_loss, 100.0))
-            except OverflowError:
-                current_ppl = float("inf")
-        else:
-            current_ppl = float("inf")
-
-        # KL budget check: if perplexity exceeds budget threshold, correct.
-        # Map kl_budget (0.0-2.0+) to a perplexity ceiling via exp scale so
-        # the full range is usable: 0.1→8, 0.3→13, 0.5→22, 1.0→55, 2.0→403
-        ppl_budget = math.exp(self.kl_budget * 3.0 + 1.0)
-        self.log(f"  Current perplexity: {current_ppl:.2f} (budget ceiling: {ppl_budget:.0f})")
-
-        if current_ppl > ppl_budget and current_ppl != float("inf"):
-            self.log("  KL budget exceeded — applying correction to weakest layers...")
-
-            # Sort layers by KL proxy (highest first = most damaging)
-            sorted_kl = sorted(layer_kl_proxy.items(), key=lambda x: x[1], reverse=True)
-
-            # Partially revert the weakest-signal layers (bottom third)
-            n_to_correct = max(1, len(sorted_kl) // 3)
-            correction_layers = [idx for idx, _ in sorted_kl[-n_to_correct:]]
-
-            for idx in correction_layers:
-                if idx not in self.refusal_directions:
-                    continue
-                d = self.refusal_directions[idx]
-
-                # Add back 30% of the removed refusal component.
-                #
-                # After full projection (reg=0), W_proj @ d = 0, so computing
-                # the revert from the current weights gives zero.  Instead we
-                # use the stored per-layer KL proxy (mean projection magnitude
-                # before excision) as a scale factor.  The revert adds back a
-                # fraction of the rank-1 refusal component: scale * d @ d^T
-                # applied in the appropriate orientation for each weight matrix.
-                revert_strength = 0.30
-                kl_proxy_mag = self._kl_contributions.get(idx, 0.0)
-                d_col = d.unsqueeze(-1) if d.dim() == 1 else d
-
-                def _partial_revert(module, weight_names, proxy_mag):
-                    for name in weight_names:
-                        proj = getattr(module, name, None)
-                        if proj is not None and hasattr(proj, "weight"):
-                            W = proj.weight.data
-                            d_dev = d_col.to(device=W.device, dtype=W.dtype)
-                            if W.shape[-1] == d_dev.shape[0]:
-                                # W is (out, hidden), d_dev is (hidden, 1)
-                                coeff = W @ d_dev  # (out, 1)
-                                coeff_mag = coeff.abs().mean().item()
-                                if coeff_mag < 1e-6 and proxy_mag > 0:
-                                    # Post-projection coeff ≈ 0, use proxy magnitude.
-                                    # Add uniform d^T to each row, scaled by proxy.
-                                    # d_dev.T is (1, hidden), broadcasts to (out, hidden)
-                                    W.add_(revert_strength * proxy_mag * d_dev.T)
-                                else:
-                                    # coeff is (out, 1), d_dev.T is (1, hidden)
-                                    # broadcasts to (out, hidden) — correct rank-1
-                                    W.add_(d_dev.T * (revert_strength * coeff))
-                            elif W.shape[0] == d_dev.shape[0]:
-                                # W is (hidden, out), d_row is (1, hidden)
-                                d_row = d_dev.squeeze(-1).unsqueeze(0)
-                                coeff = d_row @ W  # (1, out)
-                                coeff_mag = coeff.abs().mean().item()
-                                if coeff_mag < 1e-6 and proxy_mag > 0:
-                                    # d_row.T is (hidden, 1), broadcasts to (hidden, out)
-                                    W.add_(revert_strength * proxy_mag * d_row.T)
-                                else:
-                                    # d_row.T is (hidden, 1), coeff is (1, out)
-                                    W.add_(revert_strength * (d_row.T @ coeff))
-
-                try:
-                    attn = get_attention_module(layers[idx], arch)
-                    _partial_revert(attn, _ATTN_OUT_NAMES, kl_proxy_mag)
-                except (AttributeError, RuntimeError):
-                    pass
-                try:
-                    ffn = get_ffn_module(layers[idx], arch)
-                    _partial_revert(ffn, _FFN_OUT_NAMES, kl_proxy_mag)
-                except (AttributeError, RuntimeError):
-                    pass
-
-            self.log(
-                f"  Corrected {len(correction_layers)} layers "
-                f"(reverted {revert_strength:.0%} of projection)"
+        total_kl = 0.0
+        total_tokens = 0
+        first_token_kl = 0.0
+        first_token_count = 0
+        for prompt, pristine in zip(
+            self._kl_eval_prompts, self._baseline_token_logits, strict=True,
+        ):
+            inputs = tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=self.max_seq_length or 256,
             )
-        else:
+            inputs = {key: value.to(device) for key, value in inputs.items()}
+            with torch.no_grad():
+                current = model(**inputs).logits[0].detach().cpu().float()
+            pristine_float = pristine.float()
+            if current.shape != pristine_float.shape:
+                raise RuntimeError(
+                    "sequence-token KL logits changed shape: "
+                    f"{tuple(current.shape)} != {tuple(pristine_float.shape)}"
+                )
+            if not torch.isfinite(current).all() or not torch.isfinite(pristine_float).all():
+                raise RuntimeError("sequence-token KL logits contain non-finite values")
+            log_p = torch.nn.functional.log_softmax(pristine_float, dim=-1)
+            log_q = torch.nn.functional.log_softmax(current, dim=-1)
+            per_token = torch.nn.functional.kl_div(
+                log_q, log_p, log_target=True, reduction="none",
+            ).sum(dim=-1).clamp(min=0.0)
+            total_kl += per_token.sum().item()
+            total_tokens += per_token.numel()
+            first_token_kl += per_token[-1].item()
+            first_token_count += 1
+            del inputs, current, pristine_float, log_p, log_q, per_token
+        if total_tokens == 0:
+            raise RuntimeError("sequence-token KL baseline contains no tokens")
+        measured = total_kl / total_tokens
+        if not math.isfinite(measured):
+            raise RuntimeError("sequence-token KL measurement is non-finite")
+        self._last_first_token_kl = first_token_kl / first_token_count
+        return measured
+
+    def _kl_optimize_corrections(
+        self,
+        layers: nn.ModuleList,
+        total_modified: int,
+        pristine_snapshots: dict[int, dict[str, torch.Tensor]],
+    ):
+        """Measure sequence-token KL and exactly restore damaging layers.
+
+        Candidate layers are restricted to the weakest refusal-signal third.
+        Each candidate is restored temporarily from an exact pristine snapshot,
+        scored by measured marginal KL reduction, and returned to its modified
+        state. The best candidates are then restored exactly until the budget
+        is met. Missing or insufficient evidence fails closed before saving.
+        """
+        del total_modified  # retained in the signature for compatibility
+        if not pristine_snapshots:
+            raise RuntimeError("KL optimization has no exact rollback candidates")
+
+        self.log("Sequence-token KL co-optimization: measuring capability damage...")
+        current_kl = self._measure_sequence_token_kl()
+        self.log(f"  Current token KL: {current_kl:.4f} (budget: {self.kl_budget:.4f})")
+        if current_kl <= self.kl_budget:
             self.log("  KL within budget — no correction needed")
+            self._quality_metrics["kl_divergence"] = current_kl
+            self._quality_metrics["kl_budget"] = self.kl_budget
+            return
+
+        self.log("  KL budget exceeded — measuring exact layer rollback benefit...")
+        marginal: dict[int, float] = {}
+        for idx, pristine in pristine_snapshots.items():
+            modified = self._snapshot_layer_state(layers[idx])
+            self._restore_layer_state(layers[idx], pristine)
+            reverted_kl = self._measure_sequence_token_kl()
+            self._restore_layer_state(layers[idx], modified)
+            marginal[idx] = current_kl - reverted_kl
+            self._kl_contributions[idx] = marginal[idx]
+            self.log(f"    layer {idx}: marginal token KL reduction={marginal[idx]:+.4f}")
+
+        corrected: list[int] = []
+        for idx, benefit in sorted(marginal.items(), key=lambda item: item[1], reverse=True):
+            if benefit <= 0:
+                continue
+            modified = self._snapshot_layer_state(layers[idx])
+            self._restore_layer_state(layers[idx], pristine_snapshots[idx])
+            measured = self._measure_sequence_token_kl()
+            if measured >= current_kl:
+                self._restore_layer_state(layers[idx], modified)
+                continue
+            current_kl = measured
+            corrected.append(idx)
+            self.log(f"    restored layer {idx}: token KL={current_kl:.4f}")
+            if current_kl <= self.kl_budget:
+                break
+
+        self._quality_metrics["kl_divergence"] = current_kl
+        self._quality_metrics["kl_budget"] = self.kl_budget
+        if current_kl > self.kl_budget:
+            raise RuntimeError(
+                "KL optimization could not satisfy the configured sequence-token budget "
+                f"({current_kl:.4f} > {self.kl_budget:.4f}) using "
+                f"{len(pristine_snapshots)} exact rollback candidates"
+            )
+        corrected_set = set(corrected)
+        self._strong_layers = [
+            idx for idx in self._strong_layers if idx not in corrected_set
+        ]
+        self.log(f"  Corrected {len(corrected)} layers by exact pristine restoration")
 
         self._free_gpu_memory()
 
@@ -7212,67 +7209,35 @@ class AbliterationPipeline:
             self._quality_metrics["refusal_rate"] = None
             self.log("  Refusal rate: skipped (insufficient GPU memory for generation)")
 
-        # 4. First-token KL divergence (Heretic/Young standard metric)
+        # 4. Sequence-token KL divergence against the identical pristine
+        # prompt/token set used by optimization. First-token KL is retained as
+        # a secondary compatibility diagnostic, never as the configured gate.
         kl_divergence = None
-        if self._baseline_first_token_logits is not None and len(self._kl_eval_prompts) > 0:
-            self.log("Computing first-token KL divergence vs. baseline...")
+        if self._baseline_token_logits and len(self._kl_eval_prompts) > 0:
+            self.log("Computing sequence-token KL divergence vs. baseline...")
             try:
-                all_post_logits = []
-                for i in range(0, len(self._kl_eval_prompts), 8):
-                    batch = self._kl_eval_prompts[i:i + 8]
-                    inputs = tokenizer(
-                        batch, return_tensors="pt",
-                        padding=True, truncation=True, max_length=self.max_seq_length or 256,
-                    )
-                    inputs = {k: v.to(device) for k, v in inputs.items()}
-                    with torch.no_grad():
-                        logits = model(**inputs).logits
-                        # Padding-aware: extract at last real token position
-                        attn_mask = inputs["attention_mask"]
-                        last_idx = attn_mask.sum(dim=1) - 1
-                        batch_range = torch.arange(logits.shape[0], device=device)
-                        all_post_logits.append(logits[batch_range, last_idx].cpu())
-                    del inputs, logits
-                self._free_gpu_memory()
+                kl_divergence = self._measure_sequence_token_kl()
+                self._quality_metrics["kl_divergence"] = kl_divergence
+                self._quality_metrics["kl_budget"] = self.kl_budget
+                self._quality_metrics["kl_metric"] = "sequence_token_forward_kl_nats"
+                budget_label = "within budget" if kl_divergence <= self.kl_budget else "over budget"
+                self.log(
+                    f"  Sequence-token KL divergence: {kl_divergence:.4f} "
+                    f"({budget_label}; budget={self.kl_budget:.4f})"
+                )
 
-                post_logits = torch.cat(all_post_logits, dim=0)
-                pre_logits = self._baseline_first_token_logits[:post_logits.shape[0]]
-
-                # Check for NaN/Inf in post-ablation logits (model may be broken)
-                if torch.isnan(post_logits).any() or torch.isinf(post_logits).any():
-                    self.log("  KL divergence: inf (model produces NaN/Inf logits — weights may be destroyed)")
-                    kl_divergence = float("inf")
-                    self._quality_metrics["kl_divergence"] = kl_divergence
-                else:
-                    # Use F.kl_div for numerical stability
-                    log_p = torch.nn.functional.log_softmax(pre_logits.float(), dim=-1)
-                    log_q = torch.nn.functional.log_softmax(post_logits.float(), dim=-1)
-                    kl_per_prompt = torch.nn.functional.kl_div(
-                        log_q, log_p, log_target=True, reduction="none"
-                    ).sum(dim=-1).clamp(min=0.0)
-                    kl_divergence = kl_per_prompt.mean().item()
-
-                    # Guard against NaN from numerical issues in KL computation
-                    if math.isnan(kl_divergence) or math.isinf(kl_divergence):
-                        kl_divergence = float("inf")
-                        self.log("  First-token KL divergence: inf (numerical overflow — model may be severely damaged)")
-                    else:
-                        if kl_divergence < 0.2:
-                            kl_label = "excellent"
-                        elif kl_divergence < 0.5:
-                            kl_label = "good"
-                        elif kl_divergence < 1.0:
-                            kl_label = "moderate"
-                        else:
-                            kl_label = "high"
-                        self.log(f"  First-token KL divergence: {kl_divergence:.4f} ({kl_label})")
-                    self._quality_metrics["kl_divergence"] = kl_divergence
+                first_token_kl = self._last_first_token_kl
+                self._quality_metrics["first_token_kl_divergence"] = first_token_kl
+                self.log(f"  First-token KL divergence: {first_token_kl:.4f} (diagnostic)")
             except Exception as e:
-                self.log(f"  KL divergence computation failed (non-fatal): {e}")
+                self.log(f"  KL divergence computation failed: {e}")
+                if self.use_kl_optimization:
+                    raise RuntimeError("configured KL verification failed closed") from e
                 self._quality_metrics["kl_divergence"] = None
 
             # Free KL artifacts
             self._baseline_first_token_logits = None
+            self._baseline_token_logits = []
             self._kl_eval_prompts = []
         else:
             self._quality_metrics["kl_divergence"] = None
@@ -7398,6 +7363,8 @@ class AbliterationPipeline:
                 "float_layer_interpolation": self.float_layer_interpolation,
                 "cot_aware": self.cot_aware,
                 "use_kl_optimization": self.use_kl_optimization,
+                "kl_budget": self.kl_budget,
+                "kl_metric": "sequence_token_forward_kl_nats",
                 "use_lora_ablation": self.use_lora_ablation,
                 "som_iterations": self.som_iterations if self.direction_method == "som" else None,
                 "som_learning_rate": self.som_learning_rate if self.direction_method == "som" else None,

tests/test_app_model_load_settings.py

@@ -58,3 +58,40 @@ else:
         check=False,
     )
     assert result.returncode == 0, result.stdout + result.stderr
+
+
+@pytest.mark.operator_ui
+def test_result_card_uses_configured_sequence_token_kl_budget():
+    """The UI must not apply independent hard-coded KL thresholds."""
+    script = r'''
+from types import SimpleNamespace
+
+import app
+
+pipeline = SimpleNamespace(
+    _quality_metrics={
+        "kl_divergence": 0.24,
+        "kl_budget": 0.50,
+        "kl_metric": "sequence_token_forward_kl_nats",
+    },
+    _strong_layers=[1, 2],
+    kl_budget=0.50,
+)
+card = app._format_obliteration_metrics(pipeline, "advanced", "1s")
+assert "Token KL / Budget" in card
+assert "0.2400 / 0.5000" in card
+assert "🟢" in card
+
+pipeline._quality_metrics["kl_divergence"] = 0.51
+card = app._format_obliteration_metrics(pipeline, "advanced", "1s")
+assert "0.5100 / 0.5000" in card
+assert "🔴" in card
+'''
+    result = subprocess.run(
+        [sys.executable, "-c", script],
+        capture_output=True,
+        text=True,
+        timeout=60,
+        check=False,
+    )
+    assert result.returncode == 0, result.stdout + result.stderr

tests/test_kl_optimization_contracts.py

@@ -0,0 +1,159 @@
+"""Contracts for measured KL optimization and exact rollback."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from obliteratus.abliterate import AbliterationPipeline
+
+
+class _Tokenizer:
+    def __call__(self, prompt, **_kwargs):
+        length = 2 if prompt == "short" else 3
+        return {
+            "input_ids": torch.arange(length).unsqueeze(0),
+            "attention_mask": torch.ones(1, length, dtype=torch.long),
+        }
+
+
+class _LogitModel(nn.Module):
+    def __init__(self, logits_by_length):
+        super().__init__()
+        self.anchor = nn.Parameter(torch.zeros(1))
+        self.logits_by_length = logits_by_length
+
+    def forward(self, input_ids, **_kwargs):
+        logits = self.logits_by_length[input_ids.shape[1]].unsqueeze(0)
+        return SimpleNamespace(logits=logits)
+
+
+def _bare_pipeline() -> AbliterationPipeline:
+    pipeline = AbliterationPipeline.__new__(AbliterationPipeline)
+    pipeline.max_seq_length = 32
+    pipeline._quality_metrics = {}
+    pipeline._kl_contributions = {}
+    pipeline._strong_layers = []
+    pipeline._free_gpu_memory = lambda: None
+    pipeline.log = lambda _message: None
+    return pipeline
+
+
+def test_sequence_token_kl_matches_hand_computed_forward_kl():
+    pristine_short = torch.tensor([[2.0, 0.0], [0.0, 2.0]])
+    pristine_long = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
+    current_short = torch.tensor([[1.0, 1.0], [0.5, 1.5]])
+    current_long = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
+
+    pipeline = _bare_pipeline()
+    pipeline.handle = SimpleNamespace(
+        model=_LogitModel({2: current_short, 3: current_long}),
+        tokenizer=_Tokenizer(),
+    )
+    pipeline._kl_eval_prompts = ["short", "long"]
+    pipeline._baseline_token_logits = [pristine_short, pristine_long]
+
+    expected_parts = []
+    for pristine, current in (
+        (pristine_short, current_short),
+        (pristine_long, current_long),
+    ):
+        log_p = torch.log_softmax(pristine, dim=-1)
+        log_q = torch.log_softmax(current, dim=-1)
+        expected_parts.append(
+            torch.nn.functional.kl_div(
+                log_q, log_p, log_target=True, reduction="none",
+            ).sum(dim=-1)
+        )
+    expected = torch.cat(expected_parts).mean().item()
+    expected_first = torch.stack([part[-1] for part in expected_parts]).mean().item()
+
+    assert pipeline._measure_sequence_token_kl() == pytest.approx(expected)
+    assert pipeline._last_first_token_kl == pytest.approx(expected_first)
+
+
+@pytest.mark.parametrize("failure", ["missing", "shape", "nonfinite"])
+def test_sequence_token_kl_fails_closed_on_invalid_baseline(failure):
+    current = torch.zeros(2, 3)
+    pipeline = _bare_pipeline()
+    pipeline.handle = SimpleNamespace(
+        model=_LogitModel({2: current}),
+        tokenizer=_Tokenizer(),
+    )
+    pipeline._kl_eval_prompts = ["short"]
+    if failure == "missing":
+        pipeline._baseline_token_logits = []
+    elif failure == "shape":
+        pipeline._baseline_token_logits = [torch.zeros(3, 3)]
+    else:
+        pipeline._baseline_token_logits = [torch.full((2, 3), float("nan"))]
+
+    with pytest.raises(RuntimeError, match="KL baseline|changed shape|non-finite"):
+        pipeline._measure_sequence_token_kl()
+
+
+def test_layer_snapshot_restore_is_bit_exact():
+    layer = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4))
+    snapshot = AbliterationPipeline._snapshot_layer_state(layer)
+    expected = {name: value.clone() for name, value in snapshot.items()}
+
+    with torch.no_grad():
+        for value in layer.state_dict(keep_vars=True).values():
+            value.add_(torch.randn_like(value))
+    AbliterationPipeline._restore_layer_state(layer, snapshot)
+
+    for name, value in layer.state_dict().items():
+        assert torch.equal(value.cpu(), expected[name])
+
+
+def test_optimizer_restores_layers_until_measured_budget_is_met():
+    layers = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)])
+    pristine = {}
+    for index, layer in enumerate(layers):
+        with torch.no_grad():
+            layer.weight.zero_()
+        pristine[index] = AbliterationPipeline._snapshot_layer_state(layer)
+        with torch.no_grad():
+            layer.weight.fill_(1.0)
+
+    pipeline = _bare_pipeline()
+    pipeline.kl_budget = 0.5
+    pipeline._strong_layers = [0, 1]
+    pipeline._measure_sequence_token_kl = lambda: sum(
+        layer.weight.abs().sum().item() for layer in layers
+    )
+
+    pipeline._kl_optimize_corrections(layers, 2, pristine)
+
+    assert all(torch.equal(layer.weight, torch.zeros_like(layer.weight)) for layer in layers)
+    assert pipeline._quality_metrics["kl_divergence"] == 0.0
+    assert pipeline._quality_metrics["kl_budget"] == 0.5
+    assert set(pipeline._kl_contributions) == {0, 1}
+    assert pipeline._strong_layers == []
+
+
+def test_optimizer_does_not_mutate_layers_at_budget_boundary():
+    layer = nn.Linear(2, 2, bias=False)
+    pristine = {0: AbliterationPipeline._snapshot_layer_state(layer)}
+    expected = layer.weight.detach().clone()
+    pipeline = _bare_pipeline()
+    pipeline.kl_budget = 0.5
+    pipeline._measure_sequence_token_kl = lambda: 0.5
+
+    pipeline._kl_optimize_corrections(nn.ModuleList([layer]), 1, pristine)
+
+    assert torch.equal(layer.weight, expected)
+
+
+def test_optimizer_fails_when_exact_candidates_cannot_meet_budget():
+    layer = nn.Linear(2, 2, bias=False)
+    pristine = {0: AbliterationPipeline._snapshot_layer_state(layer)}
+    pipeline = _bare_pipeline()
+    pipeline.kl_budget = 0.1
+    pipeline._measure_sequence_token_kl = lambda: 1.0
+
+    with pytest.raises(RuntimeError, match="could not satisfy"):
+        pipeline._kl_optimize_corrections(nn.ModuleList([layer]), 1, pristine)