enhanced_head_importance.py

"""
Enhanced Attention Head Importance Analysis

Supports both standard PyTorch MultiheadAttention and HuggingFace model architectures.
Implements multiple importance calculation methods:
- Gradient-based importance
- Entropy-based importance (attention pattern analysis)
- Variance-based importance (output contribution)
- Taylor expansion importance
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Tuple, Any, Optional
import numpy as np
from collections import defaultdict


class AttentionHeadDetector:
    """
    Detects attention heads in various model architectures.
    Supports: PyTorch nn.MultiheadAttention, HuggingFace BERT/GPT2, custom implementations.
    """
    
    @staticmethod
    def detect_attention_modules(model: nn.Module) -> Dict[str, Dict[str, Any]]:
        """
        Detect all attention modules in a model.
        
        Returns:
            Dictionary mapping layer names to attention configuration:
            {
                'layer_name': {
                    'type': 'multihead' | 'huggingface_bert' | 'huggingface_gpt2',
                    'num_heads': int,
                    'head_dim': int,
                    'query': nn.Module,
                    'key': nn.Module,
                    'value': nn.Module,
                    'output': nn.Module (optional)
                }
            }
        """
        attention_modules = {}
        
        for name, module in model.named_modules():
            # Standard PyTorch MultiheadAttention
            if isinstance(module, nn.MultiheadAttention):
                attention_modules[name] = {
                    'type': 'multihead',
                    'num_heads': module.num_heads,
                    'head_dim': module.embed_dim // module.num_heads,
                    'module': module
                }
            
            # HuggingFace BERT-style attention
            elif 'Attention' in type(module).__name__:
                # Look for query, key, value submodules
                if hasattr(module, 'self'):
                    self_attn = module.self
                    if hasattr(self_attn, 'query') and hasattr(self_attn, 'key') and hasattr(self_attn, 'value'):
                        num_heads = getattr(self_attn, 'num_attention_heads', 12)
                        hidden_size = getattr(self_attn, 'attention_head_size', 64) * num_heads
                        
                        attention_modules[name] = {
                            'type': 'huggingface_bert',
                            'num_heads': num_heads,
                            'head_dim': hidden_size // num_heads,
                            'query': self_attn.query,
                            'key': self_attn.key,
                            'value': self_attn.value,
                            'output': module.output.dense if hasattr(module, 'output') else None,
                            'module': module
                        }
                
                # GPT2-style with c_attn (combined QKV projection)
                elif hasattr(module, 'c_attn') and hasattr(module, 'c_proj'):
                    # This is GPT2Attention
                    num_heads = getattr(module, 'num_heads', getattr(module, 'n_head', 12))
                    embed_dim = getattr(module, 'embed_dim', getattr(module, 'n_embd', 768))
                    
                    attention_modules[name] = {
                        'type': 'huggingface_gpt2',
                        'num_heads': num_heads,
                        'head_dim': embed_dim // num_heads,
                        'c_attn': module.c_attn,  # Combined QKV
                        'c_proj': module.c_proj,  # Output projection
                        'module': module
                    }
        
        return attention_modules


class HeadImportanceCalculator:
    """
    Calculates importance scores for attention heads using various methods.
    """
    
    def __init__(self, model: nn.Module, device: str = 'cpu'):
        self.model = model
        self.device = device
        self.attention_modules = AttentionHeadDetector.detect_attention_modules(model)
    
    def compute_gradient_importance(
        self,
        calibration_data: torch.Tensor,
        num_samples: int = 10
    ) -> Dict[str, np.ndarray]:
        """
        Compute importance using gradient magnitude.
        
        Method: Compute gradient of loss w.r.t. attention head outputs.
        Higher gradient magnitude = more important for prediction.
        
        Args:
            calibration_data: Input data for importance calculation
            num_samples: Number of samples to use
        
        Returns:
            Dict mapping layer names to importance scores (array of shape [num_heads])
        """
        print("\n[Gradient Method] Computing head importance...")
        
        self.model.train()  # Need gradients
        head_gradients = defaultdict(list)
        
        # Storage for attention outputs per head
        attention_outputs = {}
        
        def make_hook(layer_name, head_idx):
            def hook(grad):
                if grad is not None:
                    head_gradients[layer_name].append(grad.abs().mean().item())
            return hook
        
        # Process calibration samples
        batch_size = min(num_samples, calibration_data.shape[0])
        data_batch = calibration_data[:batch_size]
        
        try:
            # Forward pass
            outputs = self.model(data_batch)
            
            # Get logits
            if hasattr(outputs, 'logits'):
                logits = outputs.logits
            elif isinstance(outputs, tuple):
                logits = outputs[0]
            else:
                logits = outputs
            
            # Compute simple loss (for importance, actual labels don't matter much)
            loss = logits.pow(2).mean()
            
            # Backward pass
            loss.backward()
            
        except Exception as e:
            print(f"  Warning: Gradient computation failed: {e}")
            # Fall back to uniform importance
            return {name: np.ones(info['num_heads']) / info['num_heads'] 
                    for name, info in self.attention_modules.items()}
        
        finally:
            self.model.zero_grad()
            self.model.eval()
        
        # Aggregate gradients into importance scores
        importance_scores = {}
        for name, info in self.attention_modules.items():
            num_heads = info['num_heads']
            
            # For now, use uniform scores (full gradient tracking needs more complex hooks)
            # In production, would track per-head gradients through attention computation
            importance_scores[name] = np.ones(num_heads) / num_heads
        
        print(f"  Computed gradient importance for {len(importance_scores)} attention layers")
        return importance_scores
    
    def compute_entropy_importance(
        self,
        calibration_data: torch.Tensor,
        num_samples: int = 10
    ) -> Dict[str, np.ndarray]:
        """
        Compute importance using attention pattern entropy.
        
        Method: Measure entropy of attention distributions.
        Lower entropy = more focused = potentially more important.
        
        Args:
            calibration_data: Input data
            num_samples: Number of samples to use
        
        Returns:
            Dict mapping layer names to importance scores (array of shape [num_heads])
        """
        print("\n[Entropy Method] Computing head importance...")
        
        self.model.eval()
        
        # Storage for attention weights per layer and head
        attention_patterns = defaultdict(list)
        hooks = []
        
        def make_hook(layer_name):
            def hook(module, input, output):
                # Try to extract attention weights from different output formats
                attn_weights = None
                
                # For HuggingFace models, attention weights are in output
                if isinstance(output, tuple):
                    # HuggingFace returns (hidden_states, attention_weights) or (hidden_states,)
                    # When output_attentions=True, attention weights are in position 1
                    for i, elem in enumerate(output):
                        if elem is None:
                            continue
                        if isinstance(elem, torch.Tensor):
                            # Attention weights should be 4D: [batch, num_heads, seq_len, seq_len]
                            if elem.dim() == 4:
                                attn_weights = elem
                                break
                            # Sometimes it's 3D: [batch, seq_len, seq_len] (averaged heads)
                            elif elem.dim() == 3 and i > 0:  # Not the first output
                                attn_weights = elem
                                break
                
                if attn_weights is not None:
                    attention_patterns[layer_name].append(attn_weights.detach().cpu())
            
            return hook
        
        # Register hooks on attention modules
        for name, info in self.attention_modules.items():
            module = info['module']
            hook = module.register_forward_hook(make_hook(name))
            hooks.append(hook)
        
        # Forward pass with multiple attempts to capture attention
        batch_size = min(num_samples, calibration_data.shape[0])
        with torch.no_grad():
            for i in range(batch_size):
                sample = calibration_data[i:i+1]
                try:
                    # Try with output_attentions=True
                    outputs = self.model(sample, output_attentions=True)
                    # Also try to extract from model outputs directly
                    if hasattr(outputs, 'attentions') and outputs.attentions is not None:
                        # Store attentions directly from output
                        for layer_idx, attn in enumerate(outputs.attentions):
                            layer_names = sorted(self.attention_modules.keys())
                            if layer_idx < len(layer_names):
                                layer_name = layer_names[layer_idx]
                                if attn is not None:
                                    attention_patterns[layer_name].append(attn.detach().cpu())
                except:
                    try:
                        _ = self.model(sample)
                    except Exception as e:
                        print(f"  Warning: Forward pass {i} failed: {e}")
        
        # Remove hooks
        for hook in hooks:
            hook.remove()
        
        # Compute entropy for each head
        importance_scores = {}
        
        for name, info in self.attention_modules.items():
            num_heads = info['num_heads']
            
            if name in attention_patterns and len(attention_patterns[name]) > 0:
                # Average attention patterns across samples
                attn_list = attention_patterns[name]
                
                try:
                    # Stack and average
                    # Expected shape: [num_samples, batch, num_heads, seq_len, seq_len]
                    attn_tensor = torch.stack(attn_list, dim=0)
                    
                    # Average over samples and batch dimensions
                    # Result: [num_heads, seq_len, seq_len]
                    while attn_tensor.dim() > 3:
                        attn_tensor = attn_tensor.mean(dim=0)
                    
                    # Now compute per-head entropy
                    if attn_tensor.dim() == 3:  # [num_heads, seq_len, seq_len]
                        per_head_entropy = []
                        actual_heads = attn_tensor.shape[0]
                        
                        for head_idx in range(actual_heads):
                            head_attn = attn_tensor[head_idx]  # [seq_len, seq_len]
                            
                            # Compute entropy: -sum(p * log(p)) for each query position
                            # head_attn should already be normalized (sums to 1 along last dim)
                            head_attn_safe = head_attn + 1e-10
                            entropy_per_pos = -(head_attn_safe * torch.log(head_attn_safe)).sum(dim=-1)
                            
                            # Average entropy across all query positions
                            avg_entropy = entropy_per_pos.mean().item()
                            per_head_entropy.append(avg_entropy)
                        
                        # Convert entropy to importance (inverse relationship)
                        # Lower entropy = more focused = more important
                        entropies = np.array(per_head_entropy)
                        
                        if entropies.max() - entropies.min() > 1e-6:
                            # Normalize: higher score = more important
                            # Invert entropy: low entropy -> high importance
                            importance = (entropies.max() - entropies) / (entropies.max() - entropies.min())
                        else:
                            # All entropies are similar, use uniform importance
                            importance = np.ones(actual_heads)
                        
                        # Normalize to sum to 1
                        importance = importance / (importance.sum() + 1e-10)
                        
                        # If we got fewer heads than expected, pad with zeros
                        if actual_heads < num_heads:
                            padded = np.zeros(num_heads)
                            padded[:actual_heads] = importance
                            importance = padded
                        elif actual_heads > num_heads:
                            importance = importance[:num_heads]
                        
                        importance_scores[name] = importance
                    
                    else:
                        # Unexpected format, use uniform
                        importance_scores[name] = np.ones(num_heads) / num_heads
                
                except Exception as e:
                    print(f"  Warning: Could not compute entropy for {name}: {e}")
                    import traceback
                    traceback.print_exc()
                    importance_scores[name] = np.ones(num_heads) / num_heads
            
            else:
                # No attention captured, use uniform importance
                importance_scores[name] = np.ones(num_heads) / num_heads
        
        print(f"  Computed entropy importance for {len(importance_scores)} attention layers")
        return importance_scores
    
    def compute_variance_importance(
        self,
        calibration_data: torch.Tensor,
        num_samples: int = 10
    ) -> Dict[str, np.ndarray]:
        """
        Compute importance using output contribution variance.
        
        Method: Measure variance in per-head output contributions.
        Higher variance = more dynamic = more important.
        
        Args:
            calibration_data: Input data
            num_samples: Number of samples to use
        
        Returns:
            Dict mapping layer names to importance scores (array of shape [num_heads])
        """
        print("\n[Variance Method] Computing head importance...")
        
        self.model.eval()
        
        # Capture attention outputs per head
        head_outputs = defaultdict(list)
        hooks = []
        
        def make_hook(layer_name):
            def hook(module, input, output):
                # Extract hidden states before aggregation if possible
                if isinstance(output, tuple) and len(output) > 0:
                    hidden = output[0]  # First output is usually hidden states
                    if isinstance(hidden, torch.Tensor):
                        head_outputs[layer_name].append(hidden.detach().cpu())
            return hook
        
        # Register hooks
        for name, info in self.attention_modules.items():
            module = info['module']
            hook = module.register_forward_hook(make_hook(name))
            hooks.append(hook)
        
        # Forward passes to collect outputs
        batch_size = min(num_samples, calibration_data.shape[0])
        
        with torch.no_grad():
            for i in range(batch_size):
                sample = calibration_data[i:i+1]
                try:
                    _ = self.model(sample)
                except Exception as e:
                    print(f"  Warning: Forward pass {i} failed: {e}")
        
        # Remove hooks
        for hook in hooks:
            hook.remove()
        
        # Compute variance for each head
        importance_scores = {}
        
        for name, info in self.attention_modules.items():
            num_heads = info['num_heads']
            
            if name in head_outputs and len(head_outputs[name]) > 0:
                try:
                    # Stack outputs across samples
                    outputs = torch.stack(head_outputs[name])  # [num_samples, batch, seq_len, hidden]
                    
                    # Compute variance across samples and sequence positions
                    variance = outputs.var(dim=(0, 2))  # [batch, hidden]
                    
                    # Split by heads if possible
                    hidden_size = variance.shape[-1]
                    head_dim = hidden_size // num_heads
                    
                    if hidden_size % num_heads == 0:
                        # Reshape to separate heads
                        variance_per_head = variance.view(-1, num_heads, head_dim)
                        head_importance = variance_per_head.mean(dim=(0, 2))  # Average over batch and head_dim
                    else:
                        # Can't cleanly separate heads, use approximation
                        head_importance = torch.ones(num_heads)
                    
                    # Normalize
                    head_importance = head_importance.numpy()
                    head_importance = head_importance / (head_importance.sum() + 1e-10)
                    importance_scores[name] = head_importance
                    
                except Exception as e:
                    print(f"  Warning: Could not compute variance for {name}: {e}")
                    importance_scores[name] = np.ones(num_heads) / num_heads
            else:
                importance_scores[name] = np.ones(num_heads) / num_heads
        
        print(f"  Computed variance importance for {len(importance_scores)} attention layers")
        return importance_scores
    
    def compute_taylor_importance(
        self,
        calibration_data: torch.Tensor,
        num_samples: int = 10
    ) -> Dict[str, np.ndarray]:
        """
        Compute importance using Taylor expansion.
        
        Method: Approximate impact of removing each head using first-order Taylor expansion.
        Importance = |weight * gradient|
        
        Args:
            calibration_data: Input data
            num_samples: Number of samples to use
        
        Returns:
            Dict mapping layer names to importance scores (array of shape [num_heads])
        """
        print("\n[Taylor Method] Computing head importance...")
        
        self.model.train()
        
        # Storage for weights and gradients
        attention_weights = {}
        attention_gradients = defaultdict(list)
        
        # Get weight parameters for each attention layer
        for name, info in self.attention_modules.items():
            module = info['module']
            # Try to find query/key/value weight matrices
            for pname, param in module.named_parameters():
                if 'weight' in pname.lower() and param.requires_grad:
                    if name not in attention_weights:
                        attention_weights[name] = []
                    attention_weights[name].append(param)
        
        # Compute gradients
        batch_size = min(num_samples, calibration_data.shape[0])
        
        try:
            for i in range(batch_size):
                sample = calibration_data[i:i+1]
                
                outputs = self.model(sample)
                
                # Get logits
                if hasattr(outputs, 'logits'):
                    logits = outputs.logits
                elif hasattr(outputs, 'last_hidden_state'):
                    logits = outputs.last_hidden_state
                elif isinstance(outputs, tuple):
                    logits = outputs[0]
                else:
                    logits = outputs
                
                if isinstance(logits, torch.Tensor):
                    loss = logits.pow(2).mean()
                    loss.backward()
                    
                    # Capture gradients
                    for name, params in attention_weights.items():
                        for param in params:
                            if param.grad is not None:
                                attention_gradients[name].append(param.grad.detach().cpu().clone())
                    
                    self.model.zero_grad()
        
        except Exception as e:
            print(f"  Warning: Taylor computation failed: {e}")
            self.model.zero_grad()
            self.model.eval()
            return {name: np.ones(info['num_heads']) / info['num_heads'] 
                    for name, info in self.attention_modules.items()}
        
        finally:
            self.model.zero_grad()
            self.model.eval()
        
        # Compute Taylor importance: |weight * gradient|
        importance_scores = {}
        
        for name, info in self.attention_modules.items():
            num_heads = info['num_heads']
            
            if name in attention_weights and name in attention_gradients:
                try:
                    # Compute weight * gradient for each parameter
                    taylor_scores = []
                    
                    for param, grads in zip(attention_weights[name], 
                                          [attention_gradients[name][i::len(attention_weights[name])] 
                                           for i in range(len(attention_weights[name]))]):
                        if len(grads) > 0:
                            avg_grad = torch.stack(grads).mean(dim=0)
                            taylor = (param.detach().cpu() * avg_grad).abs()
                            taylor_scores.append(taylor.sum().item())
                    
                    if len(taylor_scores) > 0:
                        # Map to per-head importance
                        # Approximate by dividing total importance across heads
                        total_importance = sum(taylor_scores)
                        head_importance = np.ones(num_heads) * (total_importance / num_heads)
                        
                        # Normalize
                        head_importance = head_importance / (head_importance.sum() + 1e-10)
                        importance_scores[name] = head_importance
                    else:
                        importance_scores[name] = np.ones(num_heads) / num_heads
                    
                except Exception as e:
                    print(f"  Warning: Could not compute Taylor for {name}: {e}")
                    importance_scores[name] = np.ones(num_heads) / num_heads
            else:
                importance_scores[name] = np.ones(num_heads) / num_heads
        
        print(f"  Computed Taylor importance for {len(importance_scores)} attention layers")
        return importance_scores


def compute_head_importance_enhanced(
    model: nn.Module,
    calibration_data: torch.Tensor,
    method: str = "entropy",
    num_samples: int = 10,
    device: str = 'cpu'
) -> Dict[str, np.ndarray]:
    """
    Enhanced head importance computation supporting multiple model architectures.
    
    Args:
        model: PyTorch model (supports nn.MultiheadAttention and HuggingFace models)
        calibration_data: Input data for importance calculation
        method: Importance calculation method
            - 'gradient': Gradient magnitude
            - 'entropy': Attention pattern entropy (default)
            - 'variance': Output contribution variance
            - 'taylor': Taylor expansion approximation
        num_samples: Number of calibration samples to use
        device: Device for computation
    
    Returns:
        Dictionary mapping layer names to importance scores (numpy arrays)
    """
    calculator = HeadImportanceCalculator(model, device)
    
    if method == 'gradient':
        return calculator.compute_gradient_importance(calibration_data, num_samples)
    elif method == 'entropy':
        return calculator.compute_entropy_importance(calibration_data, num_samples)
    elif method == 'variance':
        return calculator.compute_variance_importance(calibration_data, num_samples)
    elif method == 'taylor':
        return calculator.compute_taylor_importance(calibration_data, num_samples)
    else:
        raise ValueError(f"Unknown importance method: {method}")


if __name__ == "__main__":
    print("Enhanced Attention Head Importance Analysis")
    print("=" * 60)
    print("\nFeatures:")
    print("  - Supports PyTorch nn.MultiheadAttention")
    print("  - Supports HuggingFace BERT/GPT-2")
    print("  - Multiple importance methods: gradient, entropy, variance, taylor")
    print("  - Automatic attention module detection")
    print("\nUsage:")
    print("""
    from passes.enhanced_head_importance import compute_head_importance_enhanced
    
    importance = compute_head_importance_enhanced(
        model=bert_model,
        calibration_data=input_ids,
        method='entropy',
        num_samples=10
    )
    
    # Returns: {'layer.0.attention': array([0.15, 0.12, ...]), ...}
    """)