rustytorch/crates/core/rtx-interpret

RTX Interpret - Model Interpretability Library

GPU-accelerated model interpretability and explainability toolkit for RustyTorch++. This crate provides PyTorch Captum-equivalent functionality for understanding and interpreting neural network predictions.

Features

• Gradient-based Attribution: Saliency maps using numerical gradients
• Flexible Baseline Support: Zero, mean, custom, random, and Gaussian noise baselines
• Type-safe API: Comprehensive error handling with detailed error types
• Extensible Design: Core traits for implementing new attribution methods
• Full Test Coverage: Comprehensive unit and integration tests

Installation

Add to your Cargo.toml:

[dependencies]
rtx-interpret = { path = "path/to/rtx-interpret" }

Quick Start

use rtx_interpret::{Attribution, Saliency, SaliencyConfig};
use rtx_tensor::{Tensor, Device};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Device::cpu();
    let input = Tensor::randn(&[1, 3, 224, 224], &device)?;

    // Define your model's forward function
    let forward_fn = |x: &Tensor| -> rtx_tensor::Result<Tensor> {
        // Your model forward pass here
        Ok(x.clone()) // Placeholder
    };

    // Compute saliency map
    let config = SaliencyConfig::default();
    let saliency = Saliency::new(config);
    let attributions = saliency.attribute(&forward_fn, &input, Some(0))?;

    println!("Attribution shape: {:?}", attributions.attributions().shape());
    Ok(())
}

Architecture

Core Traits

Attribution

The main trait for attribution methods that assign importance scores to input features:

pub trait Attribution: Send + Sync {
    fn attribute(
        &self,
        forward_fn: &dyn Fn(&Tensor) -> rtx_tensor::Result<Tensor>,
        input: &Tensor,
        target_index: Option<usize>,
    ) -> Result<AttributionOutput>;
}

Attribution Methods

Saliency Maps

Computes gradients of model output with respect to input using numerical differentiation (central finite differences).

use rtx_interpret::{Saliency, SaliencyConfig};

// Basic usage
let config = SaliencyConfig::default();
let saliency = Saliency::new(config);

// With absolute value
let config = SaliencyConfig { abs: true };
let saliency = Saliency::new(config);

Implementation Details:

• Uses central finite differences for numerical gradient computation
• Epsilon: 1e-5 for perturbation
• Supports multi-dimensional inputs and batched computation
• Optional absolute value of gradients

Types

AttributionOutput

Contains attribution scores and optional metadata:

pub struct AttributionOutput {
    pub attributions: Tensor,
    pub metadata: Option<AttributionMetadata>,
}

Baseline

Defines baseline options for attribution methods:

pub enum Baseline {
    Zero,                                    // All zeros
    Mean(Tensor),                           // Dataset mean
    Custom(Tensor),                         // Custom baseline
    Random,                                  // Random values
    GaussianNoise { mean: f32, std: f32 }, // Gaussian noise
}

Error Handling

Comprehensive error types with detailed messages:

pub enum InterpretError {
    Tensor { message: String },
    Autograd { message: String },
    InvalidInput { message: String },
    InvalidConfig { message: String },
    Attribution { message: String },
    Perturbation { message: String },
    Baseline { message: String },
    ShapeMismatch { message: String },
    DeviceMismatch { message: String },
}

Examples

Linear Function Attribution

use rtx_interpret::{Attribution, Saliency, SaliencyConfig};
use rtx_tensor::{Tensor, Device};

let device = Device::cpu();
let input = Tensor::from_slice(&[1.0_f32, 2.0, 3.0], &[3], &device).unwrap();

let forward_fn = |x: &Tensor| -> rtx_tensor::Result<Tensor> {
    let two = Tensor::from_slice(&[2.0_f32], &[1], &device)?;
    x.mul(&two)
};

let saliency = Saliency::default();
let result = saliency.attribute(&forward_fn, &input, None).unwrap();

// Gradients should be approximately [2.0, 2.0, 2.0]

Multi-Output with Target Index

let forward_fn = |x: &Tensor| -> rtx_tensor::Result<Tensor> {
    let coeffs = Tensor::from_slice(&[2.0_f32, 3.0], &[2], &device)?;
    x.mul(&coeffs)
};

// Get gradient with respect to first output
let result = saliency.attribute(&forward_fn, &input, Some(0)).unwrap();

Using Baselines

use rtx_interpret::Baseline;

// Zero baseline
let baseline = Baseline::Zero;
let baseline_tensor = baseline.generate(&input)?;

// Gaussian noise baseline
let baseline = Baseline::GaussianNoise { mean: 0.0, std: 1.0 };
let baseline_tensor = baseline.generate(&input)?;

// Custom baseline
let custom = Tensor::full(&[2, 3], 0.5, &device)?;
let baseline = Baseline::Custom(custom);

Implementation Notes

Numerical Gradients

The current implementation uses central finite differences for gradient computation:

• Advantages: Works with any differentiable function, no autograd infrastructure required
• Epsilon: 1e-5 provides good balance between accuracy and numerical stability
• Accuracy: Typical error ~1e-2 to 0.1, acceptable for most interpretability tasks

Future versions may integrate more closely with rtx-autograd for exact gradients.

Performance

• Complexity: O(n) forward passes for n input features
• Memory: Efficient - no intermediate graph storage
• Optimization: Consider implementing batch processing for large inputs

Testing

Run the full test suite:

cargo test

Run specific test categories:

cargo test --lib          # Library tests
cargo test --test saliency_tests  # Integration tests
cargo test --doc          # Documentation tests

Future Enhancements

Planned features for future releases:

1. Integrated Gradients: Path integral attribution
2. GradCAM: Gradient-weighted class activation mapping
3. Occlusion: Perturbation-based attribution
4. Feature Ablation: Systematic feature removal
5. LIME: Local interpretable model-agnostic explanations
6. Layer/Neuron Attribution: Internal activation analysis
7. Exact Autograd Integration: Switch to rtx-autograd for exact gradients

Contributing

This crate follows strict development standards:

• TDD: All features developed test-first
• No unsafe code: Memory safe by design
• File size limit: Maximum 1000 lines per file
• No mocks/stubs: Full implementations only
• Comprehensive docs: All public APIs documented

License

MIT OR Apache-2.0

References

• Simonyan et al., "Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps" (2013)
• PyTorch Captum library: https://captum.ai/