Regularization

Deep dive

How it works

In: weights, gradients or activations, plus λ or p · Out: shrunk weights, masked activations or a stop signal · State: dropout mask until backward; best validation score and step for early stopping; none for weight decay · Cost: O(P) per step for decay, O(activations) for dropout

1. L2 vs decoupled decay. In SGD, w ← w − lr(g + λw) is exactly decay with rate lr·λ. In Adam, the λw term is divided by √v, so weights with large gradients barely decay. AdamW applies the decay outside the adaptive normalization.
2. Weight decay with normalization layers. A weight feeding a norm layer is scale-invariant, so decay only shrinks its norm. That raises the effective learning rate (lr/‖w‖²). Decay then acts as a learning-rate controller more than a capacity limit.
3. Dropout samples a sub-network per step. Scaling at train time (inverted) keeps the expected activation equal, so inference needs no change. Backward multiplies the gradient by the same mask and scale.
4. Early stopping for GD on a quadratic loss behaves like L2 with λ ≈ 1/(lr·steps): fewer steps, stronger regularization.

Code

/// xorshift64*: cheap RNG for masks.
pub struct Rng(pub u64);
impl Rng {
    pub fn next_u64(&mut self) -> u64 {
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        self.0.wrapping_mul(0x2545F4914F6CDD1D)
    }
    pub fn uniform(&mut self) -> f32 { (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32 }
}

/// Inverted dropout. Stores a 1-bit-per-element mask for backward (32x smaller than f32).
pub fn dropout_fwd(x: &mut [f32], p: f32, rng: &mut Rng, mask: &mut Vec<u64>) {
    let scale = 1.0 / (1.0 - p);
    let thresh = (p * (1u64 << 24) as f32) as u64;      // keep iff 24-bit uniform >= thresh
    mask.clear();
    mask.resize((x.len() + 63) / 64, 0);
    for (i, v) in x.iter_mut().enumerate() {
        let keep = (rng.next_u64() >> 40) >= thresh;
        mask[i / 64] |= (keep as u64) << (i % 64);
        *v = if keep { *v * scale } else { 0.0 };
    }
}

pub fn dropout_bwd(dy: &mut [f32], p: f32, mask: &[u64]) {
    let scale = 1.0 / (1.0 - p);
    for (i, g) in dy.iter_mut().enumerate() {
        *g = if mask[i / 64] >> (i % 64) & 1 == 1 { *g * scale } else { 0.0 };
    }
}

/// SGD with L2 folded into the gradient (identical to decay at rate lr·λ for plain SGD).
pub fn sgd_l2(w: &mut [f32], g: &[f32], lr: f32, l2: f32) {
    for (wi, gi) in w.iter_mut().zip(g) { *wi -= lr * (gi + l2 * *wi); }
}

/// Decoupled decay: shrink first, then any optimizer step (here plain SGD).
pub fn sgd_decoupled(w: &mut [f32], g: &[f32], lr: f32, wd: f32) {
    for (wi, gi) in w.iter_mut().zip(g) { *wi = *wi * (1.0 - lr * wd) - lr * gi; }
}

/// Proximal L1 (soft threshold) after a gradient step: produces exact zeros.
pub fn prox_l1(w: &mut [f32], lr: f32, l1: f32) {
    for wi in w.iter_mut() { *wi = wi.signum() * (wi.abs() - lr * l1).max(0.0); }
}

/// Mixup two examples and their one-hot or soft targets with weight lam.
pub fn mixup(xa: &mut [f32], xb: &[f32], ya: &mut [f32], yb: &[f32], lam: f32) {
    for (a, b) in xa.iter_mut().zip(xb) { *a = lam * *a + (1.0 - lam) * b; }
    for (a, b) in ya.iter_mut().zip(yb) { *a = lam * *a + (1.0 - lam) * b; }
}

pub struct EarlyStop { pub best: f32, pub best_step: usize, pub patience: usize, bad: usize }
impl EarlyStop {
    pub fn new(patience: usize) -> Self { Self { best: f32::INFINITY, best_step: 0, patience, bad: 0 } }
    /// Returns (save_checkpoint, stop).
    pub fn check(&mut self, step: usize, val_loss: f32, min_delta: f32) -> (bool, bool) {
        if val_loss < self.best - min_delta {
            self.best = val_loss; self.best_step = step; self.bad = 0;
            (true, false)
        } else { self.bad += 1; (false, self.bad >= self.patience) }
    }
}

fn main() {
    let mut rng = Rng(42);
    let mut x = vec![1.0f32; 100_000];
    let mut mask = Vec::new();
    dropout_fwd(&mut x, 0.1, &mut rng, &mut mask);
    let mean = x.iter().sum::<f32>() / x.len() as f32;
    println!("dropout p=0.1: mean activation {mean:.3} (expected 1.0)");
    let mut g = vec![1.0f32; x.len()];
    dropout_bwd(&mut g, 0.1, &mask);
    assert!(x.iter().zip(&g).all(|(a, b)| a == b));        // same mask, same scale
    let (mut a, mut b) = (vec![2.0f32], vec![2.0f32]);
    sgd_l2(&mut a, &[0.5], 0.1, 0.01);
    sgd_decoupled(&mut b, &[0.5], 0.1, 0.01);
    assert!((a[0] - b[0]).abs() < 1e-6);                    // identical for plain SGD
    let mut w = vec![0.03f32, -0.5, 0.004];
    prox_l1(&mut w, 0.1, 0.05);
    println!("prox L1: {w:?}");                             // small weights become exactly 0
    let mut es = EarlyStop::new(3);
    for (s, l) in [1.0f32, 0.8, 0.82, 0.79, 0.81, 0.80, 0.83].iter().enumerate() {
        let (save, stop) = es.check(s, *l, 1e-3);
        if save { println!("step {s}: new best {l}"); }
        if stop { println!("stop at step {s}, restore step {}", es.best_step); break; }
    }
}

Complexity and constants

• Weight decay: one multiply per parameter, fused into the optimizer kernel. Free.
• Dropout: one random number per activation. Storing the mask as bits costs N/8 bytes, vs 4N for an fp32 mask. Fused kernels (and FlashAttention's attention dropout) regenerate the mask from the RNG seed and offset in backward and store nothing.
• Early stopping: a validation pass every k steps, order of 1-5 % of training compute.
• Mixup and label smoothing: O(batch × features), negligible.

Efficiency

• Cost: weight decay is one multiply per parameter inside the optimizer kernel, effectively free. Dropout needs one random number per activation and a mask of N/8 bytes as bits (4N as fp32). Early stopping costs a validation pass every k steps, order of 1-5 % of training compute. Bound: memory bandwidth for mask traffic.
• Levers:
  • decoupled weight decay in the fused optimizer — no extra pass (Loshchilov and Hutter 2019) — always
  • regenerate the dropout mask from the RNG seed and offset in backward — stores no mask (above, FlashAttention) — attention and large activations
  • stochastic depth — skipped layers cost no compute, so training time falls (Huang et al. 2016) — deep residual nets
  • SWA or EMA of weights — one extra weight copy, no extra training steps (Izmailov et al. 2018) — late-stage training
  • mixup, label smoothing — O(batch × features), negligible (above) — small data
• Naive → standard → state of the art:
  • fp32 dropout mask stored per layer, separate elementwise kernels · 4N bytes per mask
  • bit-packed mask, fused bias + dropout + residual · N/8 bytes
  • mask regenerated from seed in backward, fused into the surrounding kernel · 0 stored bytes

Implementation notes

• Build two optimizer parameter groups: decay for weight matrices and conv kernels, no decay for 1-D parameters (biases, norm scales).
• PyTorch's SGD(weight_decay=λ) adds λw to the gradient (L2), and Adam(weight_decay=λ) does the same. Only AdamW decouples. PyTorch AdamW multiplies λ by lr: λ = 0.1 at lr 3e-4 shrinks by 3e-5 per step.
• Dropout placement in a transformer: after attention softmax, after each sublayer output before the residual add, and on embeddings. Turn all of them off at inference (model.eval()).
• Keep the RNG state in the checkpoint if you need bit-exact resume.
• Early stopping with noisy validation: smooth (EMA) the metric or require min_delta, and keep the best checkpoint rather than the last.
• In GBDTs, the equivalents are shrinkage (learning rate), λ on leaf weights, row and column subsampling, and early stopping on validation rounds.

Pitfalls

• Adam(weight_decay=…) believed to be AdamW: decay is then scaled per parameter by 1/√v and is much weaker on frequently-updated weights.
• Decaying LayerNorm γ toward 0: it slowly shrinks the signal and fights the optimizer.
• Dropout left on at evaluation, or BatchNorm and dropout interacting (variance shift): dropout before BN changes the variance BN sees between train and test.
• Tuning weight decay without retuning lr: with AdamW the effective decay per step is lr·λ, so a new lr changes regularization too.
• Early stopping on the test set: the test score is then optimistic. Use a separate validation split.

Variants and alternatives

In production

torch.optim.AdamW, torch.nn.Dropout, timm (drop_path_rate for stochastic depth, Mixup/CutMix, RandAugment), scikit-learn Ridge/Lasso/ElasticNet (alpha), XGBoost lambda/alpha/subsample, Keras EarlyStopping(patience=…, restore_best_weights=True).

Implementations

Sources

• Srivastava, Hinton, Krizhevsky, Sutskever, Salakhutdinov, Dropout: A Simple Way to Prevent Neural Networks from Overfitting, JMLR 2014. paper
• Loshchilov and Hutter, Decoupled Weight Decay Regularization, ICLR 2019. paper
• Huang, Sun, Liu, Sedra, Weinberger, Deep Networks with Stochastic Depth, ECCV 2016. paper
• Zhang, Cisse, Dauphin, Lopez-Paz, mixup: Beyond Empirical Risk Minimization, ICLR 2018. paper
• Yun et al., CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features, ICCV 2019. paper
• He, Zhang, Ren, Sun, Deep Residual Learning for Image Recognition, CVPR 2016. paper
• Vaswani et al., Attention Is All You Need, NeurIPS 2017. paper
• Izmailov et al., Averaging Weights Leads to Wider Optima and Better Generalization, UAI 2018. paper