ada_norm

ada_norm

AdaNorm

AdaNorm(base_norm, conditioning_dimension, feature_dim, use_gate=False, activation=value, init_strategy='zero')

Bases: Module

Adaptive normalization layer with conditional affine modulation.

Initialize adaptive normalization.

Parameters:

Name	Type	Description	Default
base_norm	Module	Normalization module applied before modulation.	required
conditioning_dimension	int	Dimension of the conditioning vector.	required
feature_dim	int	Feature dimension to modulate.	required
use_gate	bool	Whether to return a learned residual gate.	False
activation	str	Activation used inside the modulation projection.	value
init_strategy	Literal['zero', 'xavier']	Initialization strategy for modulation weights.	'zero'

Source code in src/versatil/models/layers/normalization/ada_norm.py

def __init__(
    self,
    base_norm: nn.Module,
    conditioning_dimension: int,
    feature_dim: int,
    use_gate: bool = False,
    activation: str = ActivationFunction.SILU.value,
    init_strategy: Literal["zero", "xavier"] = "zero",
):
    """Initialize adaptive normalization.

    Args:
        base_norm: Normalization module applied before modulation.
        conditioning_dimension: Dimension of the conditioning vector.
        feature_dim: Feature dimension to modulate.
        use_gate: Whether to return a learned residual gate.
        activation: Activation used inside the modulation projection.
        init_strategy: Initialization strategy for modulation weights.
    """
    super().__init__()
    self.norm = base_norm
    self.conditioning_dimension = conditioning_dimension
    self.feature_dim = feature_dim
    self.activation = activation
    self.modulation = ConditionalModulation(
        conditioning_dimension=conditioning_dimension,
        feature_dim=feature_dim,
        use_shift=True,
        use_gate=use_gate,
        activation=activation,
        init_strategy=init_strategy,
    )

forward

forward(x, condition)

Forward pass with conditioning.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor to normalize and modulate.	required
condition	Tensor	Conditioning tensor of shape (batch_size, conditioning_dimension).	required

Returns:

Type	Description
Tensor	Tuple of (normalized+modulated x, gate). Gate is a learned
Tensor	tensor when use_gate=True, or 1.0 when use_gate=False.

Source code in src/versatil/models/layers/normalization/ada_norm.py

def forward(
    self, x: torch.Tensor, condition: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Forward pass with conditioning.

    Args:
        x: Input tensor to normalize and modulate.
        condition: Conditioning tensor of shape ``(batch_size, conditioning_dimension)``.

    Returns:
        Tuple of (normalized+modulated x, gate). Gate is a learned
        tensor when use_gate=True, or 1.0 when use_gate=False.
    """
    x = self.norm(x)
    return self.modulation(x, condition)