encoder_layer

encoder_layer

Transformer encoder layer inspired by the original "Attention is All You Need" paper, with bidirectional self-attention and optional conditioning.

TransformerEncoderLayer

TransformerEncoderLayer(embedding_dimension, number_of_heads, number_of_key_value_heads=None, feedforward_dimension=None, dropout=0.1, attention_dropout=0.0, activation=value, normalization_type=value, attention_type=value, bias=True, normalization_epsilon=1e-06, conditioning_dimension=None, use_gating=False)

Bases: Module

Self-attention + feedforward blocks.

Note

Supports optional conditioning when constructed with adaptive normalization types.

Initialize Transformer encoder layer.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Model embedding dimension.	required
number_of_heads	int	Number of attention heads.	required
number_of_key_value_heads	int | None	Number of K/V heads (for GQA).	None
feedforward_dimension	int | None	FFN hidden dimension (defaults to 4 * embedding_dimension).	None
dropout	float	Dropout probability for residual connections.	0.1
attention_dropout	float	Dropout probability for attention weights.	0.0
activation	str	Activation function (use ActivationFunction enum values).	value
normalization_type	str	Type of normalization (use NormalizationType enum values).	value
attention_type	str	Type of attention (use AttentionType enum values).	value
bias	bool	Whether to use bias in linear layers.	True
normalization_epsilon	float	Epsilon for normalization layers.	1e-06
conditioning_dimension	int | None	Conditioning dimension for adaptive normalization. Required when normalization_type is adaptive.	None
use_gating	bool	Whether to use gating in adaptive normalization (AdaLN-Zero).	False

Source code in src/versatil/models/layers/transformer/layer/encoder_layer.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    number_of_key_value_heads: int | None = None,
    feedforward_dimension: int | None = None,
    dropout: float = 0.1,
    attention_dropout: float = 0.0,
    activation: str = ActivationFunction.SWIGLU.value,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    attention_type: str = AttentionType.GROUPED_QUERY.value,
    bias: bool = True,
    normalization_epsilon: float = 1e-6,
    conditioning_dimension: int | None = None,
    use_gating: bool = False,
):
    """Initialize Transformer encoder layer.

    Args:
        embedding_dimension: Model embedding dimension.
        number_of_heads: Number of attention heads.
        number_of_key_value_heads: Number of K/V heads (for GQA).
        feedforward_dimension: FFN hidden dimension (defaults to 4 * embedding_dimension).
        dropout: Dropout probability for residual connections.
        attention_dropout: Dropout probability for attention weights.
        activation: Activation function (use ActivationFunction enum values).
        normalization_type: Type of normalization (use NormalizationType enum values).
        attention_type: Type of attention (use AttentionType enum values).
        bias: Whether to use bias in linear layers.
        normalization_epsilon: Epsilon for normalization layers.
        conditioning_dimension: Conditioning dimension for adaptive normalization.
            Required when normalization_type is adaptive.
        use_gating: Whether to use gating in adaptive normalization (AdaLN-Zero).
    """
    super().__init__()
    self.embedding_dimension = embedding_dimension
    self.number_of_heads = number_of_heads
    if feedforward_dimension is None:
        feedforward_dimension = 4 * embedding_dimension
    self.self_attention_block = SelfAttentionBlock(
        attention=CachedAttention(
            embedding_dimension=embedding_dimension,
            number_of_heads=number_of_heads,
            number_of_key_value_heads=number_of_key_value_heads,
            dropout=attention_dropout,
            bias=bias,
            attention_type=attention_type,
        ),
        normalization=create_block_normalization(
            normalization_type=normalization_type,
            dimension=embedding_dimension,
            epsilon=normalization_epsilon,
            conditioning_dimension=conditioning_dimension,
            use_gating=use_gating,
        ),
        dropout=dropout,
    )
    self.feedforward_block = FeedforwardBlock(
        feedforward=build_feedforward(
            embedding_dimension=embedding_dimension,
            feedforward_dimension=feedforward_dimension,
            activation=activation,
            dropout=dropout,
            bias=bias,
        ),
        normalization=create_block_normalization(
            normalization_type=normalization_type,
            dimension=embedding_dimension,
            epsilon=normalization_epsilon,
            conditioning_dimension=conditioning_dimension,
            use_gating=use_gating,
        ),
        dropout=dropout,
    )

forward

forward(hidden_states, conditioning=None, attention_mask=None, positional_encoding=None)

Forward pass through encoder layer.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Input embeddings (B, S, D).	required
conditioning	Tensor | None	Conditioning vector for adaptive normalization (B, C). Ignored when constructed with plain normalization.	None
attention_mask	Tensor | None	Optional mask (B, 1, S, S) where True means masked.	None
positional_encoding	RotaryPositionalEncoding | None	Optional rotary positional encoding module.	None

Returns:

Type	Description
Tensor	Output hidden states (B, S, D).

Source code in src/versatil/models/layers/transformer/layer/encoder_layer.py

def forward(
    self,
    hidden_states: torch.Tensor,
    conditioning: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    positional_encoding: RotaryPositionalEncoding | None = None,
) -> torch.Tensor:
    """Forward pass through encoder layer.

    Args:
        hidden_states: Input embeddings (B, S, D).
        conditioning: Conditioning vector for adaptive normalization (B, C).
            Ignored when constructed with plain normalization.
        attention_mask: Optional mask (B, 1, S, S) where True means masked.
        positional_encoding: Optional rotary positional encoding module.

    Returns:
        Output hidden states (B, S, D).
    """
    hidden_states, _ = self.self_attention_block(
        hidden_states=hidden_states,
        conditioning=conditioning,
        attention_mask=attention_mask,
        positional_encoding=positional_encoding,
    )
    hidden_states = self.feedforward_block(
        hidden_states=hidden_states,
        conditioning=conditioning,
    )
    return hidden_states