TransformerLayer

transformer.TransformerLayer

A transformer block that applies multihead attention followed by a feedforward network.

The transformer can be configured to apply time-wise attention (i.e. attention over the time axis) or space-wise attention (i.e. attention over the variate axis).

The transformer block uses pre-norm, which is a variant of the transformer architecture where LayerNorm is applied before each sublayer, rather than after. This is the approach taken in LLaMA and other recent transformer-based models.

The transformer block also uses SwiGLU, which is a variant of the Gated Linear Unit (GLU) activation function. SwiGLU is a variant of the GLU activation that uses the Swish activation function. This activation function has been used extensively in recent transformer-based models and has been shown to improve performance.

__init__(self, embed_dim: int, num_heads: int, mlp_hidden_dim: int, dropout: float, rotary_emb: RotaryEmbedding = None, attention_axis: AttentionAxis = AttentionAxis.TIME, RMS_norm: bool = True, use_memory_efficient_attention: bool = True)

Parameter	Type	Default	Description
embed_dim	int	-	(undocumented)
num_heads	int	-	(undocumented)
mlp_hidden_dim	int	-	(undocumented)
dropout	float	-	(undocumented)
rotary_emb	RotaryEmbedding	None	(undocumented)
attention_axis	AttentionAxis	AttentionAxis.TIME	(undocumented)
RMS_norm	bool	True	(undocumented)
use_memory_efficient_attention	bool	True	(undocumented)

forward(self, layer_idx: int, inputs: Float[torch.Tensor, "batch variate seq_len embed_dim"], attention_mask: Optional[Union[Bool[torch.Tensor, "batch seq_len variate variate"], Bool[torch.Tensor, "batch #variate seq_len seq_len"]]] = None, kv_cache: Optional[KVCache] = None) -> Float[torch.Tensor, "batch variate seq_len embed_dim"]

Parameters:

Parameter	Type	Default	Description
layer_idx	int	-	(undocumented)
inputs	Float[torch.Tensor, "batch variate seq_len embed_dim"]	-	(undocumented)
attention_mask	Optional[Union[Bool[torch.Tensor, "batch seq_len variate variate"], Bool[torch.Tensor, "batch #variate seq_len seq_len"]]]	None	(undocumented)
kv_cache	Optional[KVCache]	None	(undocumented)

Returns: Float[torch.Tensor, "batch variate seq_len embed_dim"]

Transformer

transformer.Transformer

A stack of transformer layers. The transformer alternates between time-wise and space-wise attention to learn both temporal and cross-variate dependencies in the data.

Based on the intuition that time-wise attention is more important overall than space-wise attention (because an individual variate is more likely to be correlated with itself across time than with other variates), the transformer can be configured to apply space-wise attention less frequently than time-wise attention. This is controlled by the spacewise_every_n_layers parameter, which specifies how many time-wise transformer layers to apply between every space-wise transformer layer.

__init__(self, num_layers: int, embed_dim: int, num_heads: int, mlp_hidden_dim: int, dropout: float, spacewise_every_n_layers: int, spacewise_first: bool, use_memory_efficient_attention: bool = True)

Parameter	Type	Default	Description
num_layers	int	-	Number of transformer layers to use.
embed_dim	int	-	(undocumented)
num_heads	int	-	Number of attention heads to use in each self-attention layer.
mlp_hidden_dim	int	-	Dimension of the hidden layer in the feedforward network.
dropout	float	-	Dropout rate to use in the model.
spacewise_every_n_layers	int	-	How many time-wise transformer layers to apply between each space-wise transformer layer.
spacewise_first	bool	-	Whether to apply space-wise attention before time-wise attention.
use_memory_efficient_attention	bool	True	Whether to use memory-efficient attention. If True, the model will use the memory-efficient from xFormers.

forward(self, inputs: Float[torch.Tensor, "batch variate seq_len embed_dim"], id_mask: Float[torch.Tensor, "batch #variate seq_len"], kv_cache: Optional[KVCache] = None) -> Float[torch.Tensor, "batch variate seq_len embed_dim"]

Parameters:

Parameter	Type	Default	Description
inputs	Float[torch.Tensor, "batch variate seq_len embed_dim"]	-	(undocumented)
id_mask	Float[torch.Tensor, "batch #variate seq_len"]	-	(undocumented)
kv_cache	Optional[KVCache]	None	(undocumented)

Returns: Float[torch.Tensor, "batch variate seq_len embed_dim"]