"""Steering Vectors for inference-time refusal intervention. The existing OBLITERATUS pipeline only supports permanent weight modification. Steering vectors provide a complementary approach: modifying activations at inference time without changing any weights. This is based on: - Turner et al. (2023): "Activation Addition: Steering Language Models Without Optimization" - Rimsky et al. (2024): "Steering Llama 2 via Contrastive Activation Addition" (CAA) - Li et al. (2024): "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model" The approach: 1. Compute a steering vector from the refusal direction (or any concept direction) 2. At inference time, add or subtract scaled multiples of the vector to the residual stream at specified layers 3. This steers the model toward or away from refusal without modifying any weights Advantages over weight projection: - **Reversible**: Steering can be turned on/off per-request - **Tunable**: The steering strength (alpha) can be adjusted continuously - **Composable**: Multiple steering vectors can be combined - **Non-destructive**: Model weights are never modified Limitations (vs. weight projection): - Requires wrapping the model's forward pass (hooks) - Slight inference-time overhead per token - Effect is per-token, not permanent This module provides: 1. SteeringVector construction from refusal directions or contrastive pairs 2. Hook-based application to any HuggingFace model 3. Multi-layer steering with per-layer alpha scaling 4. Evaluation utilities for measuring steering effectiveness References: - Turner et al. (2023): Activation Addition (arXiv:2308.10248) - Rimsky et al. (2024): Contrastive Activation Addition for Llama 2 - Li et al. (2024): Inference-Time Intervention (arXiv:2306.03341) """ from __future__ import annotations from dataclasses import dataclass, field import torch import torch.nn as nn @dataclass class SteeringVector: """A steering vector that can be applied at inference time.""" direction: torch.Tensor # (hidden_dim,) unit vector source_layer: int | None # layer it was extracted from (None if synthetic) label: str # human-readable label (e.g. "refusal", "truthfulness") default_alpha: float # recommended steering strength metadata: dict = field(default_factory=dict) @dataclass class SteeringConfig: """Configuration for inference-time steering.""" vectors: list[SteeringVector] target_layers: list[int] # which layers to steer at alpha: float = 1.0 # global scaling factor per_layer_alpha: dict[int, float] = field(default_factory=dict) # per-layer overrides position: str = "all" # "all", "last", or "first" — which positions to steer normalize: bool = True # normalize vector to unit norm before scaling @dataclass class SteeringResult: """Result of applying steering vectors.""" config: SteeringConfig hooks_installed: int total_steered_layers: int class SteeringVectorFactory: """Create steering vectors from various sources.""" @staticmethod def from_refusal_direction( refusal_direction: torch.Tensor, source_layer: int | None = None, alpha: float = -1.0, ) -> SteeringVector: """Create a steering vector from a pre-computed refusal direction. By default, alpha=-1.0 steers AWAY from refusal (removes it). Use alpha=+1.0 to steer TOWARD refusal (reinforces it). Args: refusal_direction: (hidden_dim,) refusal direction vector. source_layer: Layer the direction was extracted from. alpha: Steering strength. Negative = away from refusal. Returns: SteeringVector ready for application. """ d = refusal_direction.float().squeeze() d = d / d.norm().clamp(min=1e-10) return SteeringVector( direction=d, source_layer=source_layer, label="refusal", default_alpha=alpha, ) @staticmethod def from_contrastive_pairs( positive_activations: list[torch.Tensor], negative_activations: list[torch.Tensor], label: str = "contrastive", alpha: float = 1.0, ) -> SteeringVector: """Create a steering vector from contrastive activation pairs. The vector is the difference in mean activations: vector = mean(positive) - mean(negative) Args: positive_activations: Activations from "positive" examples (e.g., harmful prompts that trigger refusal). negative_activations: Activations from "negative" examples (e.g., harmless prompts without refusal). label: Human-readable label. alpha: Default steering strength. Returns: SteeringVector from contrastive difference. """ pos_mean = torch.stack([a.float().squeeze() for a in positive_activations]).mean(dim=0) neg_mean = torch.stack([a.float().squeeze() for a in negative_activations]).mean(dim=0) diff = pos_mean - neg_mean d = diff / diff.norm().clamp(min=1e-10) return SteeringVector( direction=d, source_layer=None, label=label, default_alpha=alpha, metadata={"n_positive": len(positive_activations), "n_negative": len(negative_activations), "raw_magnitude": diff.norm().item()}, ) @staticmethod def combine( vectors: list[SteeringVector], weights: list[float] | None = None, label: str = "combined", ) -> SteeringVector: """Combine multiple steering vectors into one. Args: vectors: List of SteeringVector to combine. weights: Per-vector weights. If None, equal weights. label: Label for the combined vector. Returns: Combined SteeringVector. """ if not vectors: raise ValueError("Need at least one vector to combine") if weights is None: weights = [1.0 / len(vectors)] * len(vectors) combined = sum(w * v.direction for w, v in zip(weights, vectors)) combined = combined / combined.norm().clamp(min=1e-10) mean_alpha = sum(v.default_alpha for v in vectors) / len(vectors) return SteeringVector( direction=combined, source_layer=None, label=label, default_alpha=mean_alpha, metadata={"n_combined": len(vectors), "weights": weights}, ) class SteeringHookManager: """Manages inference-time hooks for applying steering vectors. This class installs PyTorch forward hooks on specified layers to add/subtract steering vectors from the residual stream. """ def __init__(self): self._hooks: list = [] self._active = False def install( self, model: nn.Module, config: SteeringConfig, layer_modules: list[nn.Module] | None = None, ) -> SteeringResult: """Install steering hooks on the model. Args: model: The transformer model. config: SteeringConfig specifying vectors, layers, and alphas. layer_modules: If provided, use these as the layer modules. Otherwise, attempts to find them automatically. Returns: SteeringResult with installation details. """ self.remove() # Clean up any existing hooks if layer_modules is None: layer_modules = self._find_layer_modules(model) n_installed = 0 for layer_idx in config.target_layers: if layer_idx >= len(layer_modules): continue module = layer_modules[layer_idx] alpha = config.per_layer_alpha.get(layer_idx, config.alpha) hook = self._make_hook(config.vectors, alpha, config.position, config.normalize) handle = module.register_forward_hook(hook) self._hooks.append(handle) n_installed += 1 self._active = True return SteeringResult( config=config, hooks_installed=n_installed, total_steered_layers=n_installed, ) def remove(self): """Remove all installed hooks.""" for handle in self._hooks: handle.remove() self._hooks.clear() self._active = False @property def is_active(self) -> bool: return self._active def _make_hook( self, vectors: list[SteeringVector], alpha: float, position: str, normalize: bool, ): """Create a forward hook that applies steering vectors.""" def hook(module, input, output): # output is typically (hidden_states, ...) or just hidden_states if isinstance(output, tuple): hidden = output[0] rest = output[1:] else: hidden = output rest = None for vec in vectors: d = vec.direction.to(hidden.device, hidden.dtype) if normalize: d = d / d.norm().clamp(min=1e-10) scale = alpha * vec.default_alpha steering = scale * d if hidden.ndim == 3: # (batch, seq_len, hidden_dim) — typical transformer output if position == "last": hidden = hidden.clone() hidden[:, -1, :] = hidden[:, -1, :] + steering elif position == "first": hidden = hidden.clone() hidden[:, 0, :] = hidden[:, 0, :] + steering else: # "all" hidden = hidden + steering.unsqueeze(0).unsqueeze(0) elif hidden.ndim == 2: # (batch, hidden_dim) — e.g., linear layer output hidden = hidden + steering.unsqueeze(0) else: # Unsupported shape — add along last dim as best effort hidden = hidden + steering if rest is not None: return (hidden,) + rest return hidden return hook @staticmethod def _find_layer_modules(model: nn.Module) -> list[nn.Module]: """Auto-detect transformer layer modules.""" # Common attribute paths for transformer layers for attr_path in [ "model.layers", "transformer.h", "gpt_neox.layers", "model.decoder.layers", "encoder.layer", ]: obj = model try: for part in attr_path.split("."): obj = getattr(obj, part) return list(obj) except AttributeError: continue return [] def compute_steering_effectiveness( clean_projection: float, steered_projection: float, direction: str = "remove", ) -> float: """Compute how effective steering was. Args: clean_projection: Refusal projection without steering. steered_projection: Refusal projection with steering. direction: "remove" (want to reduce) or "add" (want to increase). Returns: Effectiveness score (0-1). 1.0 = perfectly effective. """ if direction == "remove": if abs(clean_projection) < 1e-10: return 1.0 # Already no refusal return max(0.0, 1.0 - abs(steered_projection) / abs(clean_projection)) else: if abs(steered_projection) < 1e-10: return 0.0 return min(1.0, abs(steered_projection) / max(abs(clean_projection), 1e-10)) def format_steering_report(result: SteeringResult) -> str: """Format steering application report.""" lines = [] lines.append("Steering Vector Application") lines.append("=" * 35) lines.append("") lines.append(f"Hooks installed: {result.hooks_installed}") lines.append(f"Layers steered: {result.total_steered_layers}") lines.append(f"Global alpha: {result.config.alpha}") lines.append(f"Position mode: {result.config.position}") lines.append(f"Vectors applied: {len(result.config.vectors)}") for v in result.config.vectors: lines.append(f" - {v.label} (alpha={v.default_alpha:+.2f}, dim={v.direction.shape[0]})") return "\n".join(lines)