1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

import torch
import torch.nn as nn
 
def visualize_normalization_axes():
    """
    For input shape (N, C, H, W), show which elements each norm averages.
    """
    N, C, H, W = 2, 4, 3, 3  # Small dimensions for clarity
    x = torch.randn(N, C, H, W)
    
    print(f"Input shape: (N={N}, C={C}, H={H}, W={W})")
    print(f"Total elements: {N * C * H * W}")
    print()
    
    # BatchNorm: normalize over N, H, W for each channel
    # Elements per normalization: N * H * W = 18
    # Number of normalizations: C = 4
    print("BatchNorm:")
    print(f"  Elements per normalization: {N * H * W} (N × H × W)")
    print(f"  Number of normalizations: {C} (one per channel)")
    print(f"  γ, β shape: ({C},)")
    
    # LayerNorm: normalize over C, H, W for each sample
    # Elements per normalization: C * H * W = 36
    # Number of normalizations: N = 2
    print("\nLayerNorm([C, H, W]):")
    print(f"  Elements per normalization: {C * H * W} (C × H × W)")
    print(f"  Number of normalizations: {N} (one per sample)")
    print(f"  γ, β shape: ({C}, {H}, {W})")
    
    # InstanceNorm: normalize over H, W for each sample and channel
    # Elements per normalization: H * W = 9
    # Number of normalizations: N * C = 8
    print("\nInstanceNorm:")
    print(f"  Elements per normalization: {H * W} (H × W)")
    print(f"  Number of normalizations: {N * C} (one per sample-channel)")
    print(f"  γ, β shape: ({C},)  [shared across samples]")
    
    # GroupNorm: normalize over (C/G, H, W) for each sample and group
    G = 2  # 2 groups
    print(f"\nGroupNorm (G={G}):")
    print(f"  Elements per normalization: {(C // G) * H * W} (C/G × H × W)")
    print(f"  Number of normalizations: {N * G} (one per sample-group)")
    print(f"  γ, β shape: ({C},)")
 
visualize_normalization_axes()
 
# Demonstrate actual outputs
print("\n" + "="*50)
print("Actual output statistics verification:")
print("="*50)
 
N, C, H, W = 4, 8, 14, 14
x = torch.randn(N, C, H, W) * 3 + 2  # Non-normalized input
 
bn = nn.BatchNorm2d(C)
ln = nn.LayerNorm([C, H, W])
ln_c = nn.LayerNorm([H, W])  # Normalize spatial only
inn = nn.InstanceNorm2d(C, affine=True)
gn = nn.GroupNorm(num_groups=4, num_channels=C)
 
bn.eval()  # Use learned running stats
 
print(f"\nInput: mean={x.mean():.2f}, std={x.std():.2f}")
print(f"BatchNorm: output std per channel ≈ 1")
print(f"LayerNorm[C,H,W]: output std per sample ≈ 1") 
print(f"InstanceNorm: output std per sample-channel ≈ 1")
print(f"GroupNorm(G=4): output std per sample-group ≈ 1")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

import torch
import torch.nn as nn
 
def instance_norm_manual(x, gamma, beta, eps=1e-5):
    """
    Manual Instance Normalization implementation.
    
    Args:
        x: Input of shape (N, C, H, W)
        gamma: Scale, shape (C,)
        beta: Shift, shape (C,)
    """
    N, C, H, W = x.shape
    
    # Compute mean and variance per sample, per channel
    # Average over spatial dimensions (H, W)
    mean = x.mean(dim=(2, 3), keepdim=True)  # (N, C, 1, 1)
    var = x.var(dim=(2, 3), keepdim=True, unbiased=False)  # (N, C, 1, 1)
    
    # Normalize
    x_norm = (x - mean) / torch.sqrt(var + eps)
    
    # Scale and shift (gamma and beta are (C,), need to reshape)
    gamma = gamma.view(1, C, 1, 1)
    beta = beta.view(1, C, 1, 1)
    
    return gamma * x_norm + beta
 
# Compare with PyTorch
torch.manual_seed(42)
N, C, H, W = 4, 8, 32, 32
x = torch.randn(N, C, H, W)
 
inn = nn.InstanceNorm2d(C, affine=True)
y_pytorch = inn(x)
y_manual = instance_norm_manual(x, inn.weight, inn.bias)
 
print(f"Manual matches PyTorch: {torch.allclose(y_pytorch, y_manual, atol=1e-5)}")
 
# Verify per-instance, per-channel normalization
print(f"\nOutput statistics (should be ~0 mean, ~1 std per (n,c) pair):")
for n in range(min(2, N)):
    for c in range(min(3, C)):
        mean = y_pytorch[n, c].mean().item()
        std = y_pytorch[n, c].std().item()
        print(f"  Sample {n}, Channel {c}: mean={mean:.4f}, std={std:.4f}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

import torch
import torch.nn as nn
 
def group_norm_manual(x, num_groups, gamma, beta, eps=1e-5):
    """
    Manual Group Normalization implementation.
    
    Args:
        x: Input of shape (N, C, H, W)
        num_groups: Number of channel groups (G)
        gamma: Scale, shape (C,)
        beta: Shift, shape (C,)
    
    Groups C channels into G groups of C/G channels each.
    Normalizes over (C/G, H, W) for each sample and group.
    """
    N, C, H, W = x.shape
    G = num_groups
    assert C % G == 0, f"C ({C}) must be divisible by num_groups ({G})"
    
    # Reshape to (N, G, C/G, H, W)
    x = x.view(N, G, C // G, H, W)
    
    # Compute mean and variance per sample, per group
    # Average over channels-in-group and spatial: dims (2, 3, 4)
    mean = x.mean(dim=(2, 3, 4), keepdim=True)  # (N, G, 1, 1, 1)
    var = x.var(dim=(2, 3, 4), keepdim=True, unbiased=False)
    
    # Normalize
    x_norm = (x - mean) / torch.sqrt(var + eps)
    
    # Reshape back to (N, C, H, W)
    x_norm = x_norm.view(N, C, H, W)
    
    # Scale and shift
    gamma = gamma.view(1, C, 1, 1)
    beta = beta.view(1, C, 1, 1)
    
    return gamma * x_norm + beta
 
# Compare with PyTorch
torch.manual_seed(42)
N, C, H, W = 4, 32, 8, 8
num_groups = 8
x = torch.randn(N, C, H, W)
 
gn = nn.GroupNorm(num_groups=num_groups, num_channels=C)
y_pytorch = gn(x)
y_manual = group_norm_manual(x, num_groups, gn.weight, gn.bias)
 
print(f"Manual matches PyTorch: {torch.allclose(y_pytorch, y_manual, atol=1e-5)}")
 
# Analyze group statistics
print(f"\nGroup structure: {C} channels / {num_groups} groups = {C // num_groups} channels per group")
print(f"Elements per normalization: {(C // num_groups) * H * W}")
 
# Verify normalization
y_reshaped = y_pytorch.view(N, num_groups, C // num_groups, H, W)
for n in range(min(2, N)):
    for g in range(min(2, num_groups)):
        mean = y_reshaped[n, g].mean().item()
        std = y_reshaped[n, g].std().item()
        print(f"  Sample {n}, Group {g}: mean={mean:.4f}, std={std:.4f}")
 
# Compare performance with different batch sizes
print("\n--- Batch Size Comparison (BN vs GN) ---")
for batch_size in [1, 2, 4, 32]:
    x_test = torch.randn(batch_size, 32, 8, 8)
    
    bn = nn.BatchNorm2d(32)
    gn = nn.GroupNorm(8, 32)
    
    bn.train()  # BN has issues in train mode with batch=1
    
    try:
        y_bn = bn(x_test)
        bn_std = y_bn.std().item()
        bn_status = f"std={bn_std:.2f}"
    except Exception as e:
        bn_status = "Error!"
    
    y_gn = gn(x_test)
    gn_std = y_gn.std().item()
    
    print(f"Batch {batch_size:2d}: BN {bn_status:12s} | GN std={gn_std:.2f}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm, remove_weight_norm
 
class WeightNormLinear(nn.Module):
    """
    Manual implementation of Weight Normalization for understanding.
    
    Actual usage: Use torch.nn.utils.weight_norm wrapper.
    """
    def __init__(self, in_features, out_features):
        super().__init__()
        
        # Direction parameters
        self.v = nn.Parameter(torch.randn(out_features, in_features))
        
        # Magnitude parameters (one per output feature)
        self.g = nn.Parameter(torch.ones(out_features))
        
        self.bias = nn.Parameter(torch.zeros(out_features))
    
    def forward(self, x):
        # Compute normalized weights
        # w_i = g_i * v_i / ||v_i||
        v_norm = self.v.norm(dim=1, keepdim=True)  # (out, 1)
        w = self.g.unsqueeze(1) * self.v / v_norm
        
        return F.linear(x, w, self.bias)
 
# Compare with PyTorch's weight_norm
torch.manual_seed(42)
in_features, out_features = 64, 32
 
# Manual implementation
wn_manual = WeightNormLinear(in_features, out_features)
 
# PyTorch implementation
linear = nn.Linear(in_features, out_features)
wn_pytorch = weight_norm(linear, name='weight')
 
print("Weight Normalization structure:")
print(f"  Direction (v): shape={wn_manual.v.shape}")
print(f"  Magnitude (g): shape={wn_manual.g.shape}")
print(f"  Effective weight: g * v / ||v||")
 
# The reparameterization
v_norm = wn_manual.v.norm(dim=1, keepdim=True)
effective_weight = wn_manual.g.unsqueeze(1) * wn_manual.v / v_norm
print(f"\nEffective weight norm per output: {effective_weight.norm(dim=1)[:5]}")
print(f"g values (should match): {wn_manual.g[:5]}")
 
# Gradient analysis: g and v gradients are decoupled
x = torch.randn(8, in_features)
y = wn_manual(x)
loss = y.sum()
loss.backward()
 
print(f"\nGradient properties:")
print(f"  ∂L/∂g shape: {wn_manual.g.grad.shape}")
print(f"  ∂L/∂v shape: {wn_manual.v.grad.shape}")
 
# Remove weight norm to get regular linear layer
remove_weight_norm(wn_pytorch)
print(f"\nAfter removing weight norm, layer is regular Linear")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
 
def power_iteration(W, u, n_iterations=1):
    """
    Power iteration to approximate the largest singular value.
    
    The spectral norm is ||W||_2 = σ_1 (largest singular value).
    Power iteration efficiently approximates this.
    """
    for _ in range(n_iterations):
        # v = W^T u / ||W^T u||
        v = W.t() @ u
        v = v / v.norm()
        
        # u = W v / ||W v||
        u = W @ v
        u = u / u.norm()
    
    # σ_1 ≈ u^T W v
    sigma = (u @ W @ v).item()
    
    return sigma, u, v
 
class SpectralNormLinear(nn.Module):
    """
    Manual implementation of Spectral Normalization for understanding.
    
    Actual usage: Use torch.nn.utils.spectral_norm wrapper.
    """
    def __init__(self, in_features, out_features, n_power_iterations=1):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.n_power_iterations = n_power_iterations
        
        # Initialize u vector for power iteration
        self.register_buffer('u', torch.randn(out_features))
        self.u = self.u / self.u.norm()
    
    def spectral_norm(self):
        """Compute spectrally normalized weight."""
        u = self.u
        W = self.weight
        
        for _ in range(self.n_power_iterations):
            v = W.t() @ u
            v = v / v.norm()
            u = W @ v
            u = u / u.norm()
        
        # Update u buffer (detached from graph)
        self.u = u.detach()
        
        # Compute spectral norm
        sigma = (u @ W @ v).item()
        
        # Return normalized weight
        return self.weight / sigma
    
    def forward(self, x):
        W_sn = self.spectral_norm()
        return F.linear(x, W_sn, self.bias)
 
# Compare with PyTorch's spectral_norm
torch.manual_seed(42)
in_features, out_features = 64, 32
 
# PyTorch implementation
linear = nn.Linear(in_features, out_features)
sn_layer = spectral_norm(linear)
 
# Check spectral norm of weight
W = sn_layer.weight
U, S, V = torch.linalg.svd(W)
spectral_norm_value = S[0].item()
 
print(f"Spectral norm of weight: {spectral_norm_value:.4f}")
print(f"(Should be close to 1.0 after training)")
 
# Demonstrate Lipschitz property
x1 = torch.randn(1, in_features)
x2 = torch.randn(1, in_features)
 
y1 = sn_layer(x1)
y2 = sn_layer(x2)
 
input_diff = (x1 - x2).norm().item()
output_diff = (y1 - y2).norm().item()
lipschitz_ratio = output_diff / input_diff
 
print(f"\nLipschitz check:")
print(f"  ||x1 - x2|| = {input_diff:.4f}")
print(f"  ||y1 - y2|| = {output_diff:.4f}")
print(f"  Ratio = {lipschitz_ratio:.4f} (should be ≤ 1.0 for 1-Lipschitz)")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class SwitchableNorm2d(nn.Module):
    """
    Switchable Normalization: learns to combine BN, IN, and LN.
    """
    def __init__(self, num_features, eps=1e-5):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        
        # Learnable combination weights
        self.weight_bn = nn.Parameter(torch.ones(1))
        self.weight_in = nn.Parameter(torch.ones(1))
        self.weight_ln = nn.Parameter(torch.ones(1))
        
        # Scale and shift
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        
        # Running stats for BN component
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.momentum = 0.1
    
    def forward(self, x):
        N, C, H, W = x.shape
        
        # Softmax over combination weights
        weights = F.softmax(torch.stack([self.weight_bn, self.weight_in, self.weight_ln]), dim=0)
        
        # BatchNorm statistics
        mean_bn = x.mean(dim=(0, 2, 3), keepdim=True)
        var_bn = x.var(dim=(0, 2, 3), keepdim=True, unbiased=False)
        
        # InstanceNorm statistics
        mean_in = x.mean(dim=(2, 3), keepdim=True)
        var_in = x.var(dim=(2, 3), keepdim=True, unbiased=False)
        
        # LayerNorm statistics
        mean_ln = x.mean(dim=(1, 2, 3), keepdim=True)
        var_ln = x.var(dim=(1, 2, 3), keepdim=True, unbiased=False)
        
        # Combine statistics
        mean = weights[0] * mean_bn + weights[1] * mean_in + weights[2] * mean_ln
        var = weights[0] * var_bn + weights[1] * var_in + weights[2] * var_ln
        
        # Normalize
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        
        # Scale and shift
        gamma = self.gamma.view(1, C, 1, 1)
        beta = self.beta.view(1, C, 1, 1)
        
        return gamma * x_norm + beta
 
 
class FilterResponseNorm(nn.Module):
    """
    Filter Response Normalization (FRN) with Thresholded Linear Unit (TLU).
    
    Proposed as a batch-independent alternative to BatchNorm+ReLU.
    Normalizes filter responses with learned threshold.
    """
    def __init__(self, num_features, eps=1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(1, num_features, 1, 1))
        self.beta = nn.Parameter(torch.zeros(1, num_features, 1, 1))
        self.tau = nn.Parameter(torch.zeros(1, num_features, 1, 1))  # Threshold
        self.eps = eps
    
    def forward(self, x):
        # Compute mean squared value (not centered)
        nu2 = x.pow(2).mean(dim=(2, 3), keepdim=True)
        
        # Normalize by RMS
        x_norm = x / torch.sqrt(nu2 + self.eps)
        
        # Scale and shift
        y = self.gamma * x_norm + self.beta
        
        # TLU activation: max(y, tau)
        return torch.max(y, self.tau)
 
 
# Demonstrate usage
torch.manual_seed(42)
x = torch.randn(4, 32, 8, 8)
 
sn = SwitchableNorm2d(32)
frn = FilterResponseNorm(32)
 
y_sn = sn(x)
y_frn = frn(x)
 
print("Switchable Normalization:")
print(f"  Learned weights (after softmax): {F.softmax(torch.stack([sn.weight_bn, sn.weight_in, sn.weight_ln]), dim=0).squeeze()}")
 
print("\nFilter Response Normalization:")
print(f"  Works without batch statistics")
print(f"  Includes learned threshold activation")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

def select_normalization(
    architecture,
    batch_size,
    task,
    sequence_model=False,
    style_sensitive=False,
    small_batch_required=False
):
    """
    Decision tree for selecting normalization technique.
    """
    
    # Sequence models (Transformers, RNNs)
    if sequence_model:
        if architecture == "transformer":
            return "LayerNorm (possibly RMSNorm for efficiency)"
        elif architecture == "rnn" or architecture == "lstm":
            return "LayerNorm or Weight Normalization"
        else:
            return "LayerNorm"
    
    # GANs
    if task == "gan":
        if "discriminator" in architecture.lower():
            return "Spectral Normalization"
        else:  # generator
            return "Instance Normalization or Conditional BatchNorm"
    
    # Style Transfer / Image-to-Image
    if style_sensitive or task in ["style_transfer", "image_synthesis"]:
        return "Instance Normalization or AdaIN"
    
    # Object Detection / Segmentation (often small batch)
    if task in ["detection", "segmentation"] or small_batch_required:
        if batch_size < 8:
            return "Group Normalization (G=32)"
        else:
            return "BatchNorm or GroupNorm"
    
    # Standard CNN Training
    if architecture in ["cnn", "resnet", "efficientnet"]:
        if batch_size >= 16:
            return "BatchNorm"
        elif batch_size >= 4:
            return "GroupNorm"
        else:
            return "LayerNorm or GroupNorm"
    
    # Reinforcement Learning
    if task == "reinforcement_learning":
        return "LayerNorm or Weight Normalization"
    
    # Default
    return "BatchNorm (CNN) or LayerNorm (other)"
 
# Example decisions
cases = [
    {"architecture": "transformer", "batch_size": 32, "task": "language_model", 
     "sequence_model": True},
    {"architecture": "resnet", "batch_size": 2, "task": "detection", 
     "small_batch_required": True},
    {"architecture": "generator", "batch_size": 16, "task": "gan"},
    {"architecture": "cnn", "batch_size": 64, "task": "classification"},
    {"architecture": "stylegan", "batch_size": 4, "task": "style_transfer",
     "style_sensitive": True},
]
 
print("Normalization Selection Examples:")
print("=" * 60)
for case in cases:
    result = select_normalization(**case)
    print(f"\n{case}")
    print(f"→ {result}")