Principle:Sktime Pytorch forecasting Scaled Dot Product Attention

Knowledge Sources	pytorch-forecasting Attention Is All You Need
Domains	Time_Series, Forecasting, Deep_Learning, Attention_Mechanisms
Last Updated	2026-02-08 09:00 GMT

Overview

Scaled dot-product attention mechanism with a multi-head projection wrapper, supporting optional causal masking and dropout regularization, used to compute context-aware representations of sequence elements.

Description

Scaled Dot-Product Attention computes the relevance between every pair of positions in a sequence by taking the dot product of query and key vectors, scaling by the inverse square root of the key dimension, and applying a softmax to obtain attention weights. These weights are then used to produce a weighted sum of value vectors.

The mechanism is wrapped in a multi-head attention layer (AttentionLayer) that projects the input into multiple independent query, key, and value subspaces (heads). Each head computes attention in parallel, and the results are concatenated and projected back to the model dimension. This allows the model to jointly attend to information from different representation subspaces at different positions.

The FullAttention module provides two computational backends: (1) an einsum-based path that explicitly computes the score matrix and supports outputting attention weights, and (2) an efficient path that delegates to PyTorch's native scaled_dot_product_attention, which automatically selects the optimal backend (FlashAttention-2, Memory-Efficient Attention, or a C++ fallback) based on hardware and input properties.

Causal masking is supported via a TriangularCausalMask that fills the upper triangle of the score matrix with negative infinity before softmax, preventing each position from attending to future positions. When no explicit mask is supplied and the mask flag is enabled, this causal mask is generated automatically.

Usage

Use FullAttention as the inner attention module of AttentionLayer for transformer-based time series forecasting models such as TimeXer. Configure mask_flag=True for autoregressive (decoder-side) attention where future leakage must be prevented, and mask_flag=False for bidirectional encoder self-attention. Enable use_efficient_attention=True for longer sequences where memory and compute savings are important, noting that attention weight inspection is unavailable in this mode.

Theoretical Basis

Scaled Dot-Product Attention:

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\mathrm{\top }}}{\sqrt{d_k}}\right)V$

Where $Q\in {\mathbb{ℝ}}^{L\times d_k}$, $K\in {\mathbb{ℝ}}^{S\times d_k}$, $V\in {\mathbb{ℝ}}^{S\times d_v}$, and the scaling factor $\sqrt{d_k}$ prevents the dot products from growing too large in magnitude, which would push the softmax into regions of extremely small gradients.

Multi-Head Attention:

$\text{MultiHead}(Q,K,V)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_H)W^O$

where each head is:

${\text{head}}_i=\text{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$

Causal Masking:

For autoregressive tasks, a binary upper-triangular mask $M$ is applied before softmax:

${\text{score}}_{i,j}=\{\begin{array}{ll}{Q}_i{K}_j^{\mathrm{\top }}& \text{if }j\le i\\ -\mathrm{\infty }& \text{if }j>i\end{array}$

This ensures that the attention output at position $i$ depends only on positions $\le i$.

Pseudo-code for the multi-head attention wrapper:

# Multi-head attention wrapper (pseudo-code)
def multi_head_attention(queries, keys, values, n_heads):
    Q = linear_q(queries).reshape(B, L, H, d_k)
    K = linear_k(keys).reshape(B, S, H, d_k)
    V = linear_v(values).reshape(B, S, H, d_v)

    output, weights = scaled_dot_product_attention(Q, K, V, mask)

    output = output.reshape(B, L, H * d_v)
    return linear_out(output), weights

Related Pages

Implemented By

• Implementation:Sktime_Pytorch_forecasting_AttentionLayer
• Implementation:Sktime_Pytorch_forecasting_FullAttention