mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
360 lines
12 KiB
Python
360 lines
12 KiB
Python
"""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)
|