Weight Initialization, Vanishing Gradients and Optimizers: Why Deep Nets Train

You build a ten layer network, press train, and the loss does not move. Not slowly. Not noisily. It simply sits there. Nothing is wrong with your data or your code. The problem was decided before the first update, by the numbers you chose as starting weights. Today we find out exactly why.

• Loss stuck from step one
• Cause: the starting weights

In lesson 46 we forecast time series with ARIMA and exponential smoothing. Now we return to neural networks, building on lesson 26. We will derive Xavier and He initialization, see gradients vanish and explode in experiments, then compare SGD, momentum, Nesterov, RMSProp and Adam, and finish with learning rate schedules. Dropout and batch normalization wait for the next lesson.

• Initialization from variance
• Vanishing and exploding gradients
• Optimizers and schedules

Why not start every weight at zero? Because then every neuron in a layer computes the same thing and receives the same gradient. They update identically forever, so a hundred neurons behave like one. We must break this symmetry with random values. But random how big? That question is the whole lesson.

• Zero weights give identical neurons
• Symmetry never breaks

Take a layer with one hundred inputs. Each output is a sum of one hundred terms, each a weight times an input. If the weights are tiny, the sum is tiny, and the signal fades layer by layer. If the weights have standard deviation one, the sum is much larger than the inputs. Our goal is a size that keeps the signal steady.

• Std 0.01 shrinks signals
• Std 1 blows them up

Consider one neuron: z equals the sum over n inputs of w sub i times x sub i. Assume weights and inputs are independent, with zero mean. Then the variance of each product is the variance of w times the variance of x. Because the terms are independent, variances add. So the variance of z equals n times variance of w times variance of x.

• z equals sum of w times x
• Assume zero mean, independent

Look at that factor, n times the variance of w. Each layer multiplies the signal variance by it. With n equal to one hundred and weight standard deviation one, the factor is one hundred. After ten layers the variance grows by ten to the twenty. With standard deviation point zero one, the factor is one hundredth, and ten layers shrink it by ten to the minus twenty.

• Each layer multiplies variance
• Factor must equal one

To keep the factor at one going forward, set the variance of w to one over n in. Going backward, gradients flow through the transposed weights, which asks for one over n out. When these differ, Glorot and Bengio proposed the compromise: variance equals two divided by n in plus n out. That is Xavier, or Glorot, initialization. The uniform version draws from plus or minus the square root of six over n in plus n out.

• Forward: Var(w) = 1/n_in
• Backward: Var(w) = 1/n_out
• Compromise: 2/(n_in+n_out)

A concrete case. A layer takes 256 inputs and produces 64 outputs. Variance is two divided by 320, which is 0.00625. The standard deviation is the square root of that, about 0.079. Compare with a careless choice of 1: that would inflate the signal by a factor of 256 in a single layer.

• n_in 256, n_out 64
• Var = 2/320 = 0.00625
• Std about 0.079

The Xavier derivation assumed an activation that behaves roughly linearly near zero, like tanh. ReLU is different: it zeroes every negative input. For a symmetric pre-activation, half the values vanish, so the average squared output is half of the variance of z. Each ReLU layer therefore cuts the signal power by two on top of everything else.

• ReLU zeroes half the values
• Second moment halves

Redo the algebra. The variance of the next pre-activation is n times variance of w times the average squared activation. That squared activation is half the previous variance. So the factor becomes n times variance of w times one half. Setting it to one gives variance of w equals two over n in. This is He initialization, from Kaiming He and colleagues. For n in of 256, the standard deviation is about 0.088.

• Var(z_next) = n Var(w) E[a^2]
• E[a^2] = Var(z)/2
• So Var(w) = 2/n_in

Forward signals are only half the story. During backpropagation, from lesson 26, the gradient at an early layer is a product of one factor per layer: a weight matrix times an activation derivative. A long product of numbers smaller than one collapses to zero. A long product of numbers larger than one explodes. Depth turns small biases into disasters.

• Gradient multiplies per layer
• Many factors below one vanish

Sigmoid has a derivative that peaks at 0.25, at input zero, and is smaller everywhere else. Even in the best case, ten sigmoid layers multiply the gradient by at most 0.25 to the tenth power, about one millionth. Layer one receives a gradient a million times weaker than the last layer. It effectively stops learning.

• Sigmoid derivative max is 0.25
• 0.25 to the 10th is about 1e-6

Here is the experiment. Build a twenty layer network, one hundred units per layer. Measure the gradient norm at every layer after one backward pass. With sigmoid and Xavier weights, the norms fall exponentially toward the input. With ReLU and He initialization, they stay within the same order of magnitude across all twenty layers. The derivative of ReLU is exactly one for active units, so it does not shrink the product.

• 20 layers, 100 units each
• Sigmoid: tiny gradients early
• ReLU with He: stable

The opposite failure is explosion. Suppose each layer scales the gradient by 1.5. After fifty layers the gradient is 1.5 to the fiftieth, about 640 million times larger. One update then throws the weights far away, the loss jumps, and soon you see not a number in your logs. Recurrent networks are especially prone, since the same matrix repeats at every time step.

• Factor 1.5 per layer
• 1.5^50 is about 6.4e8
• Loss becomes NaN

A practical remedy for explosion is gradient clipping. Compute the global norm of the gradient. If it exceeds a threshold c, rescale the gradient by c divided by its norm. Example: the gradient is thirty, forty, with norm fifty, and c is five. Multiply by one tenth and get three, four. The direction is unchanged; only the step length is capped.

• If norm exceeds threshold c
• g becomes g times c over norm
• Direction is preserved

Two mistakes to avoid. First, clipping cannot cure vanishing gradients, since you cannot rescale a signal that is already lost; you fix those with better activations and initialization. Second, clipping each component separately changes the direction of the gradient, while clipping by global norm keeps it. Also remember biases are usually initialized to zero, which is safe because symmetry is already broken by the weights.

• Clipping does not fix vanishing
• Clip by norm, not each value blindly

Now optimizers. Take the bowl f of x and y equals one half of x squared plus ten y squared. Its contours are narrow ellipses, a steep valley. The gradient is x, ten y. Plain gradient descent subtracts the learning rate times the gradient. With a learning rate of 0.19, the y coordinate flips sign each step, since one minus 0.19 times ten is minus 0.9, while x creeps along slowly.

• f = 0.5(x² + 10y²)
• Gradient is (x, 10y)
• SGD zigzags

Momentum treats the parameters as a heavy ball rolling downhill. Keep a velocity v that is an exponential average of gradients: v equals beta times the old v plus one minus beta times the new gradient. Then step along v. Zigzag components alternate in sign and cancel in the average, while the consistent downhill component accumulates.

• v = beta v + (1-beta) g
• theta = theta - lr v

Why exponential? Unroll the recursion. The velocity at step t equals one minus beta times the sum over k of beta to the k times the gradient from k steps ago. Recent gradients get the largest weight, older ones decay geometrically. With beta of 0.9, the effective memory is about 1 over 0.1, ten steps. Many libraries drop the one minus beta factor, which only rescales the learning rate.

• Unroll: v_t = (1-β) Σ β^k g_{t-k}
• Weights sum to 1 - β^t
• Window about 1/(1-β)

Nesterov momentum makes one change: evaluate the gradient not where you are, but where momentum is about to carry you, at theta plus beta times v. If the ball is about to overshoot, the look ahead gradient already points back, so the correction happens one step sooner. In practice this often damps oscillation a little better than classical momentum.

• Compute gradient at theta + beta v
• Corrects overshoot earlier

In our valley the y gradient is ten times larger than the x gradient, yet one learning rate serves both. RMSProp keeps a running average s of the squared gradient per parameter, then divides each step by its square root. A coordinate with gradients near ten and one near 0.1 both get steps of similar size, about the learning rate. Steep directions are tamed and flat ones are boosted.

• s = ρ s + (1-ρ) g²
• step = lr · g / (√s + ε)

Adam combines both ideas. It keeps m, an exponential average of gradients, like momentum, and v, an exponential average of squared gradients, like RMSProp. The usual defaults are beta one 0.9, beta two 0.999 and epsilon ten to the minus eight. The update is the learning rate times m divided by the square root of v, plus epsilon.

• m: average of g
• v: average of g²
• β1 0.9, β2 0.999

There is a catch. Both m and v start at zero, so early on they are biased toward zero. After t steps the weights in the average sum to one minus beta to the t, not one. The fix is to divide: m hat equals m over one minus beta one to the t, and v hat equals v over one minus beta two to the t.

• m and v start at zero
• After t steps, weights sum to 1 - β^t
• Divide by 1 - β^t

Let us run Adam by hand. Start with theta equal to 1, learning rate 0.1, and gradients of 0.5 then 0.3. Step one: m is 0.1 times 0.5, which is 0.05. v is 0.001 times 0.25, which is 0.00025. Correct: m hat is 0.05 over 0.1, equal to 0.5; v hat is 0.00025 over 0.001, equal to 0.25. The square root is 0.5, so the step is 0.1 times 0.5 over 0.5, exactly 0.1. Theta becomes 0.9.

• θ = 1.0, g₁ = 0.5, lr = 0.1
• m = 0.05, v = 0.00025
• m̂ = 0.5, v̂ = 0.25

Step two with gradient 0.3. m is 0.9 times 0.05 plus 0.1 times 0.3, which is 0.075. v is 0.999 times 0.00025 plus 0.001 times 0.09, about 0.00033975. The corrections: divide m by 0.19 to get 0.3947, and v by 0.001999 to get about 0.16996, whose square root is 0.4123. The step is 0.1 times 0.3947 over 0.4123, about 0.0958. Theta becomes roughly 0.8042.

• g₂ = 0.3
• m = 0.075, v ≈ 0.00033975
• θ ≈ 0.8042

Notice what correction did. Without it, step one would use m of 0.05 over the square root of 0.00025, a ratio of about 3.16, giving a step of 0.316, over three times too large. With correction, the first step has size exactly the learning rate, whatever the gradient scale. That is a key property of Adam: step sizes are roughly bounded by the learning rate.

• Uncorrected step one: ratio 3.16
• Corrected step one: ratio 1.0

Even with Adam, one constant learning rate is rarely best. Early in training, large steps cover ground quickly. Late in training, large steps keep bouncing around the minimum, so we shrink them. A schedule is simply a rule that sets the learning rate as a function of the step number.

• Big steps early
• Small steps near the end

Step decay multiplies the rate by a factor like 0.1 at chosen milestones. Cosine decay is smoother: learning rate equals the minimum plus one half of maximum minus minimum, times one plus cosine of pi times t over T. With maximum 0.1, minimum zero and T equal to 100, at step 50 cosine of pi over two is zero, so the rate is exactly 0.05, halfway down.

• Step: multiply by 0.1 at milestones
• Cosine: lr_min + ½(max-min)(1+cos(πt/T))

Warmup starts with a tiny rate and increases it linearly over the first W steps. With W equal to 1000, step 250 uses a quarter of the maximum rate. Early on, the weights are random and Adam's second moment estimate rests on very few gradients, so big steps can be unstable. Warmup is standard in transformer training, usually followed by cosine or linear decay.

• lr = lr_max · t / W for t < W
• W = 1000, t = 250 gives 25%

All of these fit in a few lines of NumPy. SGD needs nothing. Momentum stores v. RMSProp stores s. Adam stores m, v and a step counter t, then applies the two corrections and the update. Write each as a function taking parameters, gradient and state, and returning the new parameters. Then the same loop can run any optimizer on our valley function.

• Each is a few lines
• State: v, s, m, t

Start all optimizers at the same point and give each the same number of steps. Plain SGD zigzags and lags. Momentum oscillates a little, then speeds along the valley. Nesterov overshoots less. RMSProp and Adam head almost straight toward the minimum because each coordinate is rescaled. Watch for one caveat: with a fixed learning rate, Adam can hover around the minimum instead of settling, which is why schedules matter.

• Same start, same budget
• Plain SGD slowest
• Adam and RMSProp take direct routes

What should you use? Adam is a strong default that needs little tuning, and it is common for transformers. SGD with momentum is still widely used, notably for convolutional networks, and often generalizes well with a good schedule. Whichever you pick, the learning rate is the most important hyperparameter, so tune it first, and try a few values across orders of magnitude.

• Adam: strong default
• SGD with momentum: often generalizes well
• Tune learning rate first

Recap. A layer multiplies signal variance by n times variance of w, so we choose Xavier, two over n in plus n out, or He, two over n in, for ReLU. Backward gradients are products, so sigmoid vanishes, bad scales explode, and clipping caps explosions. Momentum averages gradients, RMSProp rescales them, and Adam does both with bias correction. Schedules and warmup tame the learning rate over time.

• Var(w) = 1/n_in, 2/(n_in+n_out), 2/n_in for ReLU
• Gradients are products
• Momentum, RMSProp, Adam
• Schedules and warmup

Good initialization and optimizers get a deep network moving, but they do not stop it from overfitting or from becoming unstable as it trains. In lesson 48 we cover dropout, batch normalization, and other ways to regularize deep networks.

• Lesson 48 next
• Subscribe for the series