Natural gradient

Overview

Plain gradient descent moves every weight by the raw gradient times one scalar step size. The natural gradient replaces that step with one that is measured in the geometry of the model's output distribution: the step is multiplied by the inverse of the Fisher information matrix 1410.7455 §4.1. The Fisher matrix describes how much the output changes for each direction in weight space, so directions the model is very sensitive to get smaller steps.

The Fisher matrix is too large to form or invert for a DNN. The common approximation in this source is to use one block per weight matrix, and to write each block as a Kronecker product of two smaller positive definite factors 1410.7455 §4.3 §4.4. Inverting the two small factors is then cheap compared with inverting the full block. The factors are estimated from the inputs and output gradients that the layer sees during training [arxiv:1410.7455 §4.7]. Every variant below applies a per-matrix preconditioner built this way; the variants differ in how the factors are estimated and kept up to date.

Variants

Simple minibatch estimate

The factors are estimated from the other samples in the current minibatch, with no state carried across minibatches 1410.7455 app. A. The paper measures this only at one job, where its word error rate is 23.16%, close to the online method's 23.19% 1410.7455 tab. 2. The matrix inversion dominates the time on GPUs, so the authors judge it more costly there 1410.7455 app. A.3. Card: Parallel training of DNNs with Natural Gradient and Parameter Averaging.

Online low-rank estimate (the authors' default)

Each factor is kept as a running low-rank estimate: a few orthonormal directions with their own eigenvalues, plus a scalar for the remaining directions. Each new minibatch updates this state with a forgetting factor, so older data fades out 1410.7455 app. B.1 app. B.2. The forgetting factor is set from the minibatch size and an approximate sample count of S = 2000 [arxiv:1410.7455 app. B.4]. The reported ranks are 20 on the input side and 80 on the output side of each weight matrix, and the factors are refreshed on the first 10 minibatches and then every 4th one, with only the cheap preconditioning step in between 1410.7455 app. B.4 app. B.5. The paper uses this method for sequence training as well 1410.7455 app. C.8. Card: Parallel training of DNNs with Natural Gradient and Parameter Averaging.

Plain SGD baseline

No preconditioner. The step is the gradient times a per-job learning rate that decays exponentially, from 0.01 to 0.001 in the 6-job setting 1410.7455 §3.2.1 §6.2. This is the baseline the paper compares against at 1, 2 and 4 jobs 1410.7455 tab. 2.

Safeguards that apply to any variant

Two safeguards appear in the paper and are reusable on other methods. The preconditioned output is rescaled so that its norm matches the norm of its input, which stops the preconditioner from changing the overall step size 1410.7455 §4.8. A separate cap scales each layer's update down when its bound on the change per sample exceeds a limit, with the cap set to 0.075 per sample in the recipe 1410.7455 app. C.3. The paper says the cap matters mostly early in training, for layers near the output 1410.7455 app. C.3.

When to use which

The paper's evidence is one speech-recognition task with author-reported numbers 1410.7455 §6. The rows below follow from what that evidence shows.

Learning-rate scaling with the number of jobs is not given as a formula. The paper gives the rate for 6 jobs and says the rate should scale with the job count 1410.7455 §3.1 §6.2.

Implementation sketch

Preconditioning one weight matrix with online Kronecker factors. The helper names stand for the low-rank updates in the paper's appendices, and each one is cited where it is described.

# Preconditioned update for one weight matrix W, shape (out, in)  [§4.7]
# Input:  grad (out, in) from the minibatch; inputs X (N, in); output
#         gradients Y (N, out); per-layer factor states S_in, S_out
# Output: updated W

def natural_gradient_step(W, grad, X, Y, S_in, S_out, lr):
    # Refresh the factors on a schedule (first 10 steps, then every 4) [app. B.5]
    if refresh_due(step):
        S_in = update_low_rank_factor(S_in, X)             # [app. B.2]
        S_out = update_low_rank_factor(S_out, Y)           # [app. B.2]
    # Kronecker-product preconditioning: P = B^-1 grad A^-1  [§4.4] [§4.7]
    precond = apply_inverse_right(S_in, apply_inverse_left(S_out, grad))
    # Restore the norm of the raw gradient so the scale is unchanged  [§4.8]
    precond = precond * norm(grad) / norm(precond)
    # Cap the update size for this layer  [app. C.3]
    scale = min(1.0, cap_limit(X, Y, lr) / bound(X, Y, lr))
    W = W - lr * scale * precond
    return W, S_in, S_out

The helper that inverts a low-rank-plus-identity factor is in app. B.3. The released Kaldi code differs from the paper in several defaults and details; those differences are listed on the card's Paper vs code section (/papers/1410.7455), and they matter if you copy its settings.

Contradictions

None across the cards that cite this concept, since one source is compiled so far. Within that source, the smoothing constant appears as 1e-20 in the text and 1e-10 in the online method 1410.7455 §4.9 app. B.2. Table 1's DNN2 result (23.79%) matches no entry in Table 2, and the evaluation set for Table 2 is not stated 1410.7455 tab. 1 tab. 2.

Papers

• Parallel training of DNNs with Natural Gradient and Parameter Averaging: online low-rank Kronecker natural gradient for Kaldi DNN training, combined with parameter averaging; author-reported results on Fisher English.