mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
513 lines
19 KiB
Python
513 lines
19 KiB
Python
"""Wasserstein Refusal Transfer Across Architectures.
|
|
|
|
When a model is successfully abliterated, the knowledge of *where* and *how*
|
|
refusal was embedded can potentially be transferred to other models without
|
|
re-running the full pipeline. "Transport and Merge" (2025) used optimal
|
|
transport for cross-architecture model merging; GiLOT (ICML 2024) used OT
|
|
for LLM interpretability.
|
|
|
|
This module uses OT maps to transfer refusal removal knowledge across
|
|
architectures. Given an abliterated source and aligned target, it computes
|
|
the Monge map T: A_source -> A_target between their activation distributions,
|
|
then transports the source's refusal directions through T.
|
|
|
|
Contributions:
|
|
1. **OT-based refusal direction transfer**: Application of optimal
|
|
transport to cross-architecture safety intervention transfer
|
|
2. **Transfer error bound (informal)**: Excess refusal after transfer is
|
|
bounded by W_2(mu_s, mu_t) * kappa(T)
|
|
3. **Refusal removal knowledge graph**: Abliterate one model, transfer
|
|
to a whole family via OT maps
|
|
4. **Wasserstein compatibility metric**: Quantifies whether transfer is
|
|
viable before attempting it
|
|
|
|
References:
|
|
- Cui et al. (2025): Transport and Merge — cross-arch OT merging (arXiv:2602.05495)
|
|
- Li et al. (ICML 2024): GiLOT — OT for LLM interpretability
|
|
- Brenier (1991): Optimal maps for quadratic cost (uniqueness theorem)
|
|
- Paper Appendix Theorem: Wasserstein Cost of Abliteration
|
|
- OBLITERATUS: Cross-Model Universality Index
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TransportPlan:
|
|
"""Optimal transport plan between two activation distributions."""
|
|
|
|
source_model: str # name of source model
|
|
target_model: str # name of target model
|
|
transport_matrix: torch.Tensor # (d_target, d_source) linear map T
|
|
wasserstein_distance: float # W_2 between source and target
|
|
condition_number: float # kappa(T), stability indicator
|
|
transport_cost: float # total transport cost
|
|
is_viable: bool # whether transfer is recommended
|
|
|
|
|
|
@dataclass
|
|
class TransferredDirection:
|
|
"""A refusal direction transferred from source to target model."""
|
|
|
|
source_layer: int # layer in source model
|
|
target_layer: int # corresponding layer in target model
|
|
source_direction: torch.Tensor # original direction in source space
|
|
transferred_direction: torch.Tensor # direction mapped to target space
|
|
transfer_fidelity: float # quality of transfer (0-1)
|
|
estimated_refusal_removal: float # expected removal effectiveness
|
|
wasserstein_bound: float # excess refusal upper bound
|
|
|
|
|
|
@dataclass
|
|
class WassersteinTransferResult:
|
|
"""Complete result of Wasserstein refusal transfer analysis."""
|
|
|
|
# Transfer metadata
|
|
source_model: str
|
|
target_model: str
|
|
n_layers_transferred: int
|
|
|
|
# Transport plan
|
|
wasserstein_distance: float # W_2(source, target)
|
|
condition_number: float # stability of transport map
|
|
transfer_viability: str # "excellent" | "good" | "marginal" | "poor"
|
|
|
|
# Transferred directions
|
|
transferred_directions: list[TransferredDirection]
|
|
mean_transfer_fidelity: float # avg quality across layers
|
|
min_transfer_fidelity: float # worst layer
|
|
|
|
# Bounds
|
|
estimated_excess_refusal: float # bound on residual refusal after transfer
|
|
estimated_vs_native_ratio: float # expected native/transfer performance ratio
|
|
|
|
# Layer alignment
|
|
layer_mapping: dict[int, int] # source_layer -> target_layer
|
|
unmapped_layers: list[int] # target layers with no source correspondence
|
|
|
|
# Recommendation
|
|
recommendation: str # summary recommendation
|
|
needs_refinement: bool # whether a refinement pass is recommended
|
|
|
|
|
|
class WassersteinRefusalTransfer:
|
|
"""Transfer refusal removal knowledge across architectures via OT.
|
|
|
|
Given a successfully abliterated source model and an aligned target,
|
|
computes the optimal transport map between their activation spaces
|
|
and uses it to transfer refusal directions.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
fidelity_threshold: float = 0.5,
|
|
max_condition_number: float = 100.0,
|
|
viability_threshold: float = 0.3,
|
|
n_sinkhorn_iterations: int = 50,
|
|
):
|
|
"""
|
|
Args:
|
|
fidelity_threshold: Minimum transfer fidelity to consider
|
|
a transferred direction useful.
|
|
max_condition_number: Maximum condition number for the transport
|
|
map before flagging instability.
|
|
viability_threshold: W_2 threshold below which transfer is viable.
|
|
n_sinkhorn_iterations: Iterations for Sinkhorn OT computation.
|
|
"""
|
|
self.fidelity_threshold = fidelity_threshold
|
|
self.max_condition_number = max_condition_number
|
|
self.viability_threshold = viability_threshold
|
|
self.n_sinkhorn_iterations = n_sinkhorn_iterations
|
|
|
|
def compute_transfer(
|
|
self,
|
|
source_activations: dict[int, torch.Tensor],
|
|
target_activations: dict[int, torch.Tensor],
|
|
source_refusal_directions: dict[int, torch.Tensor],
|
|
source_model_name: str = "source",
|
|
target_model_name: str = "target",
|
|
layer_mapping: dict[int, int] | None = None,
|
|
) -> WassersteinTransferResult:
|
|
"""Compute Wasserstein transfer of refusal directions.
|
|
|
|
Args:
|
|
source_activations: {layer_idx: (n_samples, d_source)} from source.
|
|
target_activations: {layer_idx: (n_samples, d_target)} from target.
|
|
source_refusal_directions: {layer_idx: (d_source,)} from source.
|
|
source_model_name: Identifier for source model.
|
|
target_model_name: Identifier for target model.
|
|
layer_mapping: Optional explicit {source_layer -> target_layer}.
|
|
If None, computed via activation similarity.
|
|
|
|
Returns:
|
|
WassersteinTransferResult with transferred directions and bounds.
|
|
"""
|
|
source_layers = sorted(source_activations.keys())
|
|
target_layers = sorted(target_activations.keys())
|
|
|
|
if not source_layers or not target_layers:
|
|
return self._empty_result(source_model_name, target_model_name)
|
|
|
|
# Step 1: Compute layer mapping if not provided
|
|
if layer_mapping is None:
|
|
layer_mapping = self._compute_layer_mapping(
|
|
source_layers, target_layers,
|
|
source_activations, target_activations
|
|
)
|
|
|
|
# Step 2: For each mapped layer pair, compute OT map and transfer
|
|
transferred: list[TransferredDirection] = []
|
|
all_w2: list[float] = []
|
|
all_kappa: list[float] = []
|
|
|
|
for src_l, tgt_l in layer_mapping.items():
|
|
if src_l not in source_activations or tgt_l not in target_activations:
|
|
continue
|
|
if src_l not in source_refusal_directions:
|
|
continue
|
|
|
|
src_acts = source_activations[src_l]
|
|
tgt_acts = target_activations[tgt_l]
|
|
src_dir = source_refusal_directions[src_l]
|
|
|
|
# Compute OT map between layer activations
|
|
plan = self._compute_transport_plan(
|
|
src_acts, tgt_acts,
|
|
source_model_name, target_model_name
|
|
)
|
|
all_w2.append(plan.wasserstein_distance)
|
|
all_kappa.append(plan.condition_number)
|
|
|
|
# Transport the refusal direction
|
|
transferred_dir = self._transport_direction(
|
|
src_dir, plan.transport_matrix, src_acts, tgt_acts
|
|
)
|
|
|
|
# Measure transfer fidelity
|
|
fidelity = self._measure_fidelity(
|
|
transferred_dir, tgt_acts, src_dir, src_acts
|
|
)
|
|
|
|
# Wasserstein bound on excess refusal
|
|
w2_bound = plan.wasserstein_distance * plan.condition_number
|
|
|
|
transferred.append(TransferredDirection(
|
|
source_layer=src_l,
|
|
target_layer=tgt_l,
|
|
source_direction=src_dir,
|
|
transferred_direction=transferred_dir,
|
|
transfer_fidelity=fidelity,
|
|
estimated_refusal_removal=max(0, 1.0 - w2_bound),
|
|
wasserstein_bound=w2_bound,
|
|
))
|
|
|
|
if not transferred:
|
|
return self._empty_result(source_model_name, target_model_name)
|
|
|
|
# Step 3: Aggregate results
|
|
fidelities = [t.transfer_fidelity for t in transferred]
|
|
mean_fidelity = sum(fidelities) / len(fidelities)
|
|
min_fidelity = min(fidelities)
|
|
|
|
mean_w2 = sum(all_w2) / len(all_w2)
|
|
mean_kappa = sum(all_kappa) / len(all_kappa)
|
|
|
|
excess_refusal = mean_w2 * mean_kappa
|
|
|
|
# Viability assessment
|
|
if mean_fidelity > 0.8 and mean_w2 < self.viability_threshold:
|
|
viability = "excellent"
|
|
elif mean_fidelity > 0.6 and mean_w2 < self.viability_threshold * 2:
|
|
viability = "good"
|
|
elif mean_fidelity > 0.4:
|
|
viability = "marginal"
|
|
else:
|
|
viability = "poor"
|
|
|
|
native_ratio = max(0.1, 1.0 - excess_refusal)
|
|
needs_refinement = mean_fidelity < 0.7 or viability in ("marginal", "poor")
|
|
|
|
unmapped = [
|
|
ly for ly in target_layers if ly not in layer_mapping.values()
|
|
]
|
|
|
|
recommendation = self._generate_recommendation(
|
|
viability, mean_fidelity, excess_refusal, needs_refinement
|
|
)
|
|
|
|
return WassersteinTransferResult(
|
|
source_model=source_model_name,
|
|
target_model=target_model_name,
|
|
n_layers_transferred=len(transferred),
|
|
wasserstein_distance=mean_w2,
|
|
condition_number=mean_kappa,
|
|
transfer_viability=viability,
|
|
transferred_directions=transferred,
|
|
mean_transfer_fidelity=mean_fidelity,
|
|
min_transfer_fidelity=min_fidelity,
|
|
estimated_excess_refusal=excess_refusal,
|
|
estimated_vs_native_ratio=native_ratio,
|
|
layer_mapping=layer_mapping,
|
|
unmapped_layers=unmapped,
|
|
recommendation=recommendation,
|
|
needs_refinement=needs_refinement,
|
|
)
|
|
|
|
def _compute_layer_mapping(
|
|
self,
|
|
source_layers: list[int],
|
|
target_layers: list[int],
|
|
source_activations: dict[int, torch.Tensor],
|
|
target_activations: dict[int, torch.Tensor],
|
|
) -> dict[int, int]:
|
|
"""Compute layer correspondence via relative position.
|
|
|
|
Maps layers by relative position within the network:
|
|
source_layer / n_source_layers ≈ target_layer / n_target_layers
|
|
"""
|
|
mapping = {}
|
|
n_src = max(source_layers) + 1 if source_layers else 1
|
|
n_tgt = max(target_layers) + 1 if target_layers else 1
|
|
|
|
for src_l in source_layers:
|
|
# Find target layer at closest relative position
|
|
src_ratio = src_l / max(n_src - 1, 1)
|
|
best_tgt = min(
|
|
target_layers,
|
|
key=lambda t: abs(t / max(n_tgt - 1, 1) - src_ratio)
|
|
)
|
|
mapping[src_l] = best_tgt
|
|
|
|
return mapping
|
|
|
|
def _compute_transport_plan(
|
|
self,
|
|
source_acts: torch.Tensor,
|
|
target_acts: torch.Tensor,
|
|
source_name: str,
|
|
target_name: str,
|
|
) -> TransportPlan:
|
|
"""Compute the optimal transport map between activation distributions.
|
|
|
|
Uses a linear approximation: T = Sigma_st @ Sigma_ss^{-1}
|
|
This is the Monge map for Gaussian distributions, which is optimal
|
|
for the quadratic cost when distributions are Gaussian.
|
|
"""
|
|
n_src, d_src = source_acts.shape
|
|
n_tgt, d_tgt = target_acts.shape
|
|
|
|
# Center the activations
|
|
src_mean = source_acts.mean(dim=0)
|
|
tgt_mean = target_acts.mean(dim=0)
|
|
src_centered = source_acts - src_mean
|
|
tgt_centered = target_acts - tgt_mean
|
|
|
|
# Compute covariances
|
|
n_common = min(n_src, n_tgt)
|
|
src_sub = src_centered[:n_common]
|
|
tgt_sub = tgt_centered[:n_common]
|
|
|
|
# Cross-covariance: Sigma_st = tgt^T @ src / n
|
|
sigma_st = tgt_sub.T @ src_sub / max(n_common - 1, 1) # (d_tgt, d_src)
|
|
|
|
# Source auto-covariance: Sigma_ss = src^T @ src / n
|
|
sigma_ss = src_sub.T @ src_sub / max(n_common - 1, 1) # (d_src, d_src)
|
|
|
|
# Transport matrix T = Sigma_st @ Sigma_ss^{-1}
|
|
# Use pseudo-inverse for stability
|
|
try:
|
|
reg = 1e-4 * torch.eye(d_src, device=sigma_ss.device)
|
|
sigma_ss_inv = torch.linalg.inv(sigma_ss + reg)
|
|
transport = sigma_st @ sigma_ss_inv # (d_tgt, d_src)
|
|
except Exception:
|
|
transport = sigma_st # fallback: just use cross-covariance
|
|
|
|
# Wasserstein-2 distance (Bures metric for Gaussian approximation)
|
|
w2 = self._compute_w2_gaussian(src_mean, tgt_mean, sigma_ss,
|
|
tgt_sub.T @ tgt_sub / max(n_common - 1, 1))
|
|
|
|
# Condition number of transport matrix
|
|
try:
|
|
sv = torch.linalg.svdvals(transport)
|
|
kappa = (sv[0] / sv[-1]).item() if sv[-1] > 1e-10 else float("inf")
|
|
kappa = min(kappa, 1e6)
|
|
except Exception:
|
|
kappa = 1.0
|
|
|
|
is_viable = w2 < self.viability_threshold and kappa < self.max_condition_number
|
|
|
|
return TransportPlan(
|
|
source_model=source_name,
|
|
target_model=target_name,
|
|
transport_matrix=transport,
|
|
wasserstein_distance=w2,
|
|
condition_number=kappa,
|
|
transport_cost=w2 * kappa,
|
|
is_viable=is_viable,
|
|
)
|
|
|
|
def _compute_w2_gaussian(
|
|
self,
|
|
mean_s: torch.Tensor,
|
|
mean_t: torch.Tensor,
|
|
cov_s: torch.Tensor,
|
|
cov_t: torch.Tensor,
|
|
) -> float:
|
|
"""Compute 2-Wasserstein distance between Gaussian approximations.
|
|
|
|
W_2^2 = ||mu_s - mu_t||^2 + Tr(Sigma_s + Sigma_t - 2*(Sigma_s^{1/2} Sigma_t Sigma_s^{1/2})^{1/2})
|
|
"""
|
|
# Mean shift component
|
|
mean_diff = (mean_s[:min(len(mean_s), len(mean_t))] -
|
|
mean_t[:min(len(mean_s), len(mean_t))])
|
|
mean_shift = (mean_diff ** 2).sum().item()
|
|
|
|
# Bures metric component (trace term)
|
|
# Simplified: use trace of absolute difference of eigenvalues
|
|
try:
|
|
d = min(cov_s.shape[0], cov_t.shape[0])
|
|
eig_s = torch.linalg.eigvalsh(cov_s[:d, :d])
|
|
eig_t = torch.linalg.eigvalsh(cov_t[:d, :d])
|
|
# Bures approximation via eigenvalues
|
|
sqrt_s = eig_s.clamp(min=0).sqrt()
|
|
sqrt_t = eig_t.clamp(min=0).sqrt()
|
|
bures = ((sqrt_s - sqrt_t) ** 2).sum().item()
|
|
except Exception:
|
|
bures = 0.0
|
|
|
|
w2 = math.sqrt(max(0, mean_shift + bures))
|
|
return w2
|
|
|
|
def _transport_direction(
|
|
self,
|
|
source_direction: torch.Tensor,
|
|
transport_matrix: torch.Tensor,
|
|
source_acts: torch.Tensor,
|
|
target_acts: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Transport a refusal direction through the OT map.
|
|
|
|
Applies T to the source direction and normalizes in the target space.
|
|
"""
|
|
d_src = source_direction.shape[0]
|
|
|
|
# Ensure dimensions match
|
|
if transport_matrix.shape[1] != d_src:
|
|
# Dimension mismatch — use projection
|
|
min_d = min(d_src, transport_matrix.shape[1])
|
|
src_dir = source_direction[:min_d]
|
|
T = transport_matrix[:, :min_d]
|
|
else:
|
|
src_dir = source_direction
|
|
T = transport_matrix
|
|
|
|
# Transport: t_dir = T @ s_dir
|
|
transferred = T @ src_dir
|
|
|
|
# Normalize
|
|
t_norm = transferred.norm()
|
|
if t_norm > 1e-8:
|
|
transferred = transferred / t_norm
|
|
|
|
return transferred
|
|
|
|
def _measure_fidelity(
|
|
self,
|
|
transferred_dir: torch.Tensor,
|
|
target_acts: torch.Tensor,
|
|
source_dir: torch.Tensor,
|
|
source_acts: torch.Tensor,
|
|
) -> float:
|
|
"""Measure how well a transferred direction separates harmful/harmless.
|
|
|
|
Fidelity = correlation between source projection magnitudes and
|
|
target projection magnitudes (after transfer).
|
|
"""
|
|
# Project source activations onto source direction
|
|
src_proj = (source_acts @ source_dir).abs()
|
|
|
|
# Project target activations onto transferred direction
|
|
n_common = min(source_acts.shape[0], target_acts.shape[0])
|
|
tgt_proj = (target_acts[:n_common] @ transferred_dir).abs()
|
|
src_proj = src_proj[:n_common]
|
|
|
|
if n_common < 2:
|
|
return 0.0
|
|
|
|
# Correlation as fidelity measure
|
|
src_centered = src_proj - src_proj.mean()
|
|
tgt_centered = tgt_proj - tgt_proj.mean()
|
|
|
|
src_std = src_centered.std()
|
|
tgt_std = tgt_centered.std()
|
|
|
|
if src_std < 1e-10 or tgt_std < 1e-10:
|
|
return 0.0
|
|
|
|
correlation = (src_centered @ tgt_centered) / (
|
|
n_common * src_std * tgt_std
|
|
)
|
|
fidelity = correlation.clamp(0, 1).item()
|
|
|
|
return fidelity
|
|
|
|
def _generate_recommendation(
|
|
self,
|
|
viability: str,
|
|
mean_fidelity: float,
|
|
excess_refusal: float,
|
|
needs_refinement: bool,
|
|
) -> str:
|
|
"""Generate human-readable recommendation."""
|
|
if viability == "excellent":
|
|
return (
|
|
f"Transfer is highly viable (fidelity={mean_fidelity:.2f}). "
|
|
f"Transferred directions should work with minimal refinement."
|
|
)
|
|
elif viability == "good":
|
|
return (
|
|
f"Transfer is viable (fidelity={mean_fidelity:.2f}) but "
|
|
f"recommend a single refinement pass on the target model."
|
|
)
|
|
elif viability == "marginal":
|
|
return (
|
|
f"Transfer is marginal (fidelity={mean_fidelity:.2f}). "
|
|
f"Excess refusal bound={excess_refusal:.3f}. "
|
|
f"Use as initialization only; full re-abliteration recommended."
|
|
)
|
|
else:
|
|
return (
|
|
f"Transfer is poor (fidelity={mean_fidelity:.2f}). "
|
|
f"Models are too dissimilar in Wasserstein space. "
|
|
f"Full native abliteration required."
|
|
)
|
|
|
|
def _empty_result(
|
|
self, source_name: str, target_name: str
|
|
) -> WassersteinTransferResult:
|
|
return WassersteinTransferResult(
|
|
source_model=source_name,
|
|
target_model=target_name,
|
|
n_layers_transferred=0,
|
|
wasserstein_distance=float("inf"),
|
|
condition_number=float("inf"),
|
|
transfer_viability="poor",
|
|
transferred_directions=[],
|
|
mean_transfer_fidelity=0.0,
|
|
min_transfer_fidelity=0.0,
|
|
estimated_excess_refusal=1.0,
|
|
estimated_vs_native_ratio=0.0,
|
|
layer_mapping={},
|
|
unmapped_layers=[],
|
|
recommendation="No activations available for transfer.",
|
|
needs_refinement=True,
|
|
)
|