Overview
Transformer sub-module implementations for the TimeXer model, including attention layers, embeddings, encoder components, and the output head.
Description
This module provides the core nn.Module building blocks for the TimeXer architecture. FullAttention implements scaled dot-product attention with optional causal masking and an efficient PyTorch SDPA backend. AttentionLayer wraps an attention mechanism with linear projections for queries, keys, and values. DataEmbedding_inverted embeds exogenous variables by transposing the variate and time dimensions. PositionalEmbedding provides sinusoidal positional encodings. EnEmbedding handles endogenous variable patch embedding with a learnable global token. Encoder stacks multiple EncoderLayer instances, each containing self-attention on endogenous patches and cross-attention between the global token and exogenous embeddings, followed by a convolutional feedforward network.
Usage
These sub-modules are internal components used by both TimeXer (v1) and TimeXer (v2). They can also be reused when building custom Transformer-based time series models that need patch embeddings, dual attention, or inverted data embeddings.
Code Reference
Source Location
Signature: TriangularCausalMask
class TriangularCausalMask:
def __init__(self, B, L, device="cpu"):
Signature: FullAttention
class FullAttention(nn.Module):
def __init__(
self,
mask_flag=True,
factor=5,
scale=None,
attention_dropout=0.1,
output_attention=False,
use_efficient_attention=False,
):
Signature: AttentionLayer
class AttentionLayer(nn.Module):
def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None):
Signature: DataEmbedding_inverted
class DataEmbedding_inverted(nn.Module):
def __init__(self, c_in, d_model, embed_type="fixed", freq="h", dropout=0.1):
Signature: PositionalEmbedding
class PositionalEmbedding(nn.Module):
def __init__(self, d_model, max_len=5000):
Signature: FlattenHead
class FlattenHead(nn.Module):
def __init__(self, n_vars, nf, target_window, head_dropout=0, n_quantiles=1):
Signature: EnEmbedding
class EnEmbedding(nn.Module):
def __init__(self, n_vars, d_model, patch_len, dropout):
Signature: Encoder
class Encoder(nn.Module):
def __init__(self, layers, norm_layer=None, projection=None):
Signature: EncoderLayer
class EncoderLayer(nn.Module):
def __init__(
self,
self_attention,
cross_attention,
d_model,
d_ff=None,
dropout=0.1,
activation="relu",
):
Import
from pytorch_forecasting.models.timexer.sub_modules import (
FullAttention,
AttentionLayer,
DataEmbedding_inverted,
PositionalEmbedding,
FlattenHead,
EnEmbedding,
Encoder,
EncoderLayer,
TriangularCausalMask,
)
I/O Contract
FullAttention Inputs
| Name |
Type |
Required |
Description
|
| mask_flag |
bool |
No |
Whether to apply causal masking (default True)
|
| factor |
int |
No |
Scaling factor for attention (default 5)
|
| scale |
float |
No |
Custom scaling factor; defaults to 1/sqrt(d_keys)
|
| attention_dropout |
float |
No |
Dropout rate for attention scores (default 0.1)
|
| output_attention |
bool |
No |
Whether to return attention weights (default False)
|
| use_efficient_attention |
bool |
No |
Use PyTorch native SDPA (default False)
|
EncoderLayer Forward Inputs
| Name |
Type |
Required |
Description
|
| x |
torch.Tensor |
Yes |
Endogenous patch embeddings of shape (batch*n_vars, patch_num+1, d_model)
|
| cross |
torch.Tensor |
Yes |
Exogenous variate embeddings of shape (batch, n_exog_vars, d_model)
|
| x_mask |
torch.Tensor |
No |
Mask for self-attention
|
| cross_mask |
torch.Tensor |
No |
Mask for cross-attention
|
EncoderLayer Forward Outputs
| Name |
Type |
Description
|
| output |
torch.Tensor |
Encoded representation of shape (batch*n_vars, patch_num+1, d_model)
|
FlattenHead Outputs
| Name |
Type |
Description
|
| output |
torch.Tensor |
Prediction tensor of shape (batch_size, n_vars, target_window, n_quantiles)
|
Usage Examples
import torch
from pytorch_forecasting.models.timexer.sub_modules import (
FullAttention, AttentionLayer, EncoderLayer, Encoder, EnEmbedding,
DataEmbedding_inverted, FlattenHead,
)
# Build an encoder layer with dual attention
self_attn = AttentionLayer(
FullAttention(False, factor=5, attention_dropout=0.1, output_attention=False),
d_model=256, n_heads=4,
)
cross_attn = AttentionLayer(
FullAttention(False, factor=5, attention_dropout=0.1, output_attention=False),
d_model=256, n_heads=4,
)
layer = EncoderLayer(self_attn, cross_attn, d_model=256, d_ff=1024, dropout=0.1)
# Stack into encoder
encoder = Encoder([layer], norm_layer=torch.nn.LayerNorm(256))
# Create embeddings
en_embed = EnEmbedding(n_vars=1, d_model=256, patch_len=16, dropout=0.1)
ex_embed = DataEmbedding_inverted(c_in=96, d_model=256)
Related Pages