Spaces:
Running on Zero
Running on Zero
File size: 3,265 Bytes
9d24374 b3d11b8 393bb89 9d24374 b3d11b8 9d24374 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
import torch
@dataclass(frozen=True)
class InterventionSpec:
mode: str
coefficient: float
def delta_activation(self, original_activation: float) -> float:
mode = self.mode.lower().strip()
original = float(original_activation)
coefficient = float(self.coefficient)
if mode == 'ablate':
return -original
if mode == 'scale':
return (coefficient - 1.0) * original
if mode == 'inject':
return coefficient
raise ValueError("mode must be one of: 'ablate', 'scale', 'inject'.")
def residual_delta(
decoder_direction: torch.Tensor,
original_activation: float,
spec: InterventionSpec,
) -> torch.Tensor:
"""Return the reconstruction-preserving SAE delta applied to the original residual."""
return decoder_direction * spec.delta_activation(original_activation)
def joint_residual_delta(
decoder_directions: torch.Tensor,
original_activations: Sequence[float] | torch.Tensor,
spec: InterventionSpec,
) -> tuple[torch.Tensor, list[float]]:
"""
Sum reconstruction-preserving deltas for a set of SAE features.
``decoder_directions`` must have shape ``[n_features, d_model]``. The same
ablation/scale intervention is applied to every selected feature. ``inject``
is intentionally rejected for feature sets because a shared additive
coefficient has ambiguous semantics across unrelated decoder directions.
"""
mode = spec.mode.lower().strip()
if mode not in {'ablate', 'scale'}:
raise ValueError("Feature-set interventions support only 'ablate' or 'scale'.")
directions = decoder_directions
if directions.ndim != 2:
raise ValueError('decoder_directions must have shape [n_features, d_model].')
if isinstance(original_activations, torch.Tensor):
activations = original_activations.detach().float().reshape(-1).tolist()
else:
activations = [float(x) for x in original_activations]
if len(activations) != directions.shape[0]:
raise ValueError('Number of activations must match decoder directions.')
if not activations:
raise ValueError('Select at least one feature for a feature-set intervention.')
coefficient_deltas = [spec.delta_activation(value) for value in activations]
coeff = torch.tensor(
coefficient_deltas,
device=directions.device,
dtype=directions.dtype,
)
delta = torch.sum(directions * coeff[:, None], dim=0)
return delta, coefficient_deltas
def normalized_random_control(delta: torch.Tensor, seed: int) -> torch.Tensor:
"""Generate a deterministic random residual perturbation with identical L2 norm."""
norm = torch.linalg.vector_norm(delta.float())
if float(norm.item()) == 0.0:
return torch.zeros_like(delta)
generator = torch.Generator(device='cpu').manual_seed(int(seed))
random_vec = torch.randn(delta.shape, generator=generator, dtype=torch.float32)
random_vec = random_vec / torch.linalg.vector_norm(random_vec)
random_vec = random_vec * norm.cpu()
return random_vec.to(device=delta.device, dtype=delta.dtype)
|