Training Deep Networks from Scratch

01Keeping the signal alive through depth

Each layer of a network multiplies its input by a weight matrix and applies a nonlinearity. Stack forty of them and the activations are the product of forty random factors: if each layer scales the signal by 0.7 on average the output is 10⁻⁶ of the input, and the gradients, which pass back through the same factors, vanish with it; at 1.4 per layer the output is 10⁶ times the input and the gradients explode.^[1] Vanishing and exploding gradients are the first thing depth breaks, and the initialization scale is the first lever.

So a linear layer preserves variance when n_inσ² = 1, and the gradient's variance is preserved when n_outσ² = 1. Xavier initialization splits the difference with σ² = 2/(n_in + n_out).^[1] A ReLU zeroes half of its inputs and halves the variance of what passes, which the Kaiming scheme compensates with a gain of √2: σ = √2/√n_in.^[2]

The scale is not the only lever. The sigmoid's derivative is at most 1/4 and is close to zero away from the origin, so a sigmoid stack shrinks gradients at every layer; ReLU's derivative is 1 wherever the unit is active, which is a large part of why it displaced the sigmoid.^[1] Initialization must also break symmetry: if every unit in a layer starts with the same weights it receives the same gradient forever, so the weights are drawn at random even when their scale is chosen with care.^[1]

02Normalization layers

Rather than hoping the initialization keeps activations at a usable scale for the whole run, a normalization layer forces it. Layer normalization computes the mean and variance over the feature dimension of each example and rescales to zero mean and unit variance, followed by a learned gain and bias.^[3] It is the normalization used in the original transformer and in GPT-2; the transformer page covers where it sits in the block.

RMSNorm drops the mean subtraction and the bias and keeps only the rescaling; Llama's reference implementation is this formula; its RMSNorm class defaults to ε = 10⁻⁶ and the model is built with norm_eps = 10⁻⁵.^[4] Both are per-example operations, so they behave the same in training and inference, unlike batch normalization, which uses statistics across the batch during training.^[5]

In the signal widget, switch any exploding or vanishing configuration to LayerNorm or RMSNorm: the variance sits at 1 after every layer regardless of the weight scale. Normalization does not remove the need for a sensible initialization, since the gradient scale and the pre-normalization values still depend on it, but it removes the exponential dependence on depth that makes forty layers untrainable without it.

03Momentum and Adam

The linear models page showed plain gradient descent crawling along the flat direction of an ill-conditioned bowl because the steep direction caps the learning rate. Momentum replaces the gradient with a leaky average of past gradients: aligned components add up to an effective step of η/(1 − β), oscillating ones cancel, and the stable range of ηλ widens from 2 to 2 + 2β.^[6]

Adam goes further and normalizes each parameter's step by a running estimate of its own gradient magnitude, so a coordinate with gradients a thousand times larger does not get steps a thousand times larger.^[7] The defaults β₁ = 0.9 and β₂ = 0.999 make the second moment move much more slowly than the first.

AdamW applies weight decay as a separate shrink of θ before the moment update, so the decay does not pass through the moment averages.^[8] That is the variant nanoGPT and many other language-model trainers use.^[9]

The bias correction matters more than it looks. Both moving averages start at zero, so for the first few steps m is a fraction of the true gradient and v is a smaller fraction of its square. The ratio m/√v is therefore biased upward, not downward: without correction the first Adam step with the default betas is about three times η, and it grows to about 6.6η by step 12 because m warms up faster than v. The overshoot is still about 1.26η at step 1,000.^[7]

04Warmup, decay and clipping

A fixed learning rate is wrong at both ends of a run. At the start the weights are random and the gradient directions are poor, so large steps can diverge; a warmup ramps η up over the first steps and limits early divergence in deep networks.^[10] At the end, minibatch noise keeps the iterate jittering at a scale set by η, so the rate must decay for the loss to settle; a cosine decay, proposed by Loshchilov and Hutter, holds η high for a while and then anneals it smoothly.^[10] nanoGPT's schedule is exactly these two pieces: linear warmup, then cosine from the peak to a floor of a tenth of it.^[9]

Gradient clipping is the safety rail for the rare step whose gradient is enormous. PyTorch's clip_grad_norm_ computes the norm over all parameters' gradients as if concatenated into one vector and scales them down together when it exceeds the maximum, so the update direction is kept and only its length is capped; nanoGPT clips at 1.0 on every step.^[11][9]

05Mixed precision and loss scaling

Tensor hardware runs 16-bit matrix multiplies several times faster than 32-bit ones and moves half the bytes. Mixed precision training runs the heavy operations in a 16-bit format under an autocast rule and keeps a 32-bit master copy of the parameters for the optimizer update.^[12] Which 16-bit format decides what can go wrong.

float16 has 10 fraction bits and 5 exponent bits: its smallest normal number is 2⁻¹⁴ ≈ 6.1×10⁻⁵, its smallest subnormal 2⁻²⁴ ≈ 6×10⁻⁸, and its largest finite value 65,504. Small gradients can fall below the subnormal floor and flush to zero, so the update for those parameters is lost. The fix is loss scaling: multiply the loss by a scale before the backward pass, which multiplies every gradient by the same factor, then divide the gradients by it before the optimizer step.^[12] PyTorch's GradScaler picks the scale dynamically: it starts at 2¹⁶, skips the step and halves the scale when an inf or NaN appears, and doubles it after 2,000 consecutive clean steps.^[13] bfloat16 keeps float32's 8 exponent bits with only 7 fraction bits, so it has the same range as float32, needs no scaling, and pays in resolution; it is the format nanoGPT uses when the hardware supports it, falling back to float16 with a GradScaler otherwise.^[9] The bit layouts and rounding rules are on the floating point page.

06Scaling laws and the compute budget

Training cost for a transformer is close to 6ND floating-point operations: 2ND for the forward pass over D tokens with N parameters, and twice that for the backward.^[15] A compute budget C therefore fixes the product ND and leaves one question: a larger model on fewer tokens, or a smaller one on more. Scaling laws answer it empirically. Hoffmann et al. fit the final loss of hundreds of training runs to a function of N and D and found that, at a fixed budget, model size and token count should grow together, which implied that models of the time were under-trained for their size.^[14]

The constants are the paper's third-approach fit as reproduced in nanoGPT's scaling-law notebook, which notes that it cannot make the fit reproduce the paper's tabulated optima exactly.^[15] The widget treats the formula as what it is: a fit whose shape is the lesson, with the paper's own second-approach table drawn alongside for comparison.

07AdamW, schedule and clipping, checked

These complete programs implement AdamW with bias correction, global-norm gradient clipping, and the warmup-plus-cosine schedule, then check them: the first Adam step has magnitude exactly η per coordinate regardless of the gradient's scale, clipping preserves direction and caps the norm, the schedule hits its documented values at the boundaries, and AdamW drives the condition-50 quadratic below 10⁻⁸.

#include <cassert>
#include <cmath>
#include <cstddef>
#include <numbers>
#include <vector>
using Vector = std::vector<double>;
struct AdamW {
    double beta1, beta2, epsilon, weightDecay;
    Vector firstMoment, secondMoment;
    int step = 0;
    AdamW(std::size_t size, double decay)
        : beta1(0.9), beta2(0.999), epsilon(1e-8), weightDecay(decay), firstMoment(size, 0), secondMoment(size, 0) {}
    void update(Vector& parameters, const Vector& gradient, double scheduledRate) {
        ++step;
        const double correction1 = 1 - std::pow(beta1, step), correction2 = 1 - std::pow(beta2, step);
        for (std::size_t index = 0; index < parameters.size(); ++index) {
            parameters[index] -= scheduledRate * weightDecay * parameters[index]; // Decoupled decay, applied to θ directly.
            firstMoment[index] = beta1 * firstMoment[index] + (1 - beta1) * gradient[index];
            secondMoment[index] = beta2 * secondMoment[index] + (1 - beta2) * gradient[index] * gradient[index];
            const double correctedFirst = firstMoment[index] / correction1, correctedSecond = secondMoment[index] / correction2;
            parameters[index] -= scheduledRate * correctedFirst / (std::sqrt(correctedSecond) + epsilon);
        }
    }
};
// Scale the whole gradient so its global ℓ2 norm is at most maxNorm; returns the norm before clipping.
double clipGlobalNorm(Vector& gradient, double maxNorm) {
    double sumSquares = 0;
    for (double value : gradient) sumSquares += value * value;
    const double norm = std::sqrt(sumSquares);
    if (norm > maxNorm) for (double& value : gradient) value *= maxNorm / norm;
    return norm;
}
// Linear warmup to peak, cosine decay to floor, constant floor afterwards.
double warmupCosine(int iteration, int warmup, int total, double peak, double floor) {
    if (iteration < warmup) return peak * (iteration + 1) / (warmup + 1);
    if (iteration > total) return floor;
    const double ratio = double(iteration - warmup) / (total - warmup);
    return floor + 0.5 * (1 + std::cos(std::numbers::pi * ratio)) * (peak - floor);
}
int main() {
    Vector parameters{0.5, -0.5};
    AdamW firstStep(2, 0.0);
    firstStep.update(parameters, Vector{1000.0, -0.001}, 0.01);
    assert(std::abs(parameters[0] - 0.49) < 1e-6 && std::abs(parameters[1] + 0.49) < 1e-6); // |step| = η per coordinate, whatever the gradient scale (up to ε).
    Vector gradient{3.0, 4.0};
    const double before = clipGlobalNorm(gradient, 1.0);
    assert(std::abs(before - 5.0) < 1e-12 && std::abs(gradient[0] - 0.6) < 1e-12 && std::abs(gradient[1] - 0.8) < 1e-12); // Direction kept, norm capped at 1.
    assert(std::abs(warmupCosine(0, 2000, 20000, 6e-4, 6e-5) - 6e-4 / 2001) < 1e-15);
    assert(std::abs(warmupCosine(2000, 2000, 20000, 6e-4, 6e-5) - 6e-4) < 1e-15);   // Peak at the end of warmup.
    assert(std::abs(warmupCosine(20000, 2000, 20000, 6e-4, 6e-5) - 6e-5) < 1e-15);  // Floor at the end.
    Vector point{-2.5, 2.5};
    AdamW optimizer(2, 0.0);
    for (int iteration = 0; iteration < 2000; ++iteration) {
        Vector bowlGradient{point[0], 50.0 * point[1]}; // ∇ of ½(x² + 50 y²): condition number 50.
        clipGlobalNorm(bowlGradient, 100.0);
        optimizer.update(point, bowlGradient, warmupCosine(iteration, 50, 2000, 0.1, 0.001));
    }
    const double loss = 0.5 * (point[0] * point[0] + 50.0 * point[1] * point[1]);
    assert(loss < 1e-8);
}

struct AdamW { beta1: f64, beta2: f64, epsilon: f64, weight_decay: f64, first_moment: Vec<f64>, second_moment: Vec<f64>, step: i32 }
impl AdamW {
    fn new(size: usize, weight_decay: f64) -> AdamW {
        AdamW { beta1: 0.9, beta2: 0.999, epsilon: 1e-8, weight_decay, first_moment: vec![0.0; size], second_moment: vec![0.0; size], step: 0 }
    }
    fn update(&mut self, parameters: &mut [f64], gradient: &[f64], scheduled_rate: f64) {
        self.step += 1;
        let correction1 = 1.0 - self.beta1.powi(self.step);
        let correction2 = 1.0 - self.beta2.powi(self.step);
        for index in 0..parameters.len() {
            parameters[index] -= scheduled_rate * self.weight_decay * parameters[index]; // Decoupled decay, applied to θ directly.
            self.first_moment[index] = self.beta1 * self.first_moment[index] + (1.0 - self.beta1) * gradient[index];
            self.second_moment[index] = self.beta2 * self.second_moment[index] + (1.0 - self.beta2) * gradient[index] * gradient[index];
            let corrected_first = self.first_moment[index] / correction1;
            let corrected_second = self.second_moment[index] / correction2;
            parameters[index] -= scheduled_rate * corrected_first / (corrected_second.sqrt() + self.epsilon);
        }
    }
}
// Scale the whole gradient so its global ℓ2 norm is at most max_norm; returns the norm before clipping.
fn clip_global_norm(gradient: &mut [f64], max_norm: f64) -> f64 {
    let norm = gradient.iter().map(|value| value * value).sum::<f64>().sqrt();
    if norm > max_norm { for value in gradient.iter_mut() { *value *= max_norm / norm; } }
    norm
}
// Linear warmup to peak, cosine decay to floor, constant floor afterwards.
fn warmup_cosine(iteration: i32, warmup: i32, total: i32, peak: f64, floor: f64) -> f64 {
    if iteration < warmup { return peak * (iteration + 1) as f64 / (warmup + 1) as f64; }
    if iteration > total { return floor; }
    let ratio = (iteration - warmup) as f64 / (total - warmup) as f64;
    floor + 0.5 * (1.0 + (std::f64::consts::PI * ratio).cos()) * (peak - floor)
}
fn main() {
    let mut parameters = [0.5, -0.5];
    let mut first_step = AdamW::new(2, 0.0);
    first_step.update(&mut parameters, &[1000.0, -0.001], 0.01);
    assert!((parameters[0] - 0.49).abs() < 1e-6 && (parameters[1] + 0.49).abs() < 1e-6); // |step| = η per coordinate, whatever the gradient scale (up to ε).
    let mut gradient = [3.0, 4.0];
    let before = clip_global_norm(&mut gradient, 1.0);
    assert!((before - 5.0).abs() < 1e-12 && (gradient[0] - 0.6).abs() < 1e-12 && (gradient[1] - 0.8).abs() < 1e-12); // Direction kept, norm capped at 1.
    assert!((warmup_cosine(0, 2000, 20000, 6e-4, 6e-5) - 6e-4 / 2001.0).abs() < 1e-15);
    assert!((warmup_cosine(2000, 2000, 20000, 6e-4, 6e-5) - 6e-4).abs() < 1e-15);   // Peak at the end of warmup.
    assert!((warmup_cosine(20000, 2000, 20000, 6e-4, 6e-5) - 6e-5).abs() < 1e-15);  // Floor at the end.
    let mut point = [-2.5, 2.5];
    let mut optimizer = AdamW::new(2, 0.0);
    for iteration in 0..2000 {
        let mut bowl_gradient = [point[0], 50.0 * point[1]]; // ∇ of ½(x² + 50 y²): condition number 50.
        clip_global_norm(&mut bowl_gradient, 100.0);
        optimizer.update(&mut point, &bowl_gradient, warmup_cosine(iteration, 50, 2000, 0.1, 0.001));
    }
    let loss = 0.5 * (point[0] * point[0] + 50.0 * point[1] * point[1]);
    assert!(loss < 1e-8);
}

08Where training runs go wrong

Weight decay in AdamW is applied directly to the parameters, so decaying biases and normalization gains pulls them toward zero, which is rarely wanted. nanoGPT puts every parameter with fewer than two dimensions, biases and LayerNorm gains, in a separate group with decay disabled.^[9]

Normalization hides initialization problems rather than solving them. A model that trains only because LayerNorm is rescuing every layer has gradients whose scale still depends on the pre-normalization values, and it may train slowly or need an unusually small learning rate; check the activation statistics at initialization before and after the normalization, as the §1 widget does.

A scaling law fit is a description of one family of runs on one dataset with one recipe. Its exponents transfer better than its constants, and its optimum shifts if inference cost is part of the budget, since a smaller model trained on more tokens is cheaper to serve. Use the shape of the curve to reason; measure the level with runs of your own.

09What's next

A run that trains is not yet a model that works. Evaluation and Generalization covers how to tell: held-out splits and what leaks between them, the gap between training and validation loss, metrics for classification and for language models, calibration, and why a benchmark score has an error bar whether or not it is printed.

10Sources

The cited texts define the methods; the framework sources document how shipped trainers implement them. Widget readouts report values computed in the page from seeded simulations and the stated formulas.

1. Aston Zhang, Zachary C. Lipton, Mu Li, Alexander J. Smola, 2023. Dive into Deep Learning, Section 5.4: Numerical Stability and Initialization, Cambridge University Press. Vanishing and exploding gradients, the sigmoid's role, the product of random matrices, symmetry breaking, and the Xavier variance derivation.
2. PyTorch contributors. torch.nn.init. Kaiming initialization with std = gain/√fan and the gains √2 for ReLU and 5/3 for tanh, citing He et al. (2015) and Glorot and Bengio (2010).
3. PyTorch contributors. torch.nn.LayerNorm. Mean and variance over the normalized dimensions of each input, with learned affine parameters.
4. Meta Llama team. llama/model.py. The RMSNorm reference implementation: x · rsqrt(mean(x²) + ε) · weight, with a class default ε = 10⁻⁶ and the model's norm_eps = 10⁻⁵.
5. Stanford CS231n course staff. Setting up the data and the model. Batch normalization as normalization using batch statistics, inserted before nonlinearities.
6. Aston Zhang et al., 2023. Dive into Deep Learning, Section 12.6: Momentum. The velocity recursion, the effective step η/(1 − β), and the stability range ηλ < 2 + 2β.
7. Aston Zhang et al., 2023. Dive into Deep Learning, Section 12.10: Adam. The moment updates, the defaults β₁ = 0.9 and β₂ = 0.999, the initial bias toward zero, and the bias-correction factors.
8. PyTorch contributors. torch.optim.AdamW. The algorithm with weight decay applied to the parameters before the moment update, outside the moment averages.
9. Andrej Karpathy. nanoGPT train.py. Linear warmup then cosine decay to a floor of a tenth of the peak, AdamW with β₂ = 0.95 and weight decay 0.1, gradient clipping at 1.0, and bfloat16 autocast with a float16 GradScaler fallback. Its model.py configure_optimizers exempts biases and LayerNorm gains from weight decay.
10. Aston Zhang et al., 2023. Dive into Deep Learning, Section 12.11: Learning Rate Scheduling. The cosine schedule attributed to Loshchilov and Hutter (2016) and warmup limiting early divergence, citing Gotmare et al. (2018).
11. PyTorch contributors. torch.nn.utils.clip_grad_norm_. The norm computed over all parameter gradients as one concatenated vector.
12. PyTorch contributors. Automatic Mixed Precision package. Autocast with float16 or bfloat16, gradient underflow in float16, loss scaling and unscaling before the optimizer step, and the float16 maximum of 65,504.
13. PyTorch contributors. torch.amp.GradScaler. The dynamic scale: initial 2¹⁶, step skipped and scale reduced on inf or NaN, scale doubled after 2,000 consecutive clean iterations.
14. Jordan Hoffmann et al., 2022. Training Compute-Optimal Large Language Models. The parametric loss fit and the finding that parameters and tokens should scale together; constants taken as reproduced in source 15.
15. Andrej Karpathy. nanoGPT scaling_laws.ipynb. The Approach 3 constants E = 1.69, A = 406.4, B = 410.7, α = 0.34, β = 0.28, the Approach 2 table, the 6ND compute approximation with backward counted as twice the forward, and the note that the fit does not reproduce the paper's tabulated optima exactly.