rustytorch/crates/core/rtx-interpret/ATTENTION_MODULE.md

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:

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:

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:

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:

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

Usage Example

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:

# 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:

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