query_key_norm

query_key_norm

Query-Key Normalization for attention mechanisms.

Applies RMSNorm to query and key tensors in the attention head dimension for improved training stability at scale. Used in SD3/MMDiT architectures.

References

Esser et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis" https://arxiv.org/abs/2403.03206

QueryKeyNorm

QueryKeyNorm(head_dimension, epsilon=1e-06)

Bases: Module

Query-Key normalization using RMSNorm.

Normalizes query and key tensors independently using RMSNorm in the head dimension before computing attention. This improves training stability especially when scaling to larger models.

Initialize QueryKeyNorm.

Parameters:

Name	Type	Description	Default
head_dimension	int	Dimension of each attention head.	required
epsilon	float	Small constant for numerical stability.	1e-06

Source code in src/versatil/models/layers/transformer/attention/query_key_norm.py

def __init__(
    self,
    head_dimension: int,
    epsilon: float = 1e-6,
):
    """Initialize QueryKeyNorm.

    Args:
        head_dimension: Dimension of each attention head.
        epsilon: Small constant for numerical stability.
    """
    super().__init__()
    self.query_norm = RMSNorm(
        head_dimension, epsilon=epsilon, elementwise_affine=True
    )
    self.key_norm = RMSNorm(
        head_dimension, epsilon=epsilon, elementwise_affine=True
    )

forward

forward(query, key)

Apply RMSNorm to query and key tensors.

Parameters:

Name	Type	Description	Default
query	Tensor	Query tensor (B, number_of_heads, sequence_length, head_dimension).	required
key	Tensor	Key tensor (B, number_of_heads, sequence_length, head_dimension).	required

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of normalized (query, key) tensors with same shapes.

Source code in src/versatil/models/layers/transformer/attention/query_key_norm.py

def forward(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply RMSNorm to query and key tensors.

    Args:
        query: Query tensor (B, number_of_heads, sequence_length, head_dimension).
        key: Key tensor (B, number_of_heads, sequence_length, head_dimension).

    Returns:
        Tuple of normalized (query, key) tensors with same shapes.
    """
    return self.query_norm(query), self.key_norm(key)