Layer Normalization › Visualizing Layer Normalization's Effect

class StackedTransformerBlocks(nn.Module):
    """Stack of transformer blocks for visualization."""

    def __init__(self, d_model, n_heads, d_ff, n_layers, use_layernorm=True):
        super().__init__()
        self.use_layernorm = use_layernorm
        self.layers = nn.ModuleList(
            [
                PreNormTransformerBlock(d_model, n_heads, d_ff)
                for _ in range(n_layers)
            ]
        )
        if not use_layernorm:
            # Replace LayerNorm with identity
            for layer in self.layers:
                layer.norm1 = nn.Identity()
                layer.norm2 = nn.Identity()

    def forward_with_activations(self, x):
        """Return activations after each layer."""
        activations = [x.detach().clone()]
        for layer in self.layers:
            x = layer(x)
            activations.append(x.detach().clone())
        return activations


# Create models with and without layer normalization
d_model, n_heads, d_ff, n_layers = 64, 4, 256, 10

model_with_ln = StackedTransformerBlocks(
    d_model, n_heads, d_ff, n_layers, use_layernorm=True
)
model_without_ln = StackedTransformerBlocks(
    d_model, n_heads, d_ff, n_layers, use_layernorm=False
)

# Forward pass with larger initial variance to amplify divergence without LayerNorm
x = torch.randn(4, 16, d_model) * 2.0
acts_with_ln = model_with_ln.forward_with_activations(x)
acts_without_ln = model_without_ln.forward_with_activations(x)