psa-net

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.