"""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, )