Pre-Norm vs Post-Norm › Learning Rate Sensitivity

def simulate_lr_sensitivity(norm_type, learning_rates, n_steps=200):
    """
    Simulate training stability across learning rates using loss dynamics.
    Returns a stability score: 1.0 = fully stable, 0.0 = diverged.
    Post-norm is sensitive to large LRs because the LN Jacobian gates all
    gradient paths; pre-norm's residual highway tolerates larger updates.
    """
    stability_scores = []

    # Post-norm instability threshold: large initial gradients destabilize
    # the LN Jacobian chain. The effective sensitivity scales with LR * n_layers.
    post_instability_threshold = (
        5e-3  # LRs above this cause post-norm instability
    )

    for lr in learning_rates:
        rng = np.random.default_rng(42)
        loss = 4.0
        losses = []
        diverged = False

        for step in range(n_steps):
            # Post-norm: large LR causes gradient explosion through LN Jacobian chain
            if norm_type == "post" and lr > post_instability_threshold:
                # Instability grows with LR; higher LR means faster divergence
                instability_factor = lr / post_instability_threshold
                noise = rng.standard_normal() * lr * 500 * instability_factor
                loss = loss + abs(noise) * 0.3
                if loss > 15.0:
                    diverged = True
                    break
            else:
                # Stable convergence for pre-norm (all LRs) and
                # post-norm at low-to-moderate LRs
                decay = 1.0 - lr * 5
                decay = max(0.990, min(0.999, decay))
                loss = max(0.3, loss * decay + rng.standard_normal() * 0.02)

            losses.append(loss)

        if diverged or len(losses) == 0:
            stability_scores.append(0.0)
        else:
            final_loss = (
                np.mean(losses[-20:]) if len(losses) >= 20 else np.mean(losses)
            )
            stability_scores.append(1.0 / (1.0 + final_loss))

    return stability_scores


learning_rates = [1e-5, 1e-4, 1e-3, 5e-3, 1e-2, 5e-2]
pre_stability = simulate_lr_sensitivity("pre", learning_rates)
post_stability = simulate_lr_sensitivity("post", learning_rates)