# Attention Analysis Module

Comprehensive attention analysis functionality for transformer interpretability in rtx-interpret.

## Overview

The attention module provides three complementary methods for analyzing transformer attention mechanisms:

1. **Attention Rollout** - Visualizes information flow through attention layers
2. **Attention Flow** - Computes token importance via max-flow analysis
3. **Head Importance** - Identifies critical attention heads

## Implementation Summary

### File Structure

```
src/attention/
├── mod.rs              (209 lines) - Module documentation and exports
├── rollout.rs          (457 lines) - Attention Rollout implementation
├── flow.rs             (343 lines) - Attention Flow implementation
└── head_importance.rs  (495 lines) - Head Importance analysis

tests/
└── attention_tests.rs  (500 lines) - Comprehensive integration tests
```

**Total: 2,004 lines** (all files < 1000 lines ✓)

### Test Coverage

- **50 total tests** (18 integration + 32 unit tests)
- **100% pass rate**
- Tests cover:
  - All three analysis methods
  - Multiple configuration options
  - Error handling and validation
  - Edge cases (single layer, large sequences, etc.)

## Features

### 1. Attention Rollout

Computes accumulated attention flow from input to output by recursively multiplying attention matrices while accounting for residual connections.

**Key Features:**
- Head fusion strategies: Mean, Max, Min
- Configurable discard ratio for noise filtering
- Row normalization for probability distributions
- Identity matrix addition for residual connections

**API:**
```rust
let config = AttentionRolloutConfig {
    discard_ratio: 0.1,
    head_fusion: HeadFusion::Mean,
};
let rollout = AttentionRollout::new(config);
let flow = rollout.compute(&attention_weights)?; // [seq_len, seq_len]
```

### 2. Attention Flow

Treats attention as a flow network and computes maximum flow from target token to input tokens.

**Key Features:**
- Max-flow inspired algorithm
- Backward propagation through layers
- Normalized importance scores [0, 1]
- Configurable flow threshold

**API:**
```rust
let config = AttentionFlowConfig::default();
let flow = AttentionFlow::new(config);
let importance = flow.compute(&attention_weights, 0)?; // [seq_len]
```

### 3. Head Importance

Measures the importance of individual attention heads using multiple methods.

**Methods:**
- **Gradient**: Importance based on gradient magnitude (L2 norm)
- **Taylor Expansion**: First-order Taylor approximation
- **Entropy Surprise**: Lower entropy = higher importance

**API:**
```rust
let config = HeadImportanceConfig {
    method: HeadImportanceMethod::Gradient,
};
let head_imp = HeadImportance::new(config);
let scores = head_imp.compute(&attention_weights, Some(&gradients))?; // [layers, heads]
```

## Algorithm Details

### Attention Rollout Algorithm

1. **Head Fusion**: Aggregate multi-head attention using Mean/Max/Min
2. **Add Identity**: Add identity matrix for residual connections
3. **Normalize**: Normalize rows to sum to 1
4. **Accumulate**: Multiply matrices layer by layer
   - `A_total = A_L × A_{L-1} × ... × A_1`

### Attention Flow Algorithm

1. **Initialize**: Start with flow=1.0 at source token in final layer
2. **Backward Propagation**: For each layer (reverse order):
   - Distribute flow according to attention weights
   - `flow[from] += flow[to] × attention[to][from]`
3. **Normalize**: Scale to [0, 1] range

### Head Importance Algorithms

**Gradient Method:**
```
importance[layer, head] = ||∇attention[layer, head]||_2
```

**Taylor Method:**
```
importance[layer, head] = Σ|attention[i,j] × gradient[i,j]|
```

**Entropy Method:**
```
entropy = -Σ p_ij × log(p_ij)
importance = 1 / (entropy + ε)
```

## Integration with rtx-interpret

The attention module is fully integrated into rtx-interpret:

```rust
// Available at crate root
use rtx_interpret::{
    AttentionRollout, AttentionRolloutConfig, HeadFusion,
    AttentionFlow, AttentionFlowConfig,
    HeadImportance, HeadImportanceConfig, HeadImportanceMethod,
};
```

## Usage Example

```rust
use rtx_interpret::{AttentionRollout, AttentionRolloutConfig, HeadFusion};
use rtx_tensor::{Device, Tensor};

fn analyze_transformer_attention() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();

    // Extract attention from your transformer model
    // Shape: [num_layers, num_heads, seq_len, seq_len]
    let attention_weights = extract_attention_from_model()?;

    // Analyze attention flow
    let config = AttentionRolloutConfig {
        discard_ratio: 0.1,
        head_fusion: HeadFusion::Mean,
    };
    let rollout = AttentionRollout::new(config);
    let attention_flow = rollout.compute(&attention_weights)?;

    // Visualize or save results
    visualize_attention_flow(&attention_flow)?;

    Ok(())
}
```

## Design Principles

1. **Modular**: Works with pre-computed attention matrices (no model required)
2. **Testable**: Comprehensive test suite with 100% pass rate
3. **Documented**: Full API documentation with examples
4. **Type-safe**: Strong typing with custom error types
5. **Performance**: Efficient CPU implementations
6. **Standards-compliant**: All files < 1000 lines, no mocks/stubs/todos

## References

- Abnar & Zuidema (2020): "Quantifying Attention Flow in Transformers"
- Michel et al. (2019): "Are Sixteen Heads Really Better than One?"
- Voita et al. (2019): "Analyzing Multi-Head Self-Attention"

## Testing

Run tests:
```bash
# Integration tests
cargo test -p rtx-interpret --test attention_tests

# Unit tests
cargo test -p rtx-interpret --lib attention::

# All attention tests
cargo test -p rtx-interpret attention
```

Run example:
```bash
cargo run -p rtx-interpret --example attention_analysis
```

## Future Enhancements

Potential additions (not implemented):

- GPU acceleration for large attention matrices
- Attention pattern visualization utilities
- Attention pruning recommendations
- Cross-layer attention aggregation strategies
- Integration with popular transformer architectures