Transformer Models#

Attention-based architectures for time series representation learning.

TST#

Transformer-based model with masked reconstruction pretraining. Uses BasicEncodingMixin for inference.

class chronocratic.models.transformer.tst.model.TST(*, input_dim: int, sequence_length: int, hidden_dim: int = 64, num_heads: int = 8, depth: int = 3, feedforward_dim: int = 256, dropout_rate: float = 0.1, pos_encoding: str = 'fixed', activation: str = 'gelu', normalization_layer_type: NormalizationLayerType = NormalizationLayerType.BATCH, freeze: bool = False, learning_rate: float = 0.001, lr_step: tuple[int, ...] | None = None, lr_factor: float = 0.1, weight_decay: float = 0.0, global_reg: bool = False, sync_dist: bool = False, augmentation: Callable | None = None)#

Bases: LightningModule, BasicEncodingMixin

PyTorch Lightning module for TST.

Representation-learning model trained with a masked-reconstruction pretraining objective. The same model supports both random-mask imputation and structured-mask transduction pretraining — the masking strategy is configured upstream in the dataloader and is transparent to the model.

Batch format expected from the DataLoader:

(X, targets, target_masks, padding_masks, IDs)

where target_masks marks the positions whose reconstruction is scored, and padding_masks marks valid (non-padded) timesteps.

forward(x, padding_masks) returns transformer representations of shape (batch, seq_len, hidden_dim), not the masked-reconstruction output. The reconstruction head is internal and used only during pretraining.

For downstream classification / regression, use SupervisedModule from chronocratic.models.supervised.

Parameters:

  • input_dim – Number of input features (channels).

  • sequence_length – Maximum sequence length supported by the positional encoding.

  • hidden_dim – Transformer model (token) dimensionality.

  • num_heads – Number of attention heads.

  • depth – Number of stacked transformer encoder layers.

  • feedforward_dim – Hidden dimensionality of the transformer feed-forward block.

  • dropout_rate – Dropout probability used throughout the transformer.

  • pos_encoding – Positional-encoding type (e.g. 'fixed' or 'learnable') passed to the encoder.

  • activation – Activation function name passed to the transformer feed-forward block.

  • normalization_layer_type – Normalization layer used inside the encoder. BATCH (default) uses custom BatchNorm transformer layers. CHANNEL uses PyTorch’s LayerNorm-based TransformerEncoderLayer.

  • freeze – When True, freezes the backbone weights and only trains the output layer.

  • learning_rate – Base learning rate for the Adam optimizer.

  • lr_step – Milestones (in epochs) for the MultiStepLR scheduler. None means no decay (defaults to a single far-future milestone internally).

  • lr_factor – Multiplicative decay factor applied at each lr_step milestone.

  • weight_decay – L2 regularization coefficient. Applied to the output layer only when global_reg=False, or to all parameters (via optimizer weight decay) when global_reg=True.

  • global_reg – Whether weight_decay is applied globally as weight decay (True) or only to the output layer (False).

  • sync_dist – Whether to synchronize logged metrics across distributed processes.

  • augmentation – Optional custom augmentation function.

This model was implemented based on the code available on this GitHub repo gzerveas/mvts_transformer under MIT License.

forward(x: Tensor, padding_masks: Tensor) → Tensor#

Return transformer representations of shape (batch, seq_len, hidden_dim).

get_representations(x: Tensor, padding_masks: Tensor) → Tensor#

Run the transformer trunk, skipping the reconstruction output layer.

reconstruct(x: Tensor, padding_masks: Tensor) → Tensor#

Run the full backbone, including the reconstruction output layer.

Used during masked-reconstruction pretraining; downstream callers should use forward / get_representations instead.

training_step(batch: tuple, _batch_idx: int) → Tensor#

Compute and log the masked-reconstruction training loss for one batch.

validation_step(batch: tuple, _batch_idx: int) → Tensor#

Compute and log the masked-reconstruction validation loss for one batch.

configure_gradient_clipping(optimizer: Optimizer, gradient_clip_val: float | None = None, gradient_clip_algorithm: str | None = None) → None#

Clip gradients by global norm to stabilise training.

configure_optimizers() → OptimizerLRSchedulerConfig#

Return Adam optimizer with MultiStepLR scheduler.

property encoder: Module#

Return the transformer encoder for inspection and checkpointing.

property sequence_length: int#

Return the maximum sequence length supported by this model.

property representation_dim: int#

Feature dim of the pooled encode() representation.

Returns:

hidden_dim — the feature dimension of the vector produced by encode() with output=VECTOR.

Configuration for the TST (Time Series Transformer) model.

Provides TSTModelParameters with all settings for the transformer backbone used during masked-reconstruction pretraining.

class chronocratic.models.transformer.tst.config.TSTModelParameters(*, input_dim: int, sequence_length: int, hidden_dim: int = 64, num_heads: int = 8, depth: int = 3, feedforward_dim: int = 256, dropout_rate: float = 0.1, pos_encoding: str = 'fixed', activation: str = 'gelu', normalization_layer_type: NormalizationLayerType = NormalizationLayerType.BATCH, freeze: bool = False, learning_rate: float = 0.001, lr_step: tuple[int, ...] | None = None, lr_factor: float = 0.1, weight_decay: float = 0.0, global_reg: bool = False, sync_dist: bool = False)#

Bases: object

Configuration for the TST model.

Parameters:

  • input_dim – Number of input features (channels) in the time series.

  • sequence_length – Maximum sequence length supported by the positional encoding.

  • hidden_dim – Transformer model (token) dimensionality.

  • num_heads – Number of attention heads.

  • depth – Number of stacked transformer encoder layers.

  • feedforward_dim – Hidden dimensionality of the transformer feed-forward block.

  • dropout_rate – Dropout probability used throughout the transformer.

  • pos_encoding – Positional-encoding type (e.g. 'fixed' or 'learnable') passed to the encoder.

  • activation – Activation function name passed to the transformer feed-forward block.

  • normalization_layer_type – Normalization layer used inside the encoder. BATCH (default) uses custom BatchNorm transformer layers. CHANNEL uses PyTorch’s LayerNorm-based TransformerEncoderLayer.

  • freeze – When True, freezes the backbone weights and only trains the output layer.

  • learning_rate – Base learning rate for the Adam optimizer.

  • lr_step – Milestones (in epochs) for the MultiStepLR scheduler. None means no decay (defaults to a single far-future milestone internally).

  • lr_factor – Multiplicative decay factor applied at each lr_step milestone.

  • weight_decay – L2 regularization coefficient. Applied to the output layer only when global_reg=False, or to all parameters (via optimizer weight decay) when global_reg=True.

  • global_reg – Whether weight_decay is applied globally as weight decay (True) or only to the output layer (False).

  • sync_dist – Whether to synchronize logged metrics across distributed processes.