Model Card for PSA-Net (Patch-Spectral Attention Network)

Model Summary

• Model Name: PSA-Net (Patch-Spectral Attention Network)
• Model Type: Configurable, Multi-Branch Deep Transformer for Multivariate Time-Series Forecasting
• Framework: PyTorch
• License: MIT
• Primary Task: Multi-step ahead, multi-quantile multivariate time-series forecasting

Status: pre-release / under active evaluation. Initial evaluation on a distribution-shifted test set surfaced a calibration issue (see Evaluation Results below) that is currently being diagnosed and fixed. Numbers below are the actual current results, not aspirational ones — check back or see the linked repo for updates before relying on this model's uncertainty estimates in particular.

Architectural Overview

PSA-Net is a dataset-agnostic architecture designed for time-series forecasting on metrics exhibiting high seasonality and sudden trend spikes. It integrates four core components:

1. Patch Tokenization:

   • Divides raw input sequences into overlapping time patches (e.g., patch_len=15, patch_stride=15).
   • Reduces token length and captures local temporal patterns.

2. Dual Time / Spectral (FFT) Branch:

   • Encodes each patch both in the raw time domain and in the frequency domain via Real Fast Fourier Transform (RFFT).
   • Fuses time and frequency features to give periodic seasonality dedicated representation paths.

3. Explicit Seasonal Embeddings:

   • High-resolution Time-of-Day and Day-of-Week embedding tables added per patch timestamp position.

4. Spike Pattern Cross-Attention Bank:

   • Learned memory bank of reference "spike shape" embeddings.
   • Patch tokens cross-attend to this bank to identify early signatures of sudden workload surges.

5. Multi-Quantile Output Head:

   • Joint linear projection layer predicting uncertainty quantiles ($q_{0.1}, q_{0.5}, q_{0.9}$) for all features across future timesteps in a single forward pass.

Input History [Batch, input_window, N_features]
       │
       ├──► Patch Tokenization (patch_len, patch_stride)
       ├──► FFT Spectral Branch (n_freq_bins)
       ├──► Time-of-Day / Day-of-Week Embeddings
       └──► Spike Pattern Cross-Attention Bank (n_spike_patterns)
       │
 Transformer Blocks (n_layers, d_model, n_heads)
       │
 Quantile Output Head
       │
       ▼
Forecast Tensor [Batch, forecast_horizon, N_features, n_quantiles]

Technical Specifications & Config Parameters

Evaluation Results

Training run: 2,628,000 rows, 13 columns, 10 epochs, batch size 128, 6,265,711 parameters. Training and validation loss decreased smoothly and consistently across all 10 epochs (val loss: 0.01315 → 0.01107), showing no signs of divergence or overfitting during training itself.

Test set: synthetic_hpa_traffic_shifted_test.csv — a distribution-shifted test set (evaluating generalization under shift, not just an ordinary held-out chronological split; results below should be interpreted with that in mind).

Known issue — quantile calibration is currently broken. The 80% interval should contain the true value ~80% of the time; observed coverage ranges 1-42%, far below target across every feature. This specific pattern (healthy training loss alongside badly miscalibrated eval coverage) is consistent with a denormalization or quantile-ordering bug in the evaluation pipeline rather than necessarily a model-quality problem — this is under active investigation. Do not rely on this model's quantile/uncertainty outputs (q_0.1, q_0.9) until this is resolved and re-verified. Point forecasts (q_0.5 / MAE / RMSE) may also be affected by the same root cause and should be treated with the same caution until confirmed.

RMSE substantially exceeding MAE on requests_per_second and concurrent_users specifically (RMSE ~2-3.5x MAE) suggests a subset of large errors — plausibly missed or mistimed spike predictions — pulling the average up, rather than uniformly-distributed small errors across all predictions.

This section will be updated once the calibration issue is diagnosed and a corrected evaluation run is available. Comparison against PatchTST and Prophet baselines (trained on the same data split) is planned but not yet included below.

Direct Python Usage Examples

1. Model Initialization

import torch
from psanet.model import PSANet, PSANetConfig

# 1. Define Model Configuration
cfg = PSANetConfig(
    n_features=6,           # Works for any N_features
    input_window=120,       # 2-hour lookback @ 1-min resolution
    forecast_horizon=15,    # 15-minute future forecast
    patch_len=15,
    patch_stride=15,
    d_model=128,
    n_heads=4,
    n_layers=3,
    n_quantiles=3           # [q_0.1, q_0.5, q_0.9]
)

# 2. Instantiate PyTorch Model
device = "cuda" if torch.cuda.is_available() else "cpu"
model = PSANet(cfg).to(device)

print(f"PSA-Net Parameter Count: {model.param_count():,}")

2. Training Loop & Pinball (Quantile) Loss

import torch
from psanet.losses import quantile_loss

# Dummy input tensors: [Batch, input_window, N_features]
batch_size = 32
x_hist = torch.randn(batch_size, cfg.input_window, cfg.n_features).to(device)
y_target = torch.randn(batch_size, cfg.forecast_horizon, cfg.n_features).to(device)

# Time indices (optional for seasonal embeddings)
n_patches = (cfg.input_window - cfg.patch_len) // cfg.patch_stride + 1
tod_idx = torch.randint(0, cfg.steps_per_day, (batch_size, n_patches)).to(device)
dow_idx = torch.randint(0, cfg.days_per_week, (batch_size, n_patches)).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

# Training Step
model.train()
optimizer.zero_grad()

# Forward Pass -> Output shape: [batch_size, forecast_horizon, n_features, n_quantiles]
preds = model(x_hist, tod_idx, dow_idx)

# Compute Quantile Pinball Loss for q=[0.1, 0.5, 0.9]
loss = quantile_loss(preds, y_target, quantiles=[0.1, 0.5, 0.9])
loss.backward()
optimizer.step()

print(f"Training Step Loss: {loss.item():.4f}")

3. Inference & Quantile Unscaling

import torch
import numpy as np

# Set model to evaluation mode
model.eval()

# Assume mean and std are normalization statistics fit on training data
mean = np.zeros(cfg.n_features)
std = np.ones(cfg.n_features)

# Historical context window: [1, input_window, n_features]
context_raw = np.random.randn(cfg.input_window, cfg.n_features)
context_norm = (context_raw - mean) / std
context_tensor = torch.tensor(context_norm, dtype=torch.float32).unsqueeze(0).to(device)

with torch.no_grad():
    raw_preds = model(context_tensor)  # [1, forecast_horizon, n_features, 3]

# Unscale back to original metric units
preds_np = raw_preds.squeeze(0).cpu().numpy()  # [forecast_horizon, n_features, 3]
preds_unscaled = preds_np * std[None, :, None] + mean[None, :, None]

# Extract Quantile Curves
q_10 = preds_unscaled[..., 0]  # Lower Bound (q=0.1)
q_50 = preds_unscaled[..., 1]  # Median Point Forecast (q=0.5)
q_90 = preds_unscaled[..., 2]  # Upper Bound (q=0.9)

4. Saving and Loading Checkpoints

import torch

# Save Checkpoint
checkpoint_dict = {
    "model_state": model.state_dict(),
    "config": cfg,
    "mean": mean,
    "std": std
}
torch.save(checkpoint_dict, "psanet_checkpoint.pt")

# Load Checkpoint
loaded_ckpt = torch.load("psanet_checkpoint.pt", map_location="cpu", weights_only=False)
loaded_cfg = loaded_ckpt["config"]

loaded_model = PSANet(loaded_cfg)
loaded_model.load_state_dict(loaded_ckpt["model_state"])
loaded_model.eval()

Limitations

1. Quantile calibration currently unverified / likely broken — see Evaluation Results. Treat as the primary known limitation until resolved.
2. Minimum History Requirement: Requires at least input_window continuous historical timesteps before generating forecasts.
3. Channel-Mixing Parameter Scaling: Because the output head flattens all features and patches into a single linear projection, parameter count scales with $N_{\text{features}} \times N_{\text{patches}} \times \text{d_model}$. For setups with $>20$ features, consider using a channel-independent projection head.
4. No baseline comparison yet: results have not yet been compared against the PatchTST or Prophet baselines trained on the same data; until that comparison exists, it isn't yet established that PSA-Net's added complexity (spectral branch, spike bank) outperforms a simpler model on this data.