Skip to content

base

base

Base transformer block: normalization -> operation -> gated residual.

TransformerBlock

TransformerBlock(normalization, dropout=0.1)

Bases: ABC, Module

Composable base building block for transformer layers.

Subclasses implement a specific operation (attention, feedforward) wrapped in normalization and a gated residual connection. The normalization module determines the conditioning behavior: - UnconditionedNorm: ignores condition, gate = ones(1) - AdaNorm: conditioned via (x, condition) -> (normed, gate)

Source code in src/versatil/models/layers/transformer/block/base.py
def __init__(
    self,
    normalization: BlockNormalization,
    dropout: float = 0.1,
):
    super().__init__()
    self.normalization = normalization
    self.residual_dropout = nn.Dropout(dropout)

apply_residual

apply_residual(residual, output, gate)

Gated residual: residual + gate * dropout(output).

Source code in src/versatil/models/layers/transformer/block/base.py
def apply_residual(
    self,
    residual: torch.Tensor,
    output: torch.Tensor,
    gate: torch.Tensor,
) -> torch.Tensor:
    """Gated residual: residual + gate * dropout(output)."""
    return residual + gate * self.residual_dropout(output)