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

Temporal Encoder: Multi-layer dilated convolutions for time series feature extraction
Sparse Encoder: Linear transformation with JumpReLU/TopK activations
Complete ANOVA Decoder: Full functional decomposition using all interaction orders
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

TimeSAE Model	Datasets	AUPRC (FreqShapes)	Download Size
Transformer	6 datasets	0.950±0.011	~270MB
PatchTS	6 datasets	0.842±0.014	~260MB
Chronos	6 datasets	0.930±0.016	~265MB
TimeGPT	6 datasets	0.958±0.016	~270MB
TimeFM	6 datasets	0.945±0.012	~275MB
DLinear	6 datasets	0.825±0.018	~250MB

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)

Dataset	Transformer	PatchTS	TimeGPT	Chronos
FreqShapes	0.950±0.011	0.842±0.014	0.958±0.016	0.930±0.016
SeqComb-UV	0.811±0.022	0.715±0.019	0.837±0.023	0.807±0.022
ECG	0.950±0.011	0.980±0.011	0.912±0.020	0.894±0.015
PAM	0.998±0.012	0.981±0.012	0.957±0.022	0.939±0.017

Faithfulness Scores (Fx metric)

Dataset	Score
ECG	1.78±0.078
PAM	2.15±0.080
ETTH-1	2.12±0.072
ETTH-2	2.09±0.069

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

Checkpoint not found: Download with python scripts/download_checkpoints.py
CUDA Out of Memory: Reduce batch size or use smaller model
Dead Concepts: Adjust JumpReLU threshold or TopK gamma schedule
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)

Dataset	Transformer	PatchTS	TimeGPT	Chronos
FreqShapes	0.950±0.011	0.842±0.014	0.958±0.016	0.930±0.016
SeqComb-UV	0.811±0.022	0.715±0.019	0.837±0.023	0.807±0.022
ECG	0.950±0.011	0.980±0.011	0.912±0.020	0.894±0.015
PAM	0.998±0.012	0.981±0.012	0.957±0.022	0.939±0.017

Faithfulness Scores (Fx metric)

Dataset	Score
ECG	1.78±0.078
PAM	2.15±0.080
ETTH-1	2.12±0.072
ETTH-2	2.09±0.069

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

CUDA Out of Memory: Reduce batch size or concept dimension
Dead Concepts: Adjust JumpReLU threshold or TopK gamma schedule
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

Model Type	Input Shape	Concept Dim	Activation	Sparsity	Paper AUPRC
Transformer	(5, 100)	1000	JumpReLU	10%	0.950±0.011
PatchTS	(5, 100)	1000	JumpReLU	10%	0.842±0.014
Chronos	(5, 100)	1000	JumpReLU	10%	0.930±0.016
TimeGPT	(5, 100)	1000	JumpReLU	10%	0.958±0.016
TimeFM	(5, 100)	1000	JumpReLU	10%	0.945±0.012
DLinear	(5, 100)	1000	JumpReLU	10%	0.825±0.018

Dataset Specifications

Dataset	Shape	Classes	Task	Domain
FreqShapes	(5, 100)	5	Classification	Synthetic
SeqComb-UV	(1, 150)	3	Classification	Synthetic
ECG	(1, 187)	5	Classification	Medical
PAM	(17, 100)	12	Classification	Activity
ETTH-1	(7, 96)	-	Regression	Energy
ETTH-2	(7, 96)	-	Regression	Energy

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

Model + Dataset	AUPRC	AUP	AUR	Faithfulness
Transformer + FreqShapes	0.950±0.011	0.854±0.017	0.745±0.011	1.89±0.071
PatchTS + FreqShapes	0.842±0.014	0.746±0.014	0.664±0.014	1.85±0.080
Chronos + FreqShapes	0.930±0.016	0.851±0.016	0.753±0.016	1.78±0.078
TimeGPT + FreqShapes	0.958±0.016	0.875±0.016	0.777±0.016	2.15±0.080
Transformer + ECG	0.950±0.011	0.854±0.017	0.745±0.011	1.89±0.071
Transformer + PAM	0.998±0.012	0.926±0.034	0.505±0.019	1.85±0.080

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:

Check file integrity: Use the verification script above
Ensure compatibility: Make sure you're using the same TimeSAE version
Resource requirements: Some models require significant GPU memory
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

Name		Name	Last commit message	Last commit date
checkpoints		checkpoints
docs		docs
explanations_annotations		explanations_annotations
notebooks/google-colab		notebooks/google-colab
scripts		scripts
timesae		timesae
Makefile		Makefile
README.md		README.md
Reproducibility_guide.md		Reproducibility_guide.md
image.png		image.png
pyproject.toml		pyproject.toml
requirement.txt		requirement.txt
steup.py		steup.py

Latest commit

History

Folders and files

Repository files navigation

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

Features

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

Key Innovations

Project Updates

Installation

🚀 Quick Start with Pretrained Models

Download Pretrained Checkpoints

Instant Explanation Generation

Available Pretrained Models

Paper Parameters

Model Architecture

Loss Function Weights

Training Parameters

Sparse Activation Parameters

Usage Examples

Basic Usage

Using Pretrained Models

Explanation Generation

Visualization

Running the Code

Single GPU Training

Multi-GPU Training (3 A100s)

SLURM Distributed Training

Black-Box Model Support

Pretrained Models (Ready to Use)

Trainable Models (Train from Scratch)

Dataset Configuration

Synthetic Datasets (Paper)

Real-world Datasets (Paper)

Custom Dataset Integration

Evaluation Metrics

Explanation Quality

Faithfulness

Model Quality

Paper Results Reproduction

Expected Performance (AUPRC scores)

Faithfulness Scores (Fx metric)

Complete Reproduction

TimeSAE lib Structure

Key Implementation Details

Complete ANOVA Decomposition

Temporal Convolution Stack

Distributed Training Setup

Troubleshooting

Common Issues

Performance Optimization

Paper Parameters

Model Architecture

Loss Function Weights

Training Parameters

Sparse Activation Parameters

Usage Examples

Basic Usage

Explanation Generation

Visualization

Running the Code

Single GPU Training

Multi-GPU Training (3 A100s)

SLURM Distributed Training

Dataset Configuration

Synthetic Datasets (Paper)

Real-world Datasets (Paper)

Custom Dataset Integration

Evaluation Metrics

Explanation Quality

Faithfulness

Model Quality

Paper Results Reproduction

Expected Performance (AUPRC scores)

Faithfulness Scores (Fx metric)

File Structure

Key Implementation Details