mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
243 lines
9.1 KiB
Python
243 lines
9.1 KiB
Python
"""Pure validation contracts shared by numerical analysis routines."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProjectionResult:
|
|
"""Pure rank-1 projection result and metadata."""
|
|
|
|
weight: torch.Tensor
|
|
projected: bool
|
|
coefficient_norm_sq: float
|
|
layout: str | None
|
|
|
|
|
|
def _stable_float_dtype(*tensors: torch.Tensor) -> torch.dtype:
|
|
"""Select a supported compute dtype without downcasting float64 inputs."""
|
|
if any(tensor.dtype == torch.float64 for tensor in tensors):
|
|
return torch.float64
|
|
if any(tensor.dtype in {torch.float16, torch.bfloat16} for tensor in tensors):
|
|
return torch.float32
|
|
if all(tensor.is_floating_point() for tensor in tensors):
|
|
return tensors[0].dtype
|
|
return torch.float32
|
|
|
|
|
|
def validate_whitened_parameters(
|
|
regularization_eps: object,
|
|
min_variance_ratio: object,
|
|
) -> tuple[float, float]:
|
|
"""Validate and normalize whitened-SVD tuning parameters."""
|
|
if (
|
|
isinstance(regularization_eps, bool)
|
|
or not isinstance(regularization_eps, (int, float))
|
|
or not math.isfinite(regularization_eps)
|
|
or regularization_eps <= 0
|
|
):
|
|
raise ValueError("regularization_eps must be a finite positive number")
|
|
if (
|
|
isinstance(min_variance_ratio, bool)
|
|
or not isinstance(min_variance_ratio, (int, float))
|
|
or not math.isfinite(min_variance_ratio)
|
|
or not 0 <= min_variance_ratio < 1
|
|
):
|
|
raise ValueError("min_variance_ratio must be in the interval [0, 1)")
|
|
return float(regularization_eps), float(min_variance_ratio)
|
|
|
|
|
|
def validate_whitened_request(
|
|
harmful_count: int,
|
|
harmless_count: int,
|
|
n_directions: object,
|
|
) -> int:
|
|
"""Validate paired sample counts and requested direction count."""
|
|
if harmful_count <= 0 or harmless_count <= 0:
|
|
raise ValueError("harmful and harmless activations must both be non-empty")
|
|
if harmful_count != harmless_count:
|
|
raise ValueError(
|
|
"harmful and harmless activations must have equal sample counts, got "
|
|
f"{harmful_count} and {harmless_count}",
|
|
)
|
|
if isinstance(n_directions, bool) or not isinstance(n_directions, int):
|
|
raise ValueError("n_directions must be a positive integer")
|
|
if n_directions <= 0:
|
|
raise ValueError("n_directions must be a positive integer")
|
|
return n_directions
|
|
|
|
|
|
def orthogonalize_subspace_rows(subspace: torch.Tensor) -> torch.Tensor:
|
|
"""Orthogonalize rows of a subspace matrix with QR while preserving dtype/device."""
|
|
if subspace.shape[0] <= 1 or subspace.numel() == 0:
|
|
return subspace
|
|
if not torch.isfinite(subspace).all():
|
|
return subspace
|
|
compute_dtype = _stable_float_dtype(subspace)
|
|
work = subspace.to(dtype=compute_dtype)
|
|
if work.norm() < 1e-8:
|
|
return torch.zeros_like(subspace)
|
|
|
|
q, _ = torch.linalg.qr(work.T)
|
|
result = q[:, : subspace.shape[0]].T
|
|
if (result[0] @ work[0]) < 0:
|
|
result[0] = -result[0]
|
|
return result.to(dtype=subspace.dtype, device=subspace.device)
|
|
|
|
|
|
def remove_harmless_principal_components(
|
|
subspace: torch.Tensor,
|
|
harmless_stack: torch.Tensor,
|
|
pc_count: int,
|
|
) -> torch.Tensor:
|
|
"""Subtract dominant benign activation PCs from refusal directions."""
|
|
if pc_count <= 0 or harmless_stack.shape[0] < 3 or subspace.numel() == 0:
|
|
return subspace
|
|
|
|
compute_dtype = _stable_float_dtype(subspace, harmless_stack)
|
|
harmless_work = harmless_stack.to(dtype=compute_dtype)
|
|
centered = harmless_work - harmless_work.mean(dim=0, keepdim=True)
|
|
try:
|
|
_, _, vh = torch.linalg.svd(centered, full_matrices=False)
|
|
except Exception:
|
|
return subspace
|
|
|
|
k = min(int(pc_count), vh.shape[0], subspace.shape[1])
|
|
if k <= 0:
|
|
return subspace
|
|
|
|
original = subspace.to(dtype=compute_dtype)
|
|
pcs = vh[:k]
|
|
residual = original - (original @ pcs.T) @ pcs
|
|
row_norms = residual.norm(dim=-1, keepdim=True)
|
|
near_zero = row_norms.squeeze(-1) < 1e-8
|
|
if near_zero.any():
|
|
residual[near_zero] = original[near_zero]
|
|
row_norms = residual.norm(dim=-1, keepdim=True)
|
|
|
|
residual = residual / row_norms.clamp(min=1e-8)
|
|
if residual.shape[0] > 1:
|
|
residual = orthogonalize_subspace_rows(residual)
|
|
return residual.to(dtype=subspace.dtype, device=subspace.device)
|
|
|
|
|
|
def residualize_against_shield_atoms(
|
|
subspace: torch.Tensor,
|
|
atoms: torch.Tensor,
|
|
ridge: float,
|
|
) -> torch.Tensor:
|
|
"""Remove protected concept atoms with ridge-regularized projection."""
|
|
if atoms.numel() == 0 or subspace.numel() == 0:
|
|
return subspace
|
|
|
|
compute_dtype = _stable_float_dtype(subspace, atoms)
|
|
original = subspace.to(dtype=compute_dtype)
|
|
normalized_atoms = atoms.to(dtype=compute_dtype)
|
|
normalized_atoms = normalized_atoms / normalized_atoms.norm(dim=-1, keepdim=True).clamp(min=1e-8)
|
|
gram = normalized_atoms @ normalized_atoms.T
|
|
eye = torch.eye(gram.shape[0], dtype=gram.dtype, device=gram.device)
|
|
try:
|
|
coeff = torch.linalg.solve(gram + float(ridge) * eye, normalized_atoms @ original.T)
|
|
except Exception:
|
|
return subspace
|
|
|
|
residual = original - coeff.T @ normalized_atoms
|
|
row_norms = residual.norm(dim=-1, keepdim=True)
|
|
near_zero = row_norms.squeeze(-1) < 1e-8
|
|
if near_zero.any():
|
|
residual[near_zero] = original[near_zero]
|
|
row_norms = residual.norm(dim=-1, keepdim=True)
|
|
|
|
residual = residual / row_norms.clamp(min=1e-8)
|
|
if residual.shape[0] > 1:
|
|
residual = orthogonalize_subspace_rows(residual)
|
|
return residual.to(dtype=subspace.dtype, device=subspace.device)
|
|
|
|
|
|
def select_projection_coefficients(
|
|
coeff: torch.Tensor,
|
|
projection_row_fraction: float,
|
|
) -> torch.Tensor:
|
|
"""Keep only the strongest projection coefficients when selective projection is requested."""
|
|
if not 0.0 < projection_row_fraction <= 1.0:
|
|
raise ValueError("projection_row_fraction must be in (0.0, 1.0]")
|
|
if projection_row_fraction >= 1.0:
|
|
return coeff
|
|
|
|
flat = coeff.detach().abs().reshape(-1).float().cpu()
|
|
n_coeffs = flat.numel()
|
|
if n_coeffs == 0:
|
|
return coeff
|
|
|
|
keep = max(1, min(n_coeffs, math.ceil(n_coeffs * projection_row_fraction)))
|
|
if keep >= n_coeffs:
|
|
return coeff
|
|
|
|
idx = torch.topk(flat, keep, sorted=False).indices
|
|
mask = torch.zeros(n_coeffs, dtype=torch.bool)
|
|
mask[idx] = True
|
|
mask = mask.reshape(coeff.shape).to(device=coeff.device)
|
|
return coeff * mask.to(dtype=coeff.dtype)
|
|
|
|
|
|
def project_weight_against_direction(
|
|
weight: torch.Tensor,
|
|
direction: torch.Tensor,
|
|
*,
|
|
norm_preserve: bool = False,
|
|
regularization: float = 0.0,
|
|
projection_row_fraction: float = 1.0,
|
|
max_norm_ratio: float = 1.10,
|
|
) -> ProjectionResult:
|
|
"""Return a pure rank-1 projection update for standard or transposed weight layouts."""
|
|
compute_dtype = _stable_float_dtype(weight, direction)
|
|
work = weight.to(dtype=compute_dtype)
|
|
d = direction.to(device=weight.device, dtype=compute_dtype).reshape(-1, 1)
|
|
if not torch.isfinite(work).all() or not torch.isfinite(d).all():
|
|
return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=None)
|
|
d_norm = d.norm()
|
|
if d_norm < 1e-8:
|
|
return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=None)
|
|
d = d / d_norm
|
|
|
|
scale = 1.0 - regularization
|
|
original_norm_sq = work.pow(2).sum().item() if norm_preserve else 0.0
|
|
|
|
if work.shape[-1] == d.shape[0]:
|
|
layout = "standard"
|
|
coeff = work @ d
|
|
if not torch.isfinite(coeff).all():
|
|
return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=layout)
|
|
coeff_to_remove = select_projection_coefficients(coeff, projection_row_fraction)
|
|
coeff_norm_sq = coeff_to_remove.pow(2).sum().item() if norm_preserve else 0.0
|
|
projected = work - d.T * (scale * coeff_to_remove)
|
|
elif work.shape[0] == d.shape[0]:
|
|
layout = "transposed"
|
|
coeff = d.T @ work
|
|
if not torch.isfinite(coeff).all():
|
|
return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=layout)
|
|
coeff_to_remove = select_projection_coefficients(coeff, projection_row_fraction)
|
|
coeff_norm_sq = coeff_to_remove.pow(2).sum().item() if norm_preserve else 0.0
|
|
projected = work - (scale * d) * coeff_to_remove
|
|
else:
|
|
return ProjectionResult(weight=weight.clone(), projected=False, coefficient_norm_sq=0.0, layout=None)
|
|
|
|
if norm_preserve and original_norm_sq > 0:
|
|
new_norm_sq = max(0.0, original_norm_sq - scale * (2 - scale) * coeff_norm_sq)
|
|
if new_norm_sq > 0:
|
|
ratio = math.sqrt(original_norm_sq / new_norm_sq)
|
|
if ratio > max_norm_ratio:
|
|
ratio = max_norm_ratio
|
|
projected = projected * ratio
|
|
|
|
return ProjectionResult(
|
|
weight=projected.to(dtype=weight.dtype, device=weight.device),
|
|
projected=True,
|
|
coefficient_norm_sq=coeff_norm_sq,
|
|
layout=layout,
|
|
)
|