wisent 0.7.379__py3-none-any.whl

wisent/core/weight_modification/multi_direction.py

@@ -0,0 +1,410 @@
+"""
+Multi-direction weight modification: Bake TITAN/PRISM/PULSE directions into weights.
+
+These methods learn multiple steering directions per layer. When baking into weights,
+we lose the input-dependent gating but keep the multi-directional coverage.
+
+Baking formula:
+    W_new = W + Σ(weight_i * direction_i)
+
+Where weights are either:
+1. Learned static weights from the method
+2. Average activation weights across training data
+3. Uniform weights (simple average)
+"""
+
+from __future__ import annotations
+
+import torch
+from typing import Dict, List, Optional, Any, TYPE_CHECKING
+from dataclasses import dataclass
+
+from wisent.core.weight_modification.additive import bake_steering_into_weights
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from torch import Tensor
+    from torch.nn import Module
+
+__all__ = [
+    "MultiDirectionConfig",
+    "MultiDirectionResult", 
+    "train_and_bake_titan",
+    "train_and_bake_prism",
+    "train_and_bake_pulse",
+    "bake_multi_directions",
+    "combine_directions",
+]
+
+_LOG = setup_logger(__name__)
+
+
+@dataclass
+class MultiDirectionConfig:
+    """Configuration for multi-direction weight modification."""
+    
+    # Method selection
+    method: str = "titan"
+    """Which method to use: 'titan', 'prism', or 'pulse'"""
+    
+    # Direction combination
+    combination_strategy: str = "learned"
+    """How to combine directions: 'learned', 'uniform', 'pca_weighted'"""
+    
+    # Global scaling
+    alpha: float = 1.0
+    """Global steering strength multiplier"""
+    
+    # Components to modify
+    components: List[str] = None
+    """Components to modify. Default: ['self_attn.o_proj', 'mlp.down_proj']"""
+    
+    # Method-specific configs (passed through)
+    num_directions: int = 5
+    """Number of directions per layer"""
+    
+    optimization_steps: int = 100
+    """Training steps for direction optimization"""
+    
+    # Baking method
+    bake_method: str = "bias"
+    """How to bake: 'bias' or 'weight'"""
+    
+    def __post_init__(self):
+        if self.components is None:
+            self.components = ["self_attn.o_proj", "mlp.down_proj"]
+
+
+@dataclass
+class MultiDirectionResult:
+    """Result from multi-direction training and baking."""
+    
+    method: str
+    """Method used (titan/prism/pulse)"""
+    
+    num_directions: int
+    """Number of directions per layer"""
+    
+    layers_modified: int
+    """Number of layers modified"""
+    
+    directions_per_layer: Dict[str, int]
+    """Number of directions found per layer"""
+    
+    combination_weights: Dict[str, List[float]]
+    """Weights used to combine directions per layer"""
+    
+    effective_vectors: Dict[int, Tensor]
+    """Final combined vectors per layer (what was baked)"""
+    
+    metadata: Dict[str, Any]
+    """Additional metadata from training"""
+
+
+def combine_directions(
+    directions: Tensor,
+    weights: Optional[Tensor] = None,
+    strategy: str = "learned",
+) -> Tensor:
+    """
+    Combine multiple directions into a single effective direction.
+    
+    Args:
+        directions: [num_directions, hidden_dim] tensor
+        weights: Optional [num_directions] weights
+        strategy: 'learned', 'uniform', or 'pca_weighted'
+        
+    Returns:
+        [hidden_dim] combined direction
+    """
+    num_dirs = directions.shape[0]
+    
+    if strategy == "uniform":
+        # Simple average
+        return directions.mean(dim=0)
+    
+    elif strategy == "learned":
+        # Use provided weights (softmax normalized)
+        if weights is None:
+            weights = torch.ones(num_dirs) / num_dirs
+        else:
+            weights = torch.softmax(weights, dim=0)
+        return (weights.unsqueeze(1) * directions).sum(dim=0)
+    
+    elif strategy == "pca_weighted":
+        # Weight by PCA variance (first direction gets most weight)
+        # Assumes directions are ordered by importance
+        variance_weights = torch.tensor([1.0 / (i + 1) for i in range(num_dirs)])
+        variance_weights = variance_weights / variance_weights.sum()
+        return (variance_weights.unsqueeze(1).to(directions.device) * directions).sum(dim=0)
+    
+    else:
+        raise ValueError(f"Unknown strategy: {strategy}")
+
+
+def bake_multi_directions(
+    model: Module,
+    directions: Dict[str, Tensor],
+    weights: Optional[Dict[str, Tensor]] = None,
+    config: Optional[MultiDirectionConfig] = None,
+) -> MultiDirectionResult:
+    """
+    Bake multi-direction steering into model weights.
+    
+    Args:
+        model: Model to modify (in-place)
+        directions: Per-layer directions {layer_name: [num_dirs, hidden_dim]}
+        weights: Optional per-layer weights {layer_name: [num_dirs]}
+        config: Configuration
+        
+    Returns:
+        MultiDirectionResult with modification details
+    """
+    cfg = config or MultiDirectionConfig()
+    log = bind(_LOG, method=cfg.method)
+    
+    # Combine directions into effective vectors
+    effective_vectors: Dict[int, Tensor] = {}
+    combination_weights: Dict[str, List[float]] = {}
+    directions_per_layer: Dict[str, int] = {}
+    
+    for layer_name, layer_dirs in directions.items():
+        # Get layer index
+        try:
+            layer_idx = int(layer_name.replace("layer_", ""))
+        except ValueError:
+            layer_idx = int(layer_name)
+        
+        # Get weights for this layer
+        layer_weights = weights.get(layer_name) if weights else None
+        
+        # Combine directions
+        effective = combine_directions(
+            layer_dirs,
+            layer_weights,
+            strategy=cfg.combination_strategy,
+        )
+        
+        effective_vectors[layer_idx] = effective
+        directions_per_layer[layer_name] = layer_dirs.shape[0]
+        
+        # Record combination weights used
+        if layer_weights is not None:
+            combination_weights[layer_name] = torch.softmax(layer_weights, dim=0).tolist()
+        else:
+            num_dirs = layer_dirs.shape[0]
+            combination_weights[layer_name] = [1.0 / num_dirs] * num_dirs
+    
+    # Bake into weights
+    stats = bake_steering_into_weights(
+        model=model,
+        steering_vectors=effective_vectors,
+        components=cfg.components,
+        alpha=cfg.alpha,
+        method=cfg.bake_method,
+        verbose=True,
+    )
+    
+    return MultiDirectionResult(
+        method=cfg.method,
+        num_directions=cfg.num_directions,
+        layers_modified=stats.get("layers_modified", len(effective_vectors)),
+        directions_per_layer=directions_per_layer,
+        combination_weights=combination_weights,
+        effective_vectors=effective_vectors,
+        metadata={"bake_stats": stats},
+    )
+
+
+def train_and_bake_titan(
+    model: Module,
+    pair_set: "ContrastivePairSet",
+    config: Optional[MultiDirectionConfig] = None,
+    titan_config: Optional[Dict[str, Any]] = None,
+) -> MultiDirectionResult:
+    """
+    Train TITAN and bake directions into model weights.
+    
+    Args:
+        model: Model to modify
+        pair_set: ContrastivePairSet with activations
+        config: Multi-direction config
+        titan_config: Additional TITAN-specific config
+        
+    Returns:
+        MultiDirectionResult
+    """
+    from wisent.core.steering_methods.methods.titan import TITANMethod, TITANConfig
+    
+    cfg = config or MultiDirectionConfig(method="titan")
+    
+    # Build TITAN config
+    t_cfg = TITANConfig(
+        num_directions=cfg.num_directions,
+        optimization_steps=cfg.optimization_steps,
+    )
+    if titan_config:
+        for k, v in titan_config.items():
+            if hasattr(t_cfg, k):
+                setattr(t_cfg, k, v)
+    
+    # Train TITAN
+    print(f"\nTraining TITAN with {cfg.num_directions} directions...")
+    titan = TITANMethod(config=t_cfg)
+    result = titan.train_titan(pair_set)
+    
+    # Extract directions and weights
+    directions = result.directions
+    weights = result.direction_weights
+    
+    print(f"Trained {len(directions)} layers")
+    for layer, dirs in directions.items():
+        print(f"  {layer}: {dirs.shape[0]} directions")
+    
+    # Bake into weights
+    print(f"\nBaking into weights with strategy: {cfg.combination_strategy}")
+    return bake_multi_directions(model, directions, weights, cfg)
+
+
+def train_and_bake_prism(
+    model: Module,
+    pair_set: "ContrastivePairSet",
+    config: Optional[MultiDirectionConfig] = None,
+    prism_config: Optional[Dict[str, Any]] = None,
+) -> MultiDirectionResult:
+    """
+    Train PRISM and bake directions into model weights.
+    
+    Args:
+        model: Model to modify
+        pair_set: ContrastivePairSet with activations
+        config: Multi-direction config
+        prism_config: Additional PRISM-specific config
+        
+    Returns:
+        MultiDirectionResult
+    """
+    from wisent.core.steering_methods.methods.prism import PRISMMethod, PRISMConfig
+    
+    cfg = config or MultiDirectionConfig(method="prism")
+    
+    # Build PRISM config
+    p_cfg = PRISMConfig(
+        num_directions=cfg.num_directions,
+        optimization_steps=cfg.optimization_steps,
+    )
+    if prism_config:
+        for k, v in prism_config.items():
+            if hasattr(p_cfg, k):
+                setattr(p_cfg, k, v)
+    
+    # Train PRISM
+    print(f"\nTraining PRISM with {cfg.num_directions} directions...")
+    prism = PRISMMethod(config=p_cfg)
+    result = prism.train_prism(pair_set)
+    
+    # Extract directions (PRISM doesn't have learned weights, use uniform)
+    directions = result.directions
+    
+    print(f"Trained {len(directions)} layers")
+    for layer, dirs in directions.items():
+        print(f"  {layer}: {dirs.shape[0]} directions")
+    
+    # Bake into weights (no learned weights, use strategy)
+    print(f"\nBaking into weights with strategy: {cfg.combination_strategy}")
+    return bake_multi_directions(model, directions, None, cfg)
+
+
+def train_and_bake_pulse(
+    model: Module,
+    pair_set: "ContrastivePairSet",
+    config: Optional[MultiDirectionConfig] = None,
+    pulse_config: Optional[Dict[str, Any]] = None,
+) -> MultiDirectionResult:
+    """
+    Train PULSE and bake directions into model weights.
+    
+    PULSE uses single direction per layer with conditional gating.
+    When baking, we lose the gating but keep the behavior vectors.
+    
+    Args:
+        model: Model to modify
+        pair_set: ContrastivePairSet with activations
+        config: Multi-direction config
+        pulse_config: Additional PULSE-specific config
+        
+    Returns:
+        MultiDirectionResult
+    """
+    from wisent.core.steering_methods.methods.pulse import PULSEMethod, PULSEConfig
+    
+    cfg = config or MultiDirectionConfig(method="pulse")
+    
+    # Build PULSE config
+    pu_cfg = PULSEConfig(
+        optimization_steps=cfg.optimization_steps,
+    )
+    if pulse_config:
+        for k, v in pulse_config.items():
+            if hasattr(pu_cfg, k):
+                setattr(pu_cfg, k, v)
+    
+    # Train PULSE
+    print(f"\nTraining PULSE...")
+    pulse = PULSEMethod(config=pu_cfg)
+    result = pulse.train_pulse(pair_set)
+    
+    # Extract behavior vectors and layer weights
+    behavior_vectors = result.behavior_vectors
+    layer_scales = result.layer_scales
+    
+    print(f"Trained {len(behavior_vectors)} layers")
+    
+    # Convert to multi-direction format (single direction per layer)
+    directions = {}
+    weights = {}
+    for layer, vec in behavior_vectors.items():
+        directions[layer] = vec.unsqueeze(0)  # [1, hidden_dim]
+        if layer_scales and layer in layer_scales:
+            weights[layer] = torch.tensor([layer_scales[layer]])
+        else:
+            weights[layer] = torch.tensor([1.0])
+    
+    # Bake into weights
+    print(f"\nBaking into weights...")
+    return bake_multi_directions(model, directions, weights, cfg)
+
+
+def train_and_bake(
+    model: Module,
+    pair_set: "ContrastivePairSet",
+    method: str = "titan",
+    config: Optional[MultiDirectionConfig] = None,
+    **method_kwargs,
+) -> MultiDirectionResult:
+    """
+    Train specified method and bake into weights.
+    
+    Args:
+        model: Model to modify
+        pair_set: ContrastivePairSet with activations
+        method: 'titan', 'prism', or 'pulse'
+        config: Multi-direction config
+        **method_kwargs: Additional method-specific config
+        
+    Returns:
+        MultiDirectionResult
+    """
+    if config is None:
+        config = MultiDirectionConfig(method=method)
+    else:
+        config.method = method
+    
+    if method == "titan":
+        return train_and_bake_titan(model, pair_set, config, method_kwargs)
+    elif method == "prism":
+        return train_and_bake_prism(model, pair_set, config, method_kwargs)
+    elif method == "pulse":
+        return train_and_bake_pulse(model, pair_set, config, method_kwargs)
+    else:
+        raise ValueError(f"Unknown method: {method}. Use 'titan', 'prism', or 'pulse'")

wisent/core/weight_modification/utils.py

@@ -0,0 +1,236 @@
+"""
+Utility functions for weight modification.
+"""
+
+from __future__ import annotations
+
+import torch
+from typing import TYPE_CHECKING
+from wisent.core.cli_logger import setup_logger, bind
+
+if TYPE_CHECKING:
+    from torch.nn import Module
+    from torch import Tensor
+
+__all__ = [
+    "get_modifiable_components",
+    "verify_modification",
+    "compute_modification_metrics",
+]
+
+_LOG = setup_logger(__name__)
+
+
+def get_modifiable_components(model: Module) -> list[tuple[str, list[str]]]:
+    """
+    Get list of modifiable components in model.
+
+    Returns:
+        List of (component_path, subcomponents) tuples
+
+    Example:
+        >>> components = get_modifiable_components(model)
+        >>> for path, subcomps in components:
+        ...     print(f"{path}: {subcomps}")
+        self_attn: ['q_proj', 'k_proj', 'v_proj', 'o_proj']
+        mlp: ['gate_proj', 'up_proj', 'down_proj']
+    """
+    log = bind(_LOG)
+
+    # Get layers
+    if hasattr(model, "model"):
+        sample_layer = model.model.layers[0]
+    elif hasattr(model, "transformer"):
+        sample_layer = model.transformer.h[0]
+    else:
+        sample_layer = model.layers[0]
+
+    components = []
+
+    # Check attention
+    if hasattr(sample_layer, "self_attn"):
+        attn = sample_layer.self_attn
+        attn_subcomps = []
+        for attr in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+            if hasattr(attn, attr):
+                attn_subcomps.append(attr)
+        components.append(("self_attn", attn_subcomps))
+
+    # Check MLP
+    if hasattr(sample_layer, "mlp"):
+        mlp = sample_layer.mlp
+        mlp_subcomps = []
+        for attr in ["gate_proj", "up_proj", "down_proj", "fc1", "fc2"]:
+            if hasattr(mlp, attr):
+                mlp_subcomps.append(attr)
+        components.append(("mlp", mlp_subcomps))
+
+    log.info("Found modifiable components", extra={"components": components})
+
+    return components
+
+
+def verify_modification(
+    model: Module,
+    original_state_dict: dict | None = None,
+    expected_changes: int | None = None,
+) -> dict[str, any]:
+    """
+    Verify that model was modified correctly.
+
+    Args:
+        model: Modified model
+        original_state_dict: Optional original state dict to compare against
+        expected_changes: Expected number of parameter changes
+
+    Returns:
+        Dictionary with verification results:
+        - "verified": bool
+        - "params_changed": int
+        - "issues": list of issues found
+
+    Example:
+        >>> original_state = copy.deepcopy(model.state_dict())
+        >>> project_weights(model, steering_vectors)
+        >>> result = verify_modification(model, original_state, expected_changes=64)
+        >>> assert result["verified"]
+    """
+    log = bind(_LOG)
+
+    issues = []
+    params_changed = 0
+
+    if original_state_dict is not None:
+        current_state = model.state_dict()
+
+        for key in original_state_dict.keys():
+            if key not in current_state:
+                issues.append(f"Missing key in current state: {key}")
+                continue
+
+            orig = original_state_dict[key]
+            curr = current_state[key]
+
+            if orig.shape != curr.shape:
+                issues.append(f"Shape mismatch for {key}: {orig.shape} vs {curr.shape}")
+                continue
+
+            diff = (curr - orig).abs()
+            changed = (diff > 1e-6).sum().item()
+
+            if changed > 0:
+                params_changed += changed
+
+        if expected_changes is not None:
+            if params_changed != expected_changes:
+                issues.append(
+                    f"Expected {expected_changes} params changed, got {params_changed}"
+                )
+
+    # Check for NaN or Inf
+    for name, param in model.named_parameters():
+        if torch.isnan(param).any():
+            issues.append(f"NaN detected in {name}")
+        if torch.isinf(param).any():
+            issues.append(f"Inf detected in {name}")
+
+    verified = len(issues) == 0
+
+    result = {
+        "verified": verified,
+        "params_changed": params_changed,
+        "issues": issues,
+    }
+
+    if verified:
+        log.info("Modification verified successfully", extra=result)
+    else:
+        log.warning("Verification failed", extra=result)
+
+    return result
+
+
+def compute_modification_metrics(
+    original_model: Module,
+    modified_model: Module,
+    sample_inputs: Tensor | None = None,
+) -> dict[str, float]:
+    """
+    Compute metrics quantifying the modification.
+
+    Args:
+        original_model: Original model
+        modified_model: Modified model
+        sample_inputs: Optional sample inputs to test output differences
+
+    Returns:
+        Dictionary with metrics:
+        - "weight_l2_distance": L2 distance between weight matrices
+        - "weight_cosine_similarity": Cosine similarity of weights
+        - "output_difference": If sample_inputs provided, difference in outputs
+
+    Example:
+        >>> metrics = compute_modification_metrics(original, modified, sample_inputs)
+        >>> print(f"Weight distance: {metrics['weight_l2_distance']:.4f}")
+    """
+    log = bind(_LOG)
+
+    # Compare weights
+    orig_state = original_model.state_dict()
+    mod_state = modified_model.state_dict()
+
+    weight_diff_sum = 0.0
+    weight_norm_sum = 0.0
+    weight_dot_product = 0.0
+    weight_orig_norm = 0.0
+    weight_mod_norm = 0.0
+
+    for key in orig_state.keys():
+        if key not in mod_state:
+            continue
+
+        orig = orig_state[key].float()
+        mod = mod_state[key].float()
+
+        if orig.shape != mod.shape:
+            continue
+
+        diff = (mod - orig)
+        weight_diff_sum += diff.norm().item() ** 2
+        weight_norm_sum += orig.norm().item() ** 2
+
+        # Flatten for dot product
+        orig_flat = orig.flatten()
+        mod_flat = mod.flatten()
+
+        weight_dot_product += (orig_flat @ mod_flat).item()
+        weight_orig_norm += (orig_flat @ orig_flat).item()
+        weight_mod_norm += (mod_flat @ mod_flat).item()
+
+    l2_distance = weight_diff_sum ** 0.5
+    cosine_sim = weight_dot_product / ((weight_orig_norm * weight_mod_norm) ** 0.5)
+
+    metrics = {
+        "weight_l2_distance": l2_distance,
+        "weight_cosine_similarity": cosine_sim,
+    }
+
+    # Compare outputs if sample inputs provided
+    if sample_inputs is not None:
+        with torch.no_grad():
+            orig_outputs = original_model(sample_inputs)
+            mod_outputs = modified_model(sample_inputs)
+
+            if hasattr(orig_outputs, "logits"):
+                orig_logits = orig_outputs.logits
+                mod_logits = mod_outputs.logits
+            else:
+                orig_logits = orig_outputs
+                mod_logits = mod_outputs
+
+            output_diff = (mod_logits - orig_logits).abs().mean().item()
+            metrics["output_difference"] = output_diff
+
+    log.info("Computed modification metrics", extra=metrics)
+
+    return metrics