bridge.diffusion.models.common.normalization#

Common normalization modules for diffusion models.

Module Contents#

Classes#

RMSNorm

Root Mean Square Layer Normalization.

API#

class bridge.diffusion.models.common.normalization.RMSNorm(hidden_size: int, config=None, eps: float = 1e-06)#

Bases: torch.nn.Module

Root Mean Square Layer Normalization.

A normalization technique that normalizes the input by its root mean square, then scales by a learnable weight parameter.

Parameters:

  • hidden_size – Size of the hidden dimension.

  • config – Transformer configuration (unused, for compatibility with megatron build_module).

  • eps – Small epsilon for numerical stability.

Initialization

_norm(x)#
forward(x)#