Implementation:Sktime Pytorch forecasting AttentionLayer

Knowledge Sources	Sktime_Pytorch_forecasting
Domains	Time_Series, Forecasting, Deep_Learning
Last Updated	2026-02-08 08:00 GMT

Overview

A multi-head attention wrapper layer that combines query, key, and value linear projections with a pluggable attention mechanism.

Description

AttentionLayer is an nn.Module that implements a standard multi-head attention pattern. It creates separate linear projections for queries, keys, and values, splits them across multiple attention heads, delegates the actual attention computation to a provided inner attention module (e.g., FullAttention), and projects the result back to the model dimension. When the key/value sequence length is zero (no exogenous variables), it short-circuits the cross-attention process and applies only the query projection and output projection.

The layer is designed to be composable: different attention mechanisms can be injected via the attention parameter, making it reusable across various transformer-based forecasting architectures.

Usage

Use this layer as a building block in transformer-based forecasting models whenever multi-head attention is needed. Pair it with FullAttention or other attention implementations for self-attention or cross-attention blocks.

Code Reference

Source Location

• Repository: Sktime_Pytorch_forecasting
• File: pytorch_forecasting/layers/_attention/_attention_layer.py
• Lines: 1-58

Signature

class AttentionLayer(nn.Module):
    def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None)
    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None)

Import

from pytorch_forecasting.layers._attention._attention_layer import AttentionLayer

I/O Contract

Inputs

Name	Type	Required	Description
attention	nn.Module	Yes	Inner attention mechanism (e.g., FullAttention) to compute attention scores
d_model	int	Yes	Dimension of the model (input and output feature size)
n_heads	int	Yes	Number of attention heads
d_keys	int or None	No	Dimension of key projections per head. Defaults to d_model // n_heads
d_values	int or None	No	Dimension of value projections per head. Defaults to d_model // n_heads
queries	torch.Tensor	Yes	Query tensor of shape (B, L, d_model)
keys	torch.Tensor	Yes	Key tensor of shape (B, S, d_model)
values	torch.Tensor	Yes	Value tensor of shape (B, S, d_model)
attn_mask	object or None	Yes	Attention mask object (e.g., TriangularCausalMask)
tau	float or None	No	Optional temperature parameter passed to inner attention
delta	float or None	No	Optional delta parameter passed to inner attention

Outputs

Name	Type	Description
output	torch.Tensor	Attended output tensor of shape (B, L, d_model)
attn	torch.Tensor or None	Attention weights if the inner attention returns them, otherwise None

Usage Examples

import torch
from pytorch_forecasting.layers._attention._attention_layer import AttentionLayer
from pytorch_forecasting.layers._attention._full_attention import FullAttention

# Create a multi-head attention layer with 8 heads
attention = AttentionLayer(
    attention=FullAttention(mask_flag=False, attention_dropout=0.1),
    d_model=512,
    n_heads=8,
)

# Forward pass
B, L, S = 32, 96, 96
queries = torch.randn(B, L, 512)
keys = torch.randn(B, S, 512)
values = torch.randn(B, S, 512)

output, attn_weights = attention(queries, keys, values, attn_mask=None)
# output shape: (32, 96, 512)

Related Pages

• Principle:Sktime_Pytorch_forecasting_Scaled_Dot_Product_Attention