normalization_via_fluxions.md

# Normalization Layers via the Method of Fluxions
## BatchNorm, LayerNorm, RMSNorm: What They Actually Do

**Scott Bisset, Silicon Goddess**  
OpenTransformers Ltd  
January 2026

---

## Abstract

Normalization layers (BatchNorm, LayerNorm, RMSNorm) are presented in textbooks as "normalize then scale and shift" with formulas involving means and variances. This obscures their true purpose and makes backward pass derivation seem magical. We reformulate normalization using fluxions, revealing: (1) normalization as signal conditioning, (2) the backward pass as sensitivity redistribution, and (3) why different norms suit different architectures. The fluxion view also explains the computational structure that enables fused kernels.

---

## 1. Why Normalize?

### 1.1 The Problem

Deep networks suffer from internal covariate shift:
- Each layer's input distribution changes during training
- Later layers constantly adapt to moving targets
- Training becomes unstable

### 1.2 The Solution Intuition

Force each layer's inputs to have consistent statistics.
"Standardize the signal before processing it."

---

## 2. The General Normalization Framework

### 2.1 Forward Pass

All normalization layers follow:

```
1. Compute statistics (μ, σ) over some dimension
2. Normalize: x̂ = (x - μ) / σ
3. Scale and shift: y = γ·x̂ + β
```

What differs is WHICH dimensions we compute statistics over.

### 2.2 Fluxion View

Let x be the input signal.

```
μ = mean(x)           # Signal center
σ = std(x)            # Signal spread
x̂ = (x - μ) / σ       # Centered and scaled to unit variance
y = γ·x̂ + β           # Learnable rescaling
```

**γ** (gamma) = learned amplitude
**β** (beta) = learned offset

Without γ and β, normalization would constrain representational power.
With them, the network can learn to undo normalization if needed.

---

## 3. BatchNorm: Normalize Across Batch

### 3.1 The Idea

For each feature/channel, compute mean and variance ACROSS the batch.

```
Input: X of shape [B, D]  (B samples, D features)

For each feature d:
    μ_d = (1/B) Σᵢ X[i,d]        # Mean of feature d across batch
    σ_d = sqrt((1/B) Σᵢ (X[i,d] - μ_d)²)  # Std of feature d
    
    X̂[:,d] = (X[:,d] - μ_d) / σ_d
    Y[:,d] = γ_d · X̂[:,d] + β_d
```

### 3.2 Fluxion Forward Pass

```
μ = mean(X, dim=batch)          # Shape: [D]
σ = std(X, dim=batch)           # Shape: [D]
X̂ = (X - μ) / σ                 # Shape: [B, D]
Y = γ·X̂ + β                     # Shape: [B, D]
```

### 3.3 The Backward Pass (Fluxion Derivation)

Given L̇ʸ (upstream gradient), find L̇ˣ, L̇ᵞ, L̇ᵝ.

**Easy ones first:**

```
L̇ᵝ = sum(L̇ʸ, dim=batch)         # β gradient = sum of upstream
L̇ᵞ = sum(L̇ʸ · X̂, dim=batch)     # γ gradient = upstream weighted by normalized input
```

**L̇ˣ is tricky because μ and σ depend on ALL x values.**

Let's trace the wiggle:

```
If x[i,d] wiggles:
    1. Direct effect on X̂[i,d]: ∂X̂[i,d]/∂x[i,d] = 1/σ
    2. Indirect effect via μ: changing x[i,d] shifts μ, affects ALL X̂[:,d]  
    3. Indirect effect via σ: changing x[i,d] changes σ, affects ALL X̂[:,d]
```

Full derivative:

```
L̇ˣ̂ = L̇ʸ · γ                     # Gradient through scale

L̇σ = -sum(L̇ˣ̂ · (X-μ) / σ², dim=batch)   # How σ wiggle affects loss

L̇μ = -sum(L̇ˣ̂ / σ, dim=batch) + L̇σ · (-2/B)·sum(X-μ, dim=batch)

L̇ˣ = L̇ˣ̂/σ + L̇σ·(2/B)·(X-μ)/σ + L̇μ/B
```

### 3.4 Simplified Form

After algebra, the BatchNorm backward becomes:

```
L̇ˣ = (1/σ) · (L̇ˣ̂ - mean(L̇ˣ̂) - X̂·mean(L̇ˣ̂·X̂))
```

**Interpretation:**
- Start with scaled upstream gradient
- Subtract its mean (center the gradient)  
- Subtract correlation with normalized input (decorrelate)

"BatchNorm backward REDISTRIBUTES gradient to maintain zero-mean, unit-variance gradient flow."

### 3.5 Inference Mode

At inference, we don't have a batch. Use running averages from training:

```
μ_running = momentum·μ_running + (1-momentum)·μ_batch
σ_running = momentum·σ_running + (1-momentum)·σ_batch
```

Then normalize with running stats.

### 3.6 Problems with BatchNorm

1. **Batch size dependence**: Small batches → noisy statistics
2. **Not suitable for sequence models**: Each position needs different batch members
3. **Inference/training mismatch**: Running stats ≠ batch stats

---

## 4. LayerNorm: Normalize Across Features

### 4.1 The Idea

For each sample, compute mean and variance ACROSS features.

```
Input: X of shape [B, D]

For each sample i:
    μᵢ = (1/D) Σ_d X[i,d]        # Mean across features
    σᵢ = sqrt((1/D) Σ_d (X[i,d] - μᵢ)²)
    
    X̂[i,:] = (X[i,:] - μᵢ) / σᵢ
    Y[i,:] = γ · X̂[i,:] + β
```

### 4.2 Fluxion Forward Pass

```
μ = mean(X, dim=features)       # Shape: [B]
σ = std(X, dim=features)        # Shape: [B]
X̂ = (X - μ) / σ                 # Shape: [B, D]
Y = γ·X̂ + β                     # Shape: [B, D]
```

### 4.3 Key Difference from BatchNorm

```
BatchNorm: statistics across BATCH (each feature normalized independently)
LayerNorm: statistics across FEATURES (each sample normalized independently)
```

### 4.4 Why LayerNorm for Transformers?

1. **No batch dependence**: Each token normalized independently
2. **Works with any batch size**: Including batch=1 at inference
3. **Sequence-friendly**: Position i doesn't need position j's statistics

### 4.5 Backward Pass

Same structure as BatchNorm, but sum over features instead of batch:

```
L̇ˣ̂ = L̇ʸ · γ

L̇ˣ = (1/σ) · (L̇ˣ̂ - mean(L̇ˣ̂, dim=features) 
              - X̂·mean(L̇ˣ̂·X̂, dim=features))
```

---

## 5. RMSNorm: Skip the Mean

### 5.1 The Simplification

LayerNorm computes mean AND variance.
RMSNorm: "What if we skip the mean centering?"

```
RMS(x) = sqrt(mean(x²))
X̂ = X / RMS(X)
Y = γ · X̂
```

No β parameter (no shift), no mean subtraction.

### 5.2 Fluxion Forward Pass

```
rms = sqrt(mean(X², dim=features))   # Shape: [B]
X̂ = X / rms                          # Shape: [B, D]
Y = γ · X̂                            # Shape: [B, D]
```

### 5.3 Why It Works

Empirically, the mean-centering in LayerNorm contributes little.
RMSNorm achieves similar performance with:
- Fewer operations
- Simpler backward pass
- Better numerical stability

### 5.4 Backward Pass

Much simpler without mean:

```
L̇ˣ̂ = L̇ʸ · γ
L̇ʳᵐˢ = -sum(L̇ˣ̂ · X / rms²)
L̇ˣ = L̇ˣ̂/rms + L̇ʳᵐˢ · X/(D·rms)
```

Simplified:
```
L̇ˣ = (1/rms) · (L̇ˣ̂ - X̂·mean(L̇ˣ̂·X̂))
```

One fewer term than LayerNorm!

### 5.5 Usage

RMSNorm is used in:
- LLaMA
- Mistral
- Most modern LLMs

It's becoming the default for transformers.

---

## 6. Comparison Table

| Property | BatchNorm | LayerNorm | RMSNorm |
|----------|-----------|-----------|---------|
| Stats over | Batch | Features | Features |
| Learnable | γ, β | γ, β | γ only |
| Mean centering | Yes | Yes | No |
| Batch dependent | Yes | No | No |
| Inference mode | Running stats | Same as training | Same as training |
| Use case | CNNs | Transformers | Modern LLMs |
| Operations | Most | Medium | Fewest |

---

## 7. Pre-Norm vs Post-Norm

### 7.1 Post-Norm (Original Transformer)

```
X → Attention → Add(X) → LayerNorm → FFN → Add → LayerNorm → Output
```

Normalize AFTER residual connection.

### 7.2 Pre-Norm (Modern Default)

```
X → LayerNorm → Attention → Add(X) → LayerNorm → FFN → Add → Output
```

Normalize BEFORE attention/FFN.

### 7.3 Fluxion Analysis

**Post-Norm gradient flow:**
```
L̇ˣ = LayerNorm_backward(L̇ᵒᵘᵗ)   # Gradient must flow through norm
```

**Pre-Norm gradient flow:**
```
L̇ˣ = L̇ᵒᵘᵗ + LayerNorm_backward(Attention_backward(L̇))
    ↑
    Direct highway!
```

Pre-Norm has a direct gradient path that bypasses normalization.
This stabilizes training for deep networks.

---

## 8. Numerical Stability

### 8.1 The Variance Problem

Computing variance naively:
```
var = mean(x²) - mean(x)²
```

If mean(x²) ≈ mean(x)², subtraction causes catastrophic cancellation.

### 8.2 Welford's Algorithm

Compute variance in a single pass, numerically stable:

```python
def welford_var(x):
    n = 0
    mean = 0
    M2 = 0
    for xi in x:
        n += 1
        delta = xi - mean
        mean += delta / n
        delta2 = xi - mean
        M2 += delta * delta2
    return M2 / n
```

### 8.3 Fused Kernels

The fluxion view reveals that forward and backward are tightly coupled:

Forward needs: μ, σ, X̂
Backward needs: μ, σ, X̂ (same!)

Fused kernel can:
1. Compute μ, σ in one pass
2. Store only X̂ (derived from X, μ, σ)
3. Backward reuses X̂ directly

This is why PyTorch's native LayerNorm is much faster than naive implementation.

---

## 9. Gradient Flow Analysis

### 9.1 Without Normalization

Deep network gradient flow:
```
L̇ˣ⁽⁰⁾ = W⁽¹⁾ᵀ · W⁽²⁾ᵀ · ... · W⁽ᴸ⁾ᵀ · L̇ʸ
```

If weights are slightly > 1: gradient explodes
If weights are slightly < 1: gradient vanishes

### 9.2 With Normalization

Each layer's activations are forced to unit variance.
Gradient magnitudes stabilize.

```
L̇ˣ̂ has unit variance (approximately)
→ L̇ˣ has controlled magnitude
→ No explosion or vanishing
```

### 9.3 The Jacobian View

LayerNorm Jacobian ∂Y/∂X is NOT diagonal (because of mean/var coupling).

But it has a special structure:
```
J = (1/σ) · (I - (1/D)·1·1ᵀ - (1/D)·X̂·X̂ᵀ)
```

This projects out the mean direction and decorrelates from X̂.
Eigenvalues are bounded, preventing gradient explosion.

---

## 10. Implementation Details

### 10.1 Memory Layout Matters

LayerNorm over last dimension (features):
```python
# Contiguous memory access pattern
for b in range(B):
    for d in range(D):  # Sequential access
        ...
```

BatchNorm over batch dimension:
```python
# Strided memory access pattern  
for d in range(D):
    for b in range(B):  # Jumping through memory
        ...
```

LayerNorm is more cache-friendly for typical tensor layouts [B, ..., D].

### 10.2 Epsilon Placement

```python
# Wrong (can still divide by zero if var=0):
x_hat = (x - mean) / sqrt(var + eps)

# Right (always safe):
x_hat = (x - mean) / (sqrt(var) + eps)

# Also right (fused):
x_hat = (x - mean) * rsqrt(var + eps)
```

The `rsqrt` (reciprocal square root) is a single GPU instruction.

---

## 11. Summary

### 11.1 Fluxion View of Normalization

**Forward:**
```
μ̇, σ̇ computed from X
X̂ = (X - μ)/σ       # Standardize
Y = γ·X̂ + β          # Scale and shift
```

**Backward:**
```
L̇ᵞ = sum(L̇ʸ · X̂)    # Scale gradient
L̇ᵝ = sum(L̇ʸ)         # Shift gradient
L̇ˣ = redistributed gradient (centered, decorrelated)
```

### 11.2 Key Insight

Normalization doesn't just scale activations—it COUPLES gradient flow across the normalized dimension.

Each input's gradient depends on ALL other inputs in the normalization group.

This coupling:
- Stabilizes gradient magnitudes
- Prevents single features from dominating
- Enables deeper networks

---

## References

1. Ioffe & Szegedy (2015). "Batch Normalization: Accelerating Deep Network Training."
2. Ba, Kiros & Hinton (2016). "Layer Normalization."
3. Zhang & Sennrich (2019). "Root Mean Square Layer Normalization."
4. Xiong et al. (2020). "On Layer Normalization in the Transformer Architecture."

---

*Correspondence: scott@opentransformers.online*