From 1f428f60cf62e988e4863a7efa05b410bc0a1228 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Fri, 28 Aug 2026 23:06:32 -0400 Subject: [PATCH] fix: make KL optimization measured and reversible (#182) --- README.md | 2 +- app.py | 6 +- ci/test-risk-map.json | 1 + obliteratus/abliterate.py | 541 +++++++++++------------- tests/test_app_model_load_settings.py | 37 ++ tests/test_kl_optimization_contracts.py | 159 +++++++ 6 files changed, 456 insertions(+), 290 deletions(-) create mode 100644 tests/test_kl_optimization_contracts.py diff --git a/README.md b/README.md index 1e2da22..a921e33 100644 --- a/README.md +++ b/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 | diff --git a/app.py b/app.py index 46cfa9e..c9bae2c 100644 --- a/app.py +++ b/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: diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index c3c485f..6f13763 100644 --- a/ci/test-risk-map.json +++ b/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", diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 99e98d1..1ba9736 100644 --- a/obliteratus/abliterate.py +++ b/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, diff --git a/tests/test_app_model_load_settings.py b/tests/test_app_model_load_settings.py index 7510477..cae9203 100644 --- a/tests/test_app_model_load_settings.py +++ b/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 diff --git a/tests/test_kl_optimization_contracts.py b/tests/test_kl_optimization_contracts.py new file mode 100644 index 0000000..28de087 --- /dev/null +++ b/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)