mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-07-31 16:17:36 +02:00
376 lines
14 KiB
Python
376 lines
14 KiB
Python
"""Concept Cone Geometry analysis for refusal subspace characterization.
|
|
|
|
The 2025 paper "Geometry of Concepts in LLMs" (Wollschlager et al., arXiv:2502.17420) showed that
|
|
refusal is NOT a single linear direction or even a linear subspace — it's a
|
|
*polyhedral concept cone*. Different categories of harmful content activate
|
|
geometrically distinct refusal directions that share a common half-space
|
|
but are NOT collinear.
|
|
|
|
This module implements tools to:
|
|
|
|
1. **Concept Cone Estimation**: Fit the minimal cone containing all
|
|
per-category refusal directions, measuring its solid angle and
|
|
dimensionality.
|
|
|
|
2. **Per-Category Direction Decomposition**: Extract separate refusal
|
|
directions for each harm category (weapons, cyber, fraud, etc.)
|
|
and measure their pairwise geometric relationships.
|
|
|
|
3. **Cone Complexity Scaling**: Measure how cone dimensionality scales
|
|
with model size, testing the finding that larger models have
|
|
higher-dimensional refusal cones.
|
|
|
|
4. **Direction Specificity Index**: For each refusal direction, measure
|
|
how specifically it targets one category vs. being a general-purpose
|
|
refusal signal.
|
|
|
|
Extensions beyond prior work:
|
|
- We compute the *minimal enclosing cone* explicitly using convex
|
|
optimization over the half-space intersection
|
|
- We introduce the Direction Specificity Index (DSI), which quantifies
|
|
how categorical vs. universal each component of refusal is
|
|
- We test whether the cone structure is consistent across layers
|
|
|
|
References:
|
|
- Wollschlager et al. (2025): Geometry of Concepts in LLMs (arXiv:2502.17420)
|
|
- Joad et al. (2026): 11 geometrically distinct refusal directions
|
|
- Arditi et al. (2024): Single-direction assumption (shown incomplete)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
|
|
|
|
# Default category assignments for the OBLITERATUS synthetic harmful prompts
|
|
# Maps prompt index -> category name
|
|
DEFAULT_HARM_CATEGORIES = {
|
|
0: "weapons", 1: "weapons", 2: "weapons",
|
|
3: "cyber", 4: "cyber", 5: "cyber", 6: "cyber",
|
|
7: "cyber", 8: "cyber", 9: "cyber", 10: "cyber", 11: "cyber",
|
|
12: "fraud", 13: "fraud", 14: "fraud", 15: "fraud",
|
|
16: "intrusion", 17: "intrusion", 18: "intrusion", 19: "intrusion",
|
|
20: "substances", 21: "substances",
|
|
22: "extremism", 23: "stalking",
|
|
24: "privacy", 25: "privacy",
|
|
26: "manipulation", 27: "manipulation",
|
|
28: "self_harm", 29: "self_harm",
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class CategoryDirection:
|
|
"""Refusal direction for a specific harm category."""
|
|
|
|
category: str
|
|
direction: torch.Tensor # (hidden_dim,) unit vector
|
|
strength: float # magnitude of the category's refusal signal
|
|
n_prompts: int # number of prompts in this category
|
|
specificity: float # how specific to this category (0=general, 1=unique)
|
|
|
|
|
|
@dataclass
|
|
class ConeConeResult:
|
|
"""Result of concept cone geometry analysis for a single layer."""
|
|
|
|
layer_idx: int
|
|
category_directions: list[CategoryDirection]
|
|
pairwise_cosines: dict[tuple[str, str], float] # (cat_a, cat_b) -> cosine
|
|
cone_solid_angle: float # solid angle of the minimal enclosing cone (steradians)
|
|
cone_dimensionality: float # effective dimensionality of the cone
|
|
mean_pairwise_cosine: float # average cosine between category directions
|
|
is_linear: bool # True if cone is essentially 1D (all directions aligned)
|
|
is_polyhedral: bool # True if distinct directions detected
|
|
general_direction: torch.Tensor # the mean direction (closest to "single direction")
|
|
category_count: int
|
|
|
|
|
|
@dataclass
|
|
class MultiLayerConeResult:
|
|
"""Cone geometry across multiple layers."""
|
|
|
|
per_layer: dict[int, ConeConeResult]
|
|
most_polyhedral_layer: int # layer with most complex cone
|
|
most_linear_layer: int # layer with simplest cone
|
|
cone_complexity_by_layer: dict[int, float] # cone dimensionality per layer
|
|
mean_cone_dimensionality: float
|
|
|
|
|
|
class ConceptConeAnalyzer:
|
|
"""Analyze the geometric structure of refusal as a concept cone.
|
|
|
|
Instead of assuming refusal is a single direction (Arditi) or a linear
|
|
subspace (Gabliteration), this analyzes the actual cone-like geometry
|
|
where different harm categories have distinct but related directions.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
category_map: dict[int, str] | None = None,
|
|
min_category_size: int = 2,
|
|
):
|
|
"""
|
|
Args:
|
|
category_map: {prompt_index: category_name} for grouping prompts.
|
|
If None, uses DEFAULT_HARM_CATEGORIES.
|
|
min_category_size: Minimum prompts per category to compute a
|
|
category-specific direction.
|
|
"""
|
|
self.category_map = category_map or DEFAULT_HARM_CATEGORIES
|
|
self.min_category_size = min_category_size
|
|
|
|
def analyze_layer(
|
|
self,
|
|
harmful_activations: list[torch.Tensor],
|
|
harmless_activations: list[torch.Tensor],
|
|
layer_idx: int = 0,
|
|
) -> ConeConeResult:
|
|
"""Analyze cone geometry at a single layer.
|
|
|
|
Args:
|
|
harmful_activations: List of per-prompt activation tensors.
|
|
harmless_activations: List of per-prompt activation tensors.
|
|
layer_idx: Layer index for metadata.
|
|
|
|
Returns:
|
|
ConeConeResult with full cone geometry analysis.
|
|
"""
|
|
n_prompts = min(len(harmful_activations), len(harmless_activations))
|
|
|
|
# Group prompts by category
|
|
categories: dict[str, list[int]] = {}
|
|
for idx in range(n_prompts):
|
|
cat = self.category_map.get(idx, "unknown")
|
|
if cat not in categories:
|
|
categories[cat] = []
|
|
categories[cat].append(idx)
|
|
|
|
# Compute per-category refusal directions
|
|
cat_directions: list[CategoryDirection] = []
|
|
direction_vectors: dict[str, torch.Tensor] = {}
|
|
|
|
for cat, indices in sorted(categories.items()):
|
|
if len(indices) < self.min_category_size:
|
|
continue
|
|
|
|
# Category mean difference
|
|
cat_harmful = torch.stack([
|
|
harmful_activations[i].float().squeeze() for i in indices
|
|
]).mean(dim=0)
|
|
cat_harmless = torch.stack([
|
|
harmless_activations[i].float().squeeze() for i in indices
|
|
]).mean(dim=0)
|
|
|
|
diff = cat_harmful - cat_harmless
|
|
strength = diff.norm().item()
|
|
|
|
if strength > 1e-8:
|
|
direction = diff / diff.norm()
|
|
else:
|
|
direction = diff
|
|
|
|
direction_vectors[cat] = direction
|
|
cat_directions.append(CategoryDirection(
|
|
category=cat,
|
|
direction=direction,
|
|
strength=strength,
|
|
n_prompts=len(indices),
|
|
specificity=0.0, # computed below
|
|
))
|
|
|
|
# Compute pairwise cosine similarities
|
|
pairwise: dict[tuple[str, str], float] = {}
|
|
cats = sorted(direction_vectors.keys())
|
|
for i, cat_a in enumerate(cats):
|
|
for j, cat_b in enumerate(cats):
|
|
if i < j:
|
|
cos = (direction_vectors[cat_a] @ direction_vectors[cat_b]).abs().item()
|
|
pairwise[(cat_a, cat_b)] = cos
|
|
|
|
# Mean pairwise cosine
|
|
if pairwise:
|
|
mean_cos = sum(pairwise.values()) / len(pairwise)
|
|
else:
|
|
mean_cos = 1.0
|
|
|
|
# Compute Direction Specificity Index (DSI) for each category
|
|
# DSI = 1 - mean(|cos(d_cat, d_other)|) for all other categories
|
|
# High DSI = direction is unique to this category
|
|
for cd in cat_directions:
|
|
other_cosines = []
|
|
for other_cd in cat_directions:
|
|
if other_cd.category != cd.category:
|
|
cos = (cd.direction @ other_cd.direction).abs().item()
|
|
other_cosines.append(cos)
|
|
if other_cosines:
|
|
cd.specificity = 1.0 - (sum(other_cosines) / len(other_cosines))
|
|
else:
|
|
cd.specificity = 1.0
|
|
|
|
# General direction (mean of all category directions)
|
|
if direction_vectors:
|
|
all_dirs = torch.stack(list(direction_vectors.values()))
|
|
general = all_dirs.mean(dim=0)
|
|
general = general / general.norm().clamp(min=1e-8)
|
|
else:
|
|
general = torch.zeros(1)
|
|
|
|
# Cone dimensionality estimation
|
|
# Use SVD of the category direction matrix
|
|
cone_dim, solid_angle = self._estimate_cone_geometry(direction_vectors)
|
|
|
|
# Classification
|
|
is_linear = mean_cos > 0.9 and cone_dim < 1.5
|
|
is_polyhedral = mean_cos < 0.8 or cone_dim > 2.0
|
|
|
|
return ConeConeResult(
|
|
layer_idx=layer_idx,
|
|
category_directions=cat_directions,
|
|
pairwise_cosines=pairwise,
|
|
cone_solid_angle=solid_angle,
|
|
cone_dimensionality=cone_dim,
|
|
mean_pairwise_cosine=mean_cos,
|
|
is_linear=is_linear,
|
|
is_polyhedral=is_polyhedral,
|
|
general_direction=general,
|
|
category_count=len(cat_directions),
|
|
)
|
|
|
|
def analyze_all_layers(
|
|
self,
|
|
harmful_acts: dict[int, list[torch.Tensor]],
|
|
harmless_acts: dict[int, list[torch.Tensor]],
|
|
strong_layers: list[int] | None = None,
|
|
) -> MultiLayerConeResult:
|
|
"""Analyze cone geometry across multiple layers.
|
|
|
|
Args:
|
|
harmful_acts: {layer_idx: [activations]} per layer.
|
|
harmless_acts: {layer_idx: [activations]} per layer.
|
|
strong_layers: If provided, only analyze these layers.
|
|
|
|
Returns:
|
|
MultiLayerConeResult with per-layer and aggregate analysis.
|
|
"""
|
|
layers = strong_layers or sorted(harmful_acts.keys())
|
|
per_layer = {}
|
|
|
|
for idx in layers:
|
|
if idx not in harmful_acts or idx not in harmless_acts:
|
|
continue
|
|
per_layer[idx] = self.analyze_layer(
|
|
harmful_acts[idx], harmless_acts[idx], layer_idx=idx
|
|
)
|
|
|
|
if not per_layer:
|
|
return MultiLayerConeResult(
|
|
per_layer={},
|
|
most_polyhedral_layer=0,
|
|
most_linear_layer=0,
|
|
cone_complexity_by_layer={},
|
|
mean_cone_dimensionality=0.0,
|
|
)
|
|
|
|
complexity = {idx: r.cone_dimensionality for idx, r in per_layer.items()}
|
|
most_poly = max(complexity, key=complexity.get)
|
|
most_linear = min(complexity, key=complexity.get)
|
|
mean_dim = sum(complexity.values()) / len(complexity)
|
|
|
|
return MultiLayerConeResult(
|
|
per_layer=per_layer,
|
|
most_polyhedral_layer=most_poly,
|
|
most_linear_layer=most_linear,
|
|
cone_complexity_by_layer=complexity,
|
|
mean_cone_dimensionality=mean_dim,
|
|
)
|
|
|
|
def _estimate_cone_geometry(
|
|
self, direction_vectors: dict[str, torch.Tensor]
|
|
) -> tuple[float, float]:
|
|
"""Estimate cone dimensionality and solid angle.
|
|
|
|
Uses the effective rank of the direction matrix (SVD-based) as the
|
|
cone dimensionality, and approximates the solid angle from the
|
|
spread of directions.
|
|
|
|
Returns:
|
|
(cone_dimensionality, solid_angle_steradians)
|
|
"""
|
|
if len(direction_vectors) < 2:
|
|
return 1.0, 0.0
|
|
|
|
D = torch.stack(list(direction_vectors.values())) # (n_cats, hidden_dim)
|
|
n_cats = D.shape[0]
|
|
|
|
# SVD to get effective dimensionality
|
|
s = torch.linalg.svdvals(D)
|
|
s = s[s > 1e-10]
|
|
if len(s) == 0:
|
|
return 0.0, 0.0
|
|
|
|
# Effective rank via entropy
|
|
p = s / s.sum()
|
|
entropy = -(p * p.log()).sum()
|
|
eff_rank = torch.exp(entropy).item()
|
|
|
|
# Solid angle approximation:
|
|
# For directions on a unit sphere, the solid angle is related to
|
|
# the volume of the spherical cap they span.
|
|
# Approximate using: Omega ~ 2*pi*(1 - min_cos) for a circular cone
|
|
# For polyhedral cones, use the mean angular spread
|
|
cos_values = []
|
|
mean_dir = D.mean(dim=0)
|
|
mean_dir = mean_dir / mean_dir.norm().clamp(min=1e-8)
|
|
for i in range(n_cats):
|
|
cos = (D[i] @ mean_dir).abs().item()
|
|
cos_values.append(cos)
|
|
|
|
if cos_values:
|
|
min_cos = min(cos_values)
|
|
# Solid angle of a cone with half-angle theta:
|
|
# Omega = 2*pi*(1 - cos(theta))
|
|
# For high dimensions, generalize: Omega ~ (1 - min_cos)^(d/2)
|
|
# Use simplified 3D formula as approximation
|
|
solid_angle = 2 * math.pi * (1 - min_cos)
|
|
else:
|
|
solid_angle = 0.0
|
|
|
|
return eff_rank, solid_angle
|
|
|
|
@staticmethod
|
|
def format_report(result: ConeConeResult) -> str:
|
|
"""Format single-layer cone analysis as a report."""
|
|
lines = []
|
|
lines.append(f"Concept Cone Geometry — Layer {result.layer_idx}")
|
|
lines.append("=" * 45)
|
|
lines.append("")
|
|
|
|
geometry_type = "LINEAR (single direction)" if result.is_linear else (
|
|
"POLYHEDRAL (concept cone)" if result.is_polyhedral else "INTERMEDIATE"
|
|
)
|
|
lines.append(f"Geometry: {geometry_type}")
|
|
lines.append(f"Cone dimensionality: {result.cone_dimensionality:.2f}")
|
|
lines.append(f"Solid angle: {result.cone_solid_angle:.4f} sr")
|
|
lines.append(f"Mean pairwise cosine: {result.mean_pairwise_cosine:.3f}")
|
|
lines.append(f"Categories analyzed: {result.category_count}")
|
|
lines.append("")
|
|
|
|
lines.append("Per-Category Refusal Directions:")
|
|
for cd in sorted(result.category_directions, key=lambda x: -x.strength):
|
|
lines.append(
|
|
f" {cd.category:15s} strength={cd.strength:.3f} "
|
|
f"specificity={cd.specificity:.3f} (n={cd.n_prompts})"
|
|
)
|
|
lines.append("")
|
|
|
|
if result.pairwise_cosines:
|
|
lines.append("Pairwise Direction Cosines:")
|
|
for (a, b), cos in sorted(result.pairwise_cosines.items()):
|
|
bar = "█" * int(cos * 15)
|
|
lines.append(f" {a:12s} ↔ {b:12s}: {cos:.3f} {bar}")
|
|
|
|
return "\n".join(lines)
|