[Bounty $1500] Time Series Transformer Model Bring-Up Using TTNN APIs

Bounties
1

Tenstorrent Bounties

Issue ID: I_kwDOI9Wqc87WztU8

:memo: Background

This bounty is for bringing up the Time Series Transformer model using TTNN APIs on Tenstorrent hardware (Wormhole or Blackhole).

Time Series Transformer is a vanilla encoder-decoder Transformer architecture specifically designed for probabilistic time-series forecasting. Released in 2022 and integrated into HuggingFace Transformers, it represents a clean, straightforward application of the transformer architecture to time series data without specialized attention mechanisms or architectural modifications.

Key Capabilities:

  • Vanilla Transformer Architecture: Standard encoder-decoder transformer without modifications
    • Encoder processes historical context (past values)
    • Decoder autoregressively generates future predictions
    • Clean, well-understood architecture
  • Probabilistic Forecasting: Learns distributions rather than point estimates
    • Outputs distribution parameters (e.g., mean and scale for Student’s t-distribution)
    • Enables uncertainty quantification
    • Supports multiple distribution families (student_t, normal, negative_binomial)
  • Rich Feature Support: Comprehensive input feature handling
    • Temporal features (time encodings: day, month, hour, etc.)
    • Static categorical features (e.g., store ID, region)
    • Static real features (e.g., product embeddings)
    • Lag features for capturing historical patterns
  • Teacher-Forcing Training: Similar to sequence-to-sequence models
    • Efficient training paradigm
    • Prevents error accumulation during training
  • Autoregressive Generation: Flexible inference
    • Generate multiple samples from learned distribution
    • Confidence intervals and prediction intervals
    • Scenario analysis and risk assessment
  • Multivariate Support: Handles multiple correlated time series
  • Flexible Input/Output: Configurable context and prediction lengths

:bullseye: What Success Looks Like

A successful submission will fulfill all requirements in the following stages. Payout is made after all three stages are completed.

Stage 1 — Bring-Up

  • Implement Time Series Transformer model using TTNN APIs (Python)
  • Implements the full encoder-decoder architecture:
    • Value embedding layer with optional lag features
    • Temporal feature embeddings (past and future time features)
    • Static feature embeddings (categorical and real)
    • Standard transformer encoder with self-attention
    • Standard transformer decoder with:
      • Masked self-attention (causal masking)
      • Cross-attention to encoder outputs
    • Distribution head for probabilistic outputs
  • Model runs on Tenstorrent hardware (Wormhole or Blackhole) with no errors
  • Supports multiple distribution outputs:
    • Student’s t-distribution (default, handles outliers well)
    • Normal/Gaussian distribution (standard choice)
    • Negative binomial distribution (for count data)
  • Supports comprehensive feature inputs:
    • Past values (historical time series)
    • Future time features (known future temporal information)
    • Past time features (historical temporal information)
    • Static categorical features (unchanging categorical variables)
    • Static real features (unchanging continuous variables)
    • Observed mask (handling missing values)
  • Produces valid probabilistic predictions on standard benchmarks (Tourism or ETT datasets)
  • Output is verifiable (prediction accuracy, compare with PyTorch/HuggingFace reference)
  • Achieves baseline performance targets:
    • Inference throughput: At least 100 sequences/second for standard context
    • Latency: < 50ms for single sequence prediction (batch size 1)
    • Sample generation: 100 samples in < 1 second
  • Accuracy evaluation:
    • Negative log-likelihood (NLL) within 5% of PyTorch reference
    • CRPS (Continuous Ranked Probability Score) within 5% of reference
    • Mean prediction within 5% MAE of reference
  • Clear instructions for setup and running the model

Stage 2 — Basic Optimizations

  • Use optimal sharded/interleaved memory configs for:
    • Value and feature embedding layers
    • Encoder self-attention layers (Q, K, V projections)
    • Decoder self-attention and cross-attention layers
    • Feed-forward network layers
    • Distribution head parameters
  • Implement efficient sharding strategy for:
    • Multi-head attention computation (encoder and decoder)
    • Cross-attention between encoder and decoder
    • Autoregressive decoding steps
    • Multiple sample generation (parallel sampling)
  • Fuse simple ops where possible:
    • Embedding layers (value + temporal + static features)
    • Attention components (Q, K, V computation)
    • FFN layers (Linear + Activation + Dropout)
    • Layer normalization operations
    • Distribution parameter computation
  • Store intermediate activations in L1 where beneficial
  • Use recommended TTNN/tt-metal transformer flows
  • Leverage TT library of fused ops for:
    • Multi-head attention blocks (encoder and decoder)
    • Cross-attention blocks
    • Feed-forward network (FFN) blocks
    • Layer normalization
  • Optimize encoder-decoder specific operations:
    • Efficient encoder hidden state storage for cross-attention
    • Causal masking for decoder self-attention
    • Cross-attention computation
    • KV-cache for autoregressive generation
  • Optimize probabilistic output generation:
    • Efficient distribution parameter computation
    • Parallel sample generation from distributions
    • Sampling operations (e.g., Student’s t, Normal)

Stage 3 — Deeper Optimization

  • Maximize core counts used per inference
  • Implement deeper TT-specific optimizations:
    • Flash Attention or equivalent for encoder/decoder self-attention
    • Efficient cross-attention implementation
    • Optimized KV-cache management for autoregressive decoding
    • Parallel multi-head attention computation
    • Efficient causal masking implementation
  • Minimize prediction latency for real-time forecasting
  • Batch processing for multiple time series
  • Optimize autoregressive generation:
    • Efficient KV-cache updates
    • Minimize memory copies during decoding
    • Pipeline decoder steps
    • Speculative decoding (if applicable)
  • Optimize sample generation:
    • Parallel generation of multiple samples
    • Efficient sampling from distributions
    • Vectorized distribution operations
    • Minimize overhead for 100+ samples
  • Pipeline encoder and decoder operations:
    • Overlap encoder computation with decoder initialization
    • Pipeline decoder steps
    • Overlap distribution computation with sampling
  • Minimize memory and TM (tensor manipulation) overheads
  • Support for streaming inference (online forecasting)
  • Explore techniques for longer sequences (1024+ context length)
  • Document any advanced tuning, known limitations, or trade-offs
  • Target stretched goals:
    • 500+ sequences/second throughput for batch inference
    • < 20ms latency for single sequence prediction
    • 1000+ samples generation in < 2 seconds
    • Support for context lengths up to 2048
    • Efficient handling of 100+ time series in batch
  • Multi-distribution support (switch between distributions efficiently)

:compass: Guidance & Starting Points

Primary Resources

  • Use the TTNN model bring-up tech report as your primary reference
  • Reference Transformer implementations in tt-metal for encoder-decoder patterns
  • Use the HuggingFace Transformers Time Series Transformer (documentation) as the reference implementation
  • Refer to TT Fused ops PR #29236 for optimization opportunities
  • Check existing sequence-to-sequence transformer implementations in tt-metal

HuggingFace Implementation Reference

The HuggingFace implementation provides two main classes:

  1. TimeSeriesTransformerModel: The bare transformer encoder-decoder outputting raw hidden states
  2. TimeSeriesTransformerForPrediction: Full model with distribution head for probabilistic forecasting

Key features to implement:

Architecture:

  • Encoder: Processes past_values with context_length
  • Decoder: Autoregressively generates prediction_length forecasts
  • Teacher-forcing training: Decoder sees shifted future values during training
  • Autoregressive inference: Decoder generates one step at a time, using previous predictions

Probabilistic Forecasting:

  • Distribution outputs:
    • student_t: Student’s t-distribution (handles outliers, default)
    • normal: Gaussian distribution
    • negative_binomial: For count data (e.g., sales)
  • Distribution head: Outputs parameters (e.g., mean, scale for Student’s t)
  • Loss function: Negative log-likelihood (NLL)
  • Sampling: Generate multiple samples from learned distribution

Feature Inputs:

  • past_values: Historical time series [batch, context_length, input_size]
  • past_time_features: Temporal encodings for past [batch, context_length, num_time_features]
    • Examples: hour of day, day of week, day of month, month of year
  • future_time_features: Known temporal encodings for future [batch, prediction_length, num_time_features]
  • static_categorical_features: Static category IDs [batch, num_static_categorical]
    • Examples: store ID, product ID, region ID
  • static_real_features: Static continuous values [batch, num_static_real]
    • Examples: store size, product price, latitude/longitude
  • past_observed_mask: Mask for missing values [batch, context_length, input_size]

Additional Features:

  • lags_sequence: Lag indices for historical values (e.g., [1, 2, 3, 7, 14, 28])
  • scaling: Input normalization (“mean”, “std”, or None)
  • num_parallel_samples: Number of samples to generate (default 100)

:magnifying_glass_tilted_right: Possible Approaches

Sequential Implementation Strategy

  1. Start from HuggingFace implementation and port components sequentially:

    • Begin with input embedding layer (value + features)
    • Implement standard transformer encoder
    • Implement transformer decoder with causal masking
    • Add cross-attention mechanism
    • Implement distribution head
    • Test training with teacher-forcing
    • Test autoregressive inference
    • Optimize KV-cache for generation

  2. Validate each component against PyTorch reference before integration:

    • Test embedding layer outputs with various features
    • Validate encoder outputs
    • Validate decoder self-attention (with causal mask)
    • Validate cross-attention outputs
    • Check distribution parameter outputs
    • Validate sampling from distributions
    • Compare end-to-end predictions

  3. Leverage existing transformer implementations:

    • Use existing encoder-decoder transformers in tt-metal as template
    • Adapt for time series specific features
    • Reuse attention and FFN implementations
    • Add time series specific components

  4. Test on standard benchmarks:

    • Start with Tourism dataset (simpler, monthly data)
    • Test on ETT datasets (hourly, more complex)
    • Evaluate probabilistic metrics (CRPS, NLL)
    • Test uncertainty quantification quality
    • Compare with published results

  5. Experiment with optimizations:

    • Efficient KV-cache implementation
    • Fused encoder/decoder layers
    • Parallel sample generation
    • Optimized distribution operations
    • Batch processing strategies

  6. Use TTNN profiling tools to identify bottlenecks:

    • Measure encoder vs. decoder time
    • Profile attention computation
    • Measure autoregressive generation overhead
    • Profile sample generation
    • Identify memory bottlenecks

  7. Open a draft PR early to get feedback on your approach

Alternative Approaches

  • Modular testing:
    • Implement encoder and decoder separately
    • Test each component independently
    • Integrate gradually
  • Start simple:
    • Begin with univariate forecasting
    • Add multivariate support
    • Add feature support incrementally
  • Leverage existing code:
    • Use existing transformer implementations as starting point
    • Adapt for time series specifics
  • Progressive features:
    • Start with basic value inputs
    • Add time features
    • Add static features
    • Add lag features

:bar_chart: Result Submission Guidelines

Beyond the model implementation itself, contributors must submit the following material as proof of work. However, feel free to open a PR at any time if you want us checking that you are on the right track. Just understand that payout is only made after all 3 stages are completed.

Deliverables:

  • Functional model implementation
  • Validation logs (output correctness)
  • Performance report + header for final review

Links:

:books: Resources

Model Resources

Datasets & Benchmarks

Datasets & Benchmarks

Primary Source (Recommended):

Individual Datasets:

Benchmark Scripts:

  • TSLib provides standard evaluation scripts
  • Consistent train/val/test splits across all datasets
  • MSE, MAE metrics computed uniformly

Probabilistic Forecasting Resources

  • Evaluation Metrics:
    • CRPS (Continuous Ranked Probability Score): Measures probabilistic forecast quality
    • QuantileLoss: Evaluates specific quantiles
    • NLL (Negative Log-Likelihood): Training objective
    • Coverage: Percentage of observations within prediction intervals
  • Distribution Resources:
    • Student’s t-distribution: Heavy-tailed, handles outliers
    • Normal distribution: Standard choice
    • Negative binomial: For count data

Academic Resources

  • Transformer Architecture: Vaswani et al., “Attention is All You Need”, NeurIPS 2017
  • Probabilistic Forecasting:
    • Salinas et al., “DeepAR: Probabilistic Forecasting with Autoregressive Recurrent Networks”, IJF 2020
    • Rangapuram et al., “Deep State Space Models for Time Series Forecasting”, NeurIPS 2018
  • Related Work:
    • Temporal Fusion Transformer (TFT)
    • N-BEATS, N-HiTS
    • Autoformer, Informer

TT-Metal Resources

Helpful Tools

  • Visualization:
    • Plot predictions with confidence intervals
    • Attention weights visualization
    • Distribution visualization (histogram of samples)
  • Profiling: TTNN profiler for performance analysis
  • Testing: pytest framework for model testing
  • Evaluation: Implement CRPS, quantile loss metrics