TimeSAE: Sparse Decoding for Faithful Explanations of Black-Box Time Series Models

Official implementation of TimeSAE, accepted at ICML 2026. TimeSAE is a novel framework for explaining black-box time series models using Sparse Autoencoders with complete functional ANOVA decomposition and temporal convolution layers.

A PyTorch implementation is provided, with a JAX version available here.

Features

• Complete ANOVA Decomposition: Full functional decomposition from order 0 to D for comprehensive concept interactions
• Temporal Convolution Layers: Multi-layer dilated convolutions for capturing time series dynamics
• Distributed Training: Optimized for training on 3 A100 GPUs using PyTorch DDP
• JumpReLU & TopK Activations: Advanced sparse activation functions for better concept learning
• Faithfulness Evaluation: Comprehensive metrics for explanation quality
• Concept Interpretation: Automated concept naming and analysis tools
• Pretrained Checkpoints: Ready-to-use models for 6 black-box architectures × 6 datasets

Results of TimeSAE (with TopK activation of concepts) on Black-Box Pretrained Models [e.g. TimeGPT-1 (nxtla), Chronos (Amazon), TimeFM (google), Transformer (Huggin-Face)]

Results of TimeSAE (with JumpReLU activation of concepts) on Black-Box Pretrained Models [e.g. TimeGPT-1 (nxtla), Chronos (Amazon), TimeFM (google), Transformer (Huggin-Face)]

Architecture Overview

TimeSAE Components

1. Temporal Encoder: Multi-layer dilated convolutions for time series feature extraction
2. Sparse Encoder: Linear transformation with JumpReLU/TopK activations
3. Complete ANOVA Decoder: Full functional decomposition using all interaction orders
4. Loss Functions: Combined reconstruction, sparsity, consistency, and contrastive losses

Key Innovations

• Full ANOVA Decomposition: Unlike standard approaches, implements complete functional ANOVA from order 0 to concept_dim
• Temporal Convolutions: Captures temporal dynamics with dilated convolutions (dilations: 1, 2, 4, 8)
• Distributed Training: Scales efficiently across multiple A100 GPUs
• Counterfactual Learning: Incorporates causal reasoning through contrastive learning

Project Updates

• [2024-06-03] 🚀 Released model checkpoints for TimeSAE (trained on EliteLJ and ETTh1).
• [2024-05-04] Released the EliteLJ dataset with annotated explanations.
• [2024-05-04] Added Google Colab notebooks for easier use:
  see: /notebooks/google-colab/
• [2024-05-02] Initial release of TimeSAE.
• [2024-04-20] Completed ablation study on concept masking.
• [2024-04-10] Integrated support for multi-scale feature extraction.
• [2024-03-28] Finished cross-dataset evaluation on ETTh1 and WeatherBench.
• [2024-03-15] Added counterfactual explanation module.

Installation

git clone <repository-url>
cd TimeSAE
pip install -r requirements.txt

🚀 Quick Start with Pretrained Models

Download Pretrained Checkpoints

We provide pretrained TimeSAE models for all paper experiments. No need to train from scratch!

# List available models
python checkpoints/download_checkpoints.py --list

# Download specific model
python checkpoints/download_checkpoints.py --model transformer --dataset ecg

# Download all datasets for a model
python checkpoints/download_checkpoints.py --model chronos --all-datasets

# Download all models for a dataset
python checkpoints/download_checkpoints.py --dataset freqshapes --all-models

Anonymous Download Link: https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/pretrained/

Instant Explanation Generation

from utils.checkpoint_loader import load_explanation_pipeline
from utils import TimeSeriesVisualizer

# Load pretrained TimeSAE explaining Transformer on ECG data
timesae, transformer_model = load_explanation_pipeline('transformer', 'ecg')

# Generate sample ECG data (or use your own)
import torch
ecg_data = torch.randn(5, 1, 187)  # 5 samples, 1 lead, 187 timesteps

# Generate explanations instantly
concepts, x_recon, contributions = timesae(ecg_data)

# Get black-box predictions
predictions = transformer_model(ecg_data)

print(f"Concept activations: {concepts.shape}")
print(f"ANOVA decomposition orders: {list(contributions.keys())}")

# Visualize explanations
visualizer = TimeSeriesVisualizer()
visualizer.plot_time_series_with_concepts(
    ecg_data, x_recon, concepts, contributions,
    sample_idx=0, save_path='ecg_explanation.png'
)

Available Pretrained Models

Datasets: FreqShapes, SeqComb-UV, ECG, PAM, ETTH-1, ETTH-2

Paper Parameters

Based on the TimeSAE paper, the following parameters were used:

Model Architecture

• Concept Dimension (d): r × (D × T) where r = 2.0 (expansion ratio)
• Temporal Hidden Dim: 128
• Kernel Size: 3 for all temporal convolutions
• Dilation Rates: [1, 2, 4, 8] for 4-layer temporal encoder
• Dropout: 0.1

Loss Function Weights

• Sparsity Coefficient (η): 0.1
• Consistency Weight (α): 0.5 (fixed across all experiments)
• Contrastive Weight (λ): 0.1
• Temperature (τ): 0.1
• Label Fidelity Weight: 1.0

Training Parameters

• Learning Rate: 1e-3
• Weight Decay: 1e-4
• Batch Size: 32 per GPU (effective batch size: 96 on 3 GPUs)
• Optimizer: AdamW with β = (0.9, 0.999)
• Scheduler: CosineAnnealingLR with T_max = 100000, eta_min = 1e-6
• Gradient Clipping: max_norm = 1.0

Sparse Activation Parameters

• JumpReLU Threshold (φ): Initialized with U(0.01, 0.11)
• TopK Gamma: Start at 2.0, decay to 1.0 during training
• TopK Sparsity: k = concept_dim // 10 (10% activation)

Usage Examples

Basic Usage

import torch
from models import TimeSAE
from trainer import TimeSAETrainer, setup_model_and_optimizer
from utils import DatasetProcessor

# Create model with paper parameters
input_shape = (10, 128)  # (D=10 features, T=128 timesteps)
model, loss_fn, optimizer, scheduler = setup_model_and_optimizer(
    input_shape=input_shape,
    concept_dim=int(2.0 * 10 * 128),  # r=2.0 expansion ratio
    learning_rate=1e-3,
    weight_decay=1e-4
)

# Generate synthetic data
processor = DatasetProcessor()
data = processor.create_synthetic_data(
    n_samples=1000,
    n_features=10,
    n_timesteps=128
)

# Train model
trainer = TimeSAETrainer(model, loss_fn, optimizer, scheduler)
# trainer.train(train_dataloader, val_dataloader, num_epochs=100)

Using Pretrained Models

# Quick demo with any model/dataset combination
from utils.checkpoint_loader import quick_explanation_demo

# Try different combinations
quick_explanation_demo('transformer', 'ecg')      # Medical data
quick_explanation_demo('chronos', 'freqshapes')   # Synthetic patterns
quick_explanation_demo('timegpt', 'pam')          # Activity recognition

# Compare all models on same dataset
from utils.checkpoint_loader import compare_models_on_dataset
results = compare_models_on_dataset('ecg')

Explanation Generation

from utils.checkpoint_loader import load_timesae
from utils import ConceptInterpreter

# Load specific model
timesae = load_timesae('transformer', 'ecg')

# Generate explanations
x = torch.randn(5, 1, 187)  # ECG data
concepts, x_recon, contributions = timesae(x)

# Interpret concepts
interpreter = ConceptInterpreter()
analysis = interpreter.analyze_concept_activations(concepts)
interpreter.assign_concept_names(analysis)
summary = interpreter.get_concept_summary()

print("Top concepts:")
for name, description in list(summary.items())[:5]:
    print(f"  {name}: {description}")

Visualization

from utils import TimeSeriesVisualizer

visualizer = TimeSeriesVisualizer()

# Plot explanations
visualizer.plot_time_series_with_concepts(
    x, x_recon, concepts, contributions,
    sample_idx=0, top_k_concepts=10,
    save_path='explanation.png'
)

# Plot concept heatmap
visualizer.plot_concept_heatmap(
    concepts, save_path='concepts.png'
)

Running the Code

Single GPU Training

python main.py \
    --n_features 10 \
    --n_timesteps 128 \
    --batch_size 32 \
    --learning_rate 1e-3 \
    --num_epochs 100 \
    --synthetic_data \
    --use_black_box \
    --generate_explanations

Multi-GPU Training (3 A100s)

chmod +x run_distributed.sh
./run_distributed.sh

SLURM Distributed Training

Use the provided SLURM script for cluster environments:

sbatch run_slurm_3a100.slurm

Black-Box Model Support

Pretrained Models (Ready to Use)

• Chronos: Amazon's transformer-based foundation model
• TimeGPT: Nixtla's time series foundation model (API + local)
• TimeFM: Google's decoder-only foundation model

Trainable Models (Train from Scratch)

• Transformer: Standard transformer with time series adaptations
• PatchTS: Patch-based transformer architecture
• DLinear: Linear models with time series decomposition

# Train black-box models
python trainers/blackbox_trainer/train_transformer.py --seq_len 128 --n_vars 10
python trainers/blackbox_trainer/train_patchts.py --patch_len 16 --stride 8
python trainers/blackbox_trainer/train_dlinear.py --seq_len 96 --task_type forecasting

Dataset Configuration

Synthetic Datasets (Paper)

• FreqShapes: Frequency-based synthetic patterns
• SeqComb-UV: Sequential combinations (univariate)

Real-world Datasets (Paper)

• ECG: Arrhythmia detection (classification)
• PAM: Human activity recognition (classification)
• ETTH-1/ETTH-2: Energy demand prediction (regression)
• EliteLJ: Sports performance analysis (regression)

Custom Dataset Integration

from torch.utils.data import Dataset, DataLoader

class CustomTimeSeriesDataset(Dataset):
    def __init__(self, data, labels=None):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.labels = labels
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        if self.labels is not None:
            return self.data[idx], self.labels[idx]
        return self.data[idx]

# Use with TimeSAE
dataset = CustomTimeSeriesDataset(your_data, your_labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

Evaluation Metrics

The implementation includes all metrics from the paper:

Explanation Quality

• AUPRC: Area Under Precision-Recall Curve
• AUP: Area Under Precision
• AUR: Area Under Recall

Faithfulness

• Fx Score: ||f(x) - f(ex-)||² where concepts are removed
• Counterfactual Consistency: Measures causal effect ordering

Model Quality

• Reconstruction MSE: Mean squared error between original and reconstructed
• Sparsity Metrics: L0/L1 norms, activation ratios
• Concept Diversity: Mutual activation rates, entropy measures

Paper Results Reproduction

Expected Performance (AUPRC scores)

Faithfulness Scores (Fx metric)

Complete Reproduction

# Run all paper experiments
chmod +x scripts/reproduce_all_results.sh
./scripts/reproduce_all_results.sh

# Or use SLURM for cluster
sbatch run_slurm_3a100.slurm

TimeSAE lib Structure

TimeSAE/
├── models.py                     # Core TimeSAE architecture
├── loss.py                       # Loss functions and counterfactual generation
├── trainer.py                    # Distributed training logic
├── utils.py                      # Evaluation, visualization, and utilities
├── main.py                       # Main training script
├── blackbox_models/              # All black-box model implementations
│   ├── chronos.py               # Chronos (pretrained)
│   ├── timegpt.py               # TimeGPT (pretrained) 
│   ├── timefm.py                # TimeFM (pretrained)
│   ├── transformer.py           # Transformer (trainable)
│   ├── patchts.py               # PatchTS (trainable)
│   └── dlinear.py               # DLinear (trainable)
├── trainers/blackbox_trainer/    # Training scripts for black-box models
├── checkpoints/                  # Pretrained model checkpoints
│   ├── README.md                # Checkpoint documentation
│   ├── timesae_transformer/     # TimeSAE explaining Transformers
│   ├── timesae_chronos/         # TimeSAE explaining Chronos
│   └── blackbox_models/         # Trained black-box models
├── scripts/
│   ├── download_checkpoints.py  # Automated checkpoint downloader
│   └── reproduce_all_results.sh # Complete reproduction script
├── run_distributed.sh           # 3 GPU training script
├── run_slurm_3a100.slurm        # SLURM batch script
├── requirements.txt             # Python dependencies
├── README.md                    # This file
├── REPRODUCIBILITY.md           # Complete reproduction guide
└── README_USAGE_EXAMPLES.md     # Practical usage examples

Key Implementation Details

Complete ANOVA Decomposition

# All interaction orders from 0 to concept_dim
for order in range(1, concept_dim + 1):
    combinations = list(combinations(range(concept_dim), order))
    # Process all possible concept combinations

Temporal Convolution Stack

# Multi-layer dilated convolutions
dilations = [1, 2, 4, 8]  # Paper specification
for i, dilation in enumerate(dilations):
    conv_layer = TemporalConvBlock(
        hidden_dim, hidden_dim, 
        kernel_size=3, dilation=dilation
    )

Distributed Training Setup

# 3 A100 GPU configuration
WORLD_SIZE = 3
MASTER_PORT = 12355
backend = 'nccl'
# Effective batch size: 32 × 3 = 96

Troubleshooting

Common Issues

1. Checkpoint not found: Download with python scripts/download_checkpoints.py
2. CUDA Out of Memory: Reduce batch size or use smaller model
3. Dead Concepts: Adjust JumpReLU threshold or TopK gamma schedule
4. Training Instability: Lower learning rate or increase gradient clipping

Performance Optimization

• Use mixed precision training: torch.cuda.amp
• Enable compilation: torch.compile(model)
• Optimize data loading: Increase num_workers

Paper Parameters

Based on the TimeSAE paper, the following parameters were used:

Model Architecture

• Concept Dimension (d): r × (D × T) where r = 2.0 (expansion ratio)
• Temporal Hidden Dim: 128
• Kernel Size: 3 for all temporal convolutions
• Dilation Rates: [1, 2, 4, 8] for 4-layer temporal encoder
• Dropout: 0.1

Loss Function Weights

• Sparsity Coefficient (η): 0.1
• Consistency Weight (α): 0.5 (fixed across all experiments)
• Contrastive Weight (λ): 0.1
• Temperature (τ): 0.1
• Label Fidelity Weight: 1.0

Training Parameters

• Learning Rate: 1e-3
• Weight Decay: 1e-4
• Batch Size: 32 per GPU (effective batch size: 96 on 3 GPUs)
• Optimizer: AdamW with β = (0.9, 0.999)
• Scheduler: CosineAnnealingLR with T_max = 100000, eta_min = 1e-6
• Gradient Clipping: max_norm = 1.0

Sparse Activation Parameters

• JumpReLU Threshold (φ): Initialized with U(0.01, 0.11)
• TopK Gamma: Start at 2.0, decay to 1.0 during training
• TopK Sparsity: k = concept_dim // 10 (10% activation)

Usage Examples

Basic Usage

import torch
from models import TimeSAE
from trainer import TimeSAETrainer, setup_model_and_optimizer
from utils import DatasetProcessor

# Create model with paper parameters
input_shape = (10, 128)  # (D=10 features, T=128 timesteps)
model, loss_fn, optimizer, scheduler = setup_model_and_optimizer(
    input_shape=input_shape,
    concept_dim=int(2.0 * 10 * 128),  # r=2.0 expansion ratio
    learning_rate=1e-3,
    weight_decay=1e-4
)

# Generate synthetic data
processor = DatasetProcessor()
data = processor.create_synthetic_data(
    n_samples=1000,
    n_features=10,
    n_timesteps=128
)

# Train model
trainer = TimeSAETrainer(model, loss_fn, optimizer, scheduler)
# trainer.train(train_dataloader, val_dataloader, num_epochs=100)

Explanation Generation

# Load trained model
model = TimeSAE(input_shape=(10, 128))
model.load_state_dict(torch.load('checkpoints/best_model.pt'))

# Generate explanations
x = torch.randn(5, 10, 128)  # Sample time series
concepts, x_recon, contributions = trainer.explain(x)

# Analyze concepts
from utils import ConceptInterpreter
interpreter = ConceptInterpreter()
analysis = interpreter.analyze_concept_activations(concepts)
interpreter.assign_concept_names(analysis)
summary = interpreter.get_concept_summary()

Visualization

from utils import TimeSeriesVisualizer

visualizer = TimeSeriesVisualizer()

# Plot explanations
visualizer.plot_time_series_with_concepts(
    x, x_recon, concepts, contributions,
    sample_idx=0, top_k_concepts=10,
    save_path='explanation.png'
)

# Plot concept heatmap
visualizer.plot_concept_heatmap(
    concepts, save_path='concepts.png'
)

Running the Code

Single GPU Training

python main.py \
    --n_features 10 \
    --n_timesteps 128 \
    --batch_size 32 \
    --learning_rate 1e-3 \
    --num_epochs 100 \
    --synthetic_data \
    --use_black_box \
    --generate_explanations

Multi-GPU Training (3 A100s)

chmod +x run_distributed.sh
./run_distributed.sh

SLURM Distributed Training

Use the provided SLURM script for cluster environments:

sbatch run_slurm_3a100.slurm

Dataset Configuration

Synthetic Datasets (Paper)

• FreqShapes: Frequency-based synthetic patterns
• SeqComb-UV: Sequential combinations (univariate)

Real-world Datasets (Paper)

• ECG: Arrhythmia detection (classification)
• PAM: Human activity recognition (classification)
• ETTH-1/ETTH-2: Energy demand prediction (regression)
• EliteLJ: Sports performance analysis (regression)

Custom Dataset Integration

from torch.utils.data import Dataset, DataLoader

class CustomTimeSeriesDataset(Dataset):
    def __init__(self, data, labels=None):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.labels = labels
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        if self.labels is not None:
            return self.data[idx], self.labels[idx]
        return self.data[idx]

# Use with TimeSAE
dataset = CustomTimeSeriesDataset(your_data, your_labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

Evaluation Metrics

The implementation includes all metrics from the paper:

Explanation Quality

• AUPRC: Area Under Precision-Recall Curve
• AUP: Area Under Precision
• AUR: Area Under Recall

Faithfulness

• Fx Score: ||f(x) - f(ex-)||² where concepts are removed
• Counterfactual Consistency: Measures causal effect ordering

Model Quality

• Reconstruction MSE: Mean squared error between original and reconstructed
• Sparsity Metrics: L0/L1 norms, activation ratios
• Concept Diversity: Mutual activation rates, entropy measures

Paper Results Reproduction

Expected Performance (AUPRC scores)

Faithfulness Scores (Fx metric)

File Structure

TimeSAE/
├── models.py              # Core TimeSAE architecture
├── loss.py                # Loss functions and counterfactual generation
├── trainer.py             # Distributed training logic
├── utils.py               # Evaluation, visualization, and utilities
├── main.py                # Main training script
├── run_distributed.sh     # Shell script for 3 GPU training
├── run_slurm_3a100.slurm  # SLURM batch script
├── requirements.txt       # Python dependencies
└── README.md             # This file

Key Implementation Details

Complete ANOVA Decomposition

# All interaction orders from 0 to concept_dim
for order in range(1, concept_dim + 1):
    combinations = list(combinations(range(concept_dim), order))
    # Process all possible concept combinations

Temporal Convolution Stack

# Multi-layer dilated convolutions
dilations = [1, 2, 4, 8]  # Paper specification
for i, dilation in enumerate(dilations):
    conv_layer = TemporalConvBlock(
        hidden_dim, hidden_dim, 
        kernel_size=3, dilation=dilation
    )

Distributed Training Setup

# 3 A100 GPU configuration
WORLD_SIZE = 3
MASTER_PORT = 12355
backend = 'nccl'
# Effective batch size: 32 × 3 = 96

Troubleshooting

Common Issues

1. CUDA Out of Memory: Reduce batch size or concept dimension
2. Dead Concepts: Adjust JumpReLU threshold or TopK gamma schedule
3. Training Instability: Lower learning rate or increase gradient clipping

Performance Optimization

• Use mixed precision training: torch.cuda.amp
• Enable compilation: torch.compile(model)
• Optimize data loading: Increase num_workers

TimeSAE Pretrained Checkpoints

This directory contains pretrained TimeSAE models for explaining different black-box time series models. These checkpoints are provided for users who do not have the computational resources to train TimeSAE from scratch.

Directory checkpoints

checkpoints/
├── README.md                           # This file
├── timesae_transformer/               # TimeSAE explaining Transformers
│   ├── freqshapes_best.pt             # FreqShapes dataset
│   ├── seqcomb_best.pt                # SeqComb-UV dataset  
│   ├── ecg_best.pt                    # ECG arrhythmia dataset
│   ├── pam_best.pt                    # PAM activity dataset
│   ├── etth1_best.pt                  # ETTH-1 energy dataset
│   └── etth2_best.pt                  # ETTH-2 energy dataset
├── timesae_patchts/                   # TimeSAE explaining PatchTS
│   ├── freqshapes_best.pt
│   ├── seqcomb_best.pt
│   ├── ecg_best.pt
│   ├── pam_best.pt
│   ├── etth1_best.pt
│   └── etth2_best.pt
├── timesae_chronos/                   # TimeSAE explaining Chronos
│   ├── freqshapes_best.pt
│   ├── seqcomb_best.pt
│   ├── ecg_best.pt
│   ├── pam_best.pt
│   ├── etth1_best.pt
│   └── etth2_best.pt
├── timesae_timegpt/                   # TimeSAE explaining TimeGPT
│   ├── freqshapes_best.pt
│   ├── seqcomb_best.pt
│   ├── ecg_best.pt
│   ├── pam_best.pt
│   ├── etth1_best.pt
│   └── etth2_best.pt
├── timesae_timefm/                    # TimeSAE explaining TimeFM
│   ├── freqshapes_best.pt
│   ├── seqcomb_best.pt
│   ├── ecg_best.pt
│   ├── pam_best.pt
│   ├── etth1_best.pt
│   └── etth2_best.pt
├── timesae_dlinear/                   # TimeSAE explaining DLinear
│   ├── freqshapes_best.pt
│   ├── seqcomb_best.pt
│   ├── ecg_best.pt
│   ├── pam_best.pt
│   ├── etth1_best.pt
│   └── etth2_best.pt
└── blackbox_models/                   # Trained black-box models
    ├── transformer/
    │   ├── freqshapes_best.pt
    │   ├── ecg_best.pt
    │   └── pam_best.pt
    ├── patchts/
    │   ├── freqshapes_best.pt
    │   ├── ecg_best.pt
    │   └── pam_best.pt
    └── dlinear/
        ├── freqshapes_best.pt
        ├── ecg_best.pt
        └── pam_best.pt

Download Links

Anonymous Download Link

All pretrained checkpoints are available at: https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/

Individual Model Downloads

TimeSAE for Transformer Explanations

# Download all Transformer TimeSAE models
wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_transformer/

# Or download specific datasets
wget https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_transformer/ecg_best.pt
wget https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_transformer/freqshapes_best.pt

TimeSAE for PatchTS Explanations

wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_patchts/

TimeSAE for Chronos Explanations

wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_chronos/

TimeSAE for TimeGPT Explanations

wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_timegpt/

TimeSAE for TimeFM Explanations

wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_timefm/

TimeSAE for DLinear Explanations

wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/timesae_dlinear/

Trained Black-Box Models

wget -r -np -nH --cut-dirs=2 https://anonymous.4open.science/r/TimeSAE-Checkpoints-571D/blackbox_models/

Model Specifications

TimeSAE Architecture Parameters

Dataset Specifications

🚀 Quick Usage

Loading Pretrained TimeSAE Models

import torch
from models import TimeSAE
from blackbox_models import get_model

# Load TimeSAE explaining Transformer on ECG data
def load_timesae_checkpoint(model_type: str, dataset: str, device: str = 'cuda'):
    # Determine input shape based on dataset
    dataset_shapes = {
        'freqshapes': (5, 100),
        'seqcomb': (1, 150), 
        'ecg': (1, 187),
        'pam': (17, 100),
        'etth1': (7, 96),
        'etth2': (7, 96)
    }
    
    input_shape = dataset_shapes[dataset]
    concept_dim = int(2.0 * input_shape[0] * input_shape[1])
    
    # Create TimeSAE model
    timesae = TimeSAE(
        input_shape=input_shape,
        concept_dim=concept_dim,
        activation_type='jumprelu',
        use_temporal_conv=True
    )
    
    # Load checkpoint
    checkpoint_path = f'checkpoints/timesae_{model_type}/{dataset}_best.pt'
    checkpoint = torch.load(checkpoint_path, map_location=device)
    timesae.load_state_dict(checkpoint['model_state_dict'])
    
    timesae.to(device)
    timesae.eval()
    
    return timesae

# Example usage
timesae_transformer_ecg = load_timesae_checkpoint('transformer', 'ecg')
timesae_chronos_freqshapes = load_timesae_checkpoint('chronos', 'freqshapes')

Loading Black-Box Models

def load_blackbox_checkpoint(model_type: str, dataset: str, device: str = 'cuda'):
    dataset_configs = {
        'freqshapes': {'seq_len': 100, 'n_vars': 5, 'num_classes': 5},
        'ecg': {'seq_len': 187, 'n_vars': 1, 'num_classes': 5},
        'pam': {'seq_len': 100, 'n_vars': 17, 'num_classes': 12},
    }
    
    config = dataset_configs[dataset]
    
    # Create black-box model
    if model_type == 'transformer':
        from blackbox_models import create_transformer_model
        model = create_transformer_model(**config, task_type='classification')
    elif model_type == 'patchts':
        from blackbox_models import create_patchts_model
        model = create_patchts_model(**config, task_type='classification')
    elif model_type == 'dlinear':
        from blackbox_models import create_dlinear_model
        model = create_dlinear_model(**config, task_type='classification')
    
    # Load checkpoint
    checkpoint_path = f'checkpoints/blackbox_models/{model_type}/{dataset}_best.pt'
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    
    model.to(device)
    model.eval()
    
    return model

# Example usage
transformer_ecg = load_blackbox_checkpoint('transformer', 'ecg')
patchts_freqshapes = load_blackbox_checkpoint('patchts', 'freqshapes')

Complete Explanation Pipeline

import torch
from utils import TimeSeriesVisualizer, ConceptInterpreter

# Load models
timesae = load_timesae_checkpoint('transformer', 'ecg')
black_box = load_blackbox_checkpoint('transformer', 'ecg')

# Generate sample ECG data (or load your own)
ecg_sample = torch.randn(5, 1, 187)  # 5 samples

# Generate explanations
with torch.no_grad():
    concepts, x_recon, contributions = timesae(ecg_sample)
    predictions = black_box(ecg_sample)

print(f"Concept activations shape: {concepts.shape}")
print(f"Reconstruction MSE: {torch.mse_loss(x_recon, ecg_sample):.6f}")

# Interpret concepts
interpreter = ConceptInterpreter()
concept_analysis = interpreter.analyze_concept_activations(concepts)
interpreter.assign_concept_names(concept_analysis)

# Visualize explanations
visualizer = TimeSeriesVisualizer()
visualizer.plot_time_series_with_concepts(
    ecg_sample, x_recon, concepts, contributions,
    sample_idx=0, save_path='ecg_explanation.png'
)

print("Explanation generated and saved as 'ecg_explanation.png'")

Checkpoint Details

Training Configuration

All models were trained with exact paper parameters:

TRAINING_CONFIG = {
    'eta': 0.1,        # Sparsity coefficient
    'alpha': 0.5,      # Consistency weight
    'lam': 0.1,        # Contrastive weight
    'temperature': 0.1, # InfoNCE temperature
    'learning_rate': 1e-3,
    'batch_size': 32,  # Per GPU (96 effective on 3 A100s)
    'num_epochs': 100,
    'weight_decay': 1e-4,
    'gradient_clip': 1.0
}

Performance Metrics

Each checkpoint includes validation metrics:

Checkpoint Format

Each checkpoint file contains:

checkpoint = {
    'epoch': int,                    # Training epoch
    'step': int,                     # Training step
    'model_state_dict': dict,        # TimeSAE model weights
    'optimizer_state_dict': dict,    # Optimizer state
    'scheduler_state_dict': dict,    # LR scheduler state
    'metrics': {                     # Validation metrics
        'auprc': float,
        'aup': float, 
        'aur': float,
        'faithfulness': float,
        'val_loss': float
    },
    'config': {                      # Model configuration
        'input_shape': tuple,
        'concept_dim': int,
        'activation_type': str,
        'eta': float,
        'alpha': float,
        'lam': float
    },
    'dataset_info': {               # Dataset information
        'name': str,
        'num_samples': int,
        'num_features': int,
        'sequence_length': int
    }
}

If you encounter issues with the pretrained checkpoints:

1. Check file integrity: Use the verification script above
2. Ensure compatibility: Make sure you're using the same TimeSAE version
3. Resource requirements: Some models require significant GPU memory
4. Alternative download: Try downloading individual files if bulk download fails

Citation

If you use our work, please cite:

@article{timesae2025,
  title={TimeSAE: Sparse Decoding for Faithful Explanations of Black-Box Time Series Models},
  author={Anonymous Authors},
  year={2025}
}

---

Note: These checkpoints are provided for research purposes. For production use, we recommend training TimeSAE on your specific datasets and black-box models.

License

This project is licensed under the MIT License - see the LICENSE file for details. ANOVA Decomposition**: Unlike standard approaches, implements complete functional ANOVA from order 0 to concept_dim

• Temporal Convolutions: Captures temporal dynamics with dilated convolutions (dilations: 1, 2, 4, 8)
• Distributed Training: Scales efficiently across multiple A100 GPUs
• Counterfactual Learning: Incorporates causal reasoning through contrastive learning