Parallel training of DNNs with Natural Gradient and Parameter Averaging

In plain words

Training a large neural network on a huge pile of data is slow, so people spread the work over several machines. The usual way makes the machines swap updates after almost every small batch of examples, which needs a very fast network. This paper takes a looser route. Each machine trains its own copy of the model on its own part of the data. Every minute or two, the copies are averaged and sent back out. Little traffic is needed, but plain averaging does not work well on its own.

What makes it work is a change to how each machine takes its step. Ordinary gradient descent moves every weight by a single scalar step size. Here, each weight matrix gets its own step shape, an approximation of the inverse of the Fisher information matrix. That matrix describes how sensitive the model's output is to each weight, and it is estimated cheaply and kept up to date during training. The authors call the result natural gradient. It also speeds up training on a single machine.

With both pieces, the authors report that adding machines gives close to a linear speed-up for a few machines, and that the word error rate stays about the same as the number of machines grows. The experiments are speech recognition on conversational telephone speech, and the numbers are the authors' own.

TL;DR

Kaldi's parallel DNN training: N jobs run SGD on different data and average their parameters after every K samples. Averaging only works together with an online natural-gradient preconditioner, which keeps a low-rank estimate of the inverse Fisher matrix for each weight matrix. Author-reported WER on Fisher English stays between 22.8% and 23.4% from 1 to 16 jobs, and the linear speed-up holds only up to about 4 jobs. Check the pitfalls before reusing the defaults: the released recipe differs from the paper's experiments.

Problem

Data-parallel SGD usually synchronises the model every minibatch, which is costly when machines are connected by a slow network §1. The authors instead average the parameters only every minute or two, with each job seeing different data §1. Their observation is that plain averaging by itself does not work well for these DNNs, and that it only works when combined with an efficient natural-gradient SGD §1. The setting is DNNs for speech recognition, where the training objective is the sum over frames of the log-probability of the correct HMM state §2. The paper gives no proof for the method; it claims only that the methods work empirically §1, §5.

Method

Each job runs SGD on its own randomised subset of the data. After every job has processed K samples (typically K = 400 000), the parameters are averaged across jobs and sent back to the jobs. Repeat until a fixed number of epochs §3.1. Averaging is preferred to summing the per-job changes, because summing can overshoot in a stiff direction §3.1.

1. Per-job learning rate and effective rate. The effective learning rate of the parallel procedure is the per-job rate divided by N. Keeping the effective rate fixed while adding jobs means raising the per-job rate in proportion to N §3.1. The experiments scale the initial and final rates with the number of jobs, using 0.01 to 0.001 as the 6-job default §6.2.
2. Learning-rate schedule. The rate decays exponentially during training, by about a factor of 10 overall. Unless stated otherwise the rate starts at 0.01 and ends at 0.001 §3.2.1. The number of epochs is fixed in advance §3.2.1.
3. Gradient convention. Gradients are summed over the minibatch rather than averaged. Progress per sample then does not depend on the minibatch size app. C.1. The rate and the minibatch size are therefore tuned independently §4.7.
4. Natural-gradient update. Each weight matrix W_i gets a learning-rate matrix that approximates the inverse Fisher matrix §4.1, §4.2. The Fisher matrix is block-diagonal, one block per weight matrix, and each block is a Kronecker product A_i ⊗ B_i of two positive definite factors §4.4. Each factor has the form λI + XX^T, a low-rank part plus a multiple of the identity §4.4.
5. Per-minibatch form (online method). For a minibatch with input rows X and output-gradient rows Y, the preconditioned quantities are X̄ = X A^{-1} and Ȳ = Y B^{-1}, and the weight update is W += η X̄^T Ȳ §4.7. The interface is called once for the input side and once for the output side of each weight matrix, so 2I calls per minibatch for I weight matrices §4.7.
6. Rescaling. After preconditioning, the output is rescaled so that its Frobenius norm equals that of the input. This stops the preconditioner from changing the overall step size §4.8. The authors accept a small bias from computing the scale on the same minibatch §4.8.
7. Smoothing. The factor is smoothed by adding β I, with β = (α/D) max(tr S, ε) and α = 4 §4.9. The smoothing uses ε = 10^-20 in the text of §4.9, but the online method uses ε = 10^-10 app. B.2.
8. Two estimators of the factors. The simple method estimates the factors from the other samples in the minibatch app. A. The online method keeps a running estimate with a forgetting factor and is the one the authors normally use §4.5, app. B. Only the online method is used on sequence training app. C.8.
9. Online state. For each factor the state is a low-rank approximation F_t = R_t^T D_t R_t + ρ_t I, where R_t has R orthonormal rows, D_t is a non-negative diagonal, and ρ_t > 0 app. B.1. The preconditioner uses G_t = F_t + (α tr F_t / D) I app. B.1. The output is γ_t X G_t^{-1} with γ_t chosen to restore the norm app. B.1.
10. Forgetting factor. η = 1 − exp(−N/S), where N is the minibatch size and S = 2000 is the approximate number of samples in the Fisher estimate app. B.4. The rank is R = 20 on the input side of each weight matrix and R = 80 on the output side app. B.4.
11. Update schedule. The factor is updated on the first 10 minibatches of a process and then every J = 4 minibatches app. B.4, app. B.5. Between updates only the preconditioning is applied app. B.5.
12. Factor update. The update takes a power-method step on the blend of the minibatch covariance and the current factor, then takes the top eigen-directions as the new basis and sets the scalars so that the trace is preserved app. B.2. The efficient form that avoids forming large matrices is in app. B.3.
13. Maximum parameter change. Each layer's update is scaled by a factor α_t ≤ 1 so that an upper bound on its Frobenius norm stays below N times 0.075 [app. C.3]. This cap is applied on top of the natural-gradient update app. C.3. The paper says it is active mainly early in training, for layers near the output app. C.3.
14. Model combination (optional). At the end of training the last P = 20 averaged models are combined with a per-layer weighted sum. The weights are fitted with L-BFGS on training data app. C.4. The authors say this gives a small gain but report no experiments for it app. C.4.

Parallel SGD with periodic averaging, the outer loop:

# Parallel SGD with periodic parameter averaging  [§3.1] [app. C.2]
# Input: one data shard per job; K samples per job per outer iteration
# Output: averaged parameters after the last outer iteration

def parallel_train(shards, params, num_outer_iters, K):
    for it in range(num_outer_iters):
        # Per-job rate from the exponential schedule  [§3.2.1]
        lr = learning_rate_for(it)
        # Each job runs SGD on its next K samples  [§3.1]
        job_params = [
            sgd_on_samples(shard, params, K, lr) for shard in shards
        ]
        # Average the jobs and send the result back to all of them  [§3.1]
        params = mean_over_jobs(job_params)
    return params

Preconditioning one minibatch, in the online method:

# Apply the online preconditioner to one minibatch  [app. B.1] [app. B.3]
# Input: inputs (N, dim); basis (rank, dim) with orthonormal rows;
#        eigvals (rank,); rho (scalar); alpha = 4
# Output: preconditioned matrix (N, dim), same Frobenius norm as inputs

def precondition(inputs, basis, eigvals, rho, alpha):
    dim = inputs.shape[1]
    # beta = rho (1 + alpha) + alpha * tr(D) / dim  [app. B.3]
    beta = rho * (1 + alpha) + alpha * sum(eigvals) / dim
    # Inverse of G = basis^T diag(eigvals) basis + beta I  [app. B.3]
    scale = 1 / (beta / eigvals + 1)                   # (rank,)
    inv_G = (eye(dim) - basis.T @ diag(scale) @ basis) / beta
    # Unscaled output  [app. B.1]
    hat = beta * (inputs @ inv_G)                      # (N, dim)
    # Restore the Frobenius norm of the inputs; gamma = 1 if zero  [app. B.1]
    gamma = sqrt(sum(inputs ** 2) / sum(hat ** 2))
    return gamma * hat

Factor update after a minibatch. It is called when t < 10 or t is a multiple of J = 4 app. B.5. The formulas are the direct ones; the efficient form in [app. B.3] gives the same result:

# Update basis, eigvals and rho from one minibatch  [app. B.2]
# Input: inputs (N, dim); state basis, eigvals, rho; eta from S = 2000
# Output: new state (basis, eigvals, rho); eps = 1e-10  [app. B.2]

def update_factor(inputs, basis, eigvals, rho, eta):
    dim = inputs.shape[1]
    rank = basis.shape[0]
    count = inputs.shape[0]
    # Uncentred covariance of this minibatch  [app. B.2]
    covariance = inputs.T @ inputs / count             # (dim, dim)
    # Blend the covariance with the current factor  [app. B.2]
    factor = basis.T @ diag(eigvals) @ basis + rho * eye(dim)
    target = eta * covariance + (1 - eta) * factor
    # One power-method step  [app. B.2]
    product = basis @ target                           # (rank, dim)
    # Eigen-decomposition of the small Gram matrix  [app. B.2]
    c_values, vectors = eigh(product @ product.T)
    # Floor the eigenvalues at (1 - eta)^2 rho^2  [app. B.3]
    c_values = maximum(c_values, (1 - eta) ** 2 * rho ** 2)
    # New basis with orthonormal rows  [app. B.2]
    new_basis = diag(c_values ** -0.5) @ vectors.T @ product
    # rho' keeps the trace of the factor unchanged  [app. B.2]
    rho_prime = (
        eta * sum(inputs ** 2) / count
        + (1 - eta) * (dim * rho + sum(eigvals))
        - sum(c_values ** 0.5)
    ) / (dim - rank)
    # Non-identity part and identity part, each floored at eps  [app. B.2]
    new_eigvals = maximum(c_values ** 0.5 - rho_prime, eps)
    new_rho = maximum(eps, rho_prime)
    return new_basis, new_eigvals, new_rho

Maximum parameter change for one layer app. C.3:

# Scale one layer's update to stay under a per-minibatch cap  [app. C.3]
# Input: rows x (N, dim_in) and y (N, dim_out); learning rate lr
# Output: factor alpha in (0, 1] applied to this layer's update

def max_change_scale(x, y, lr, max_change_per_sample):
    count = x.shape[0]
    # Upper bound on the Frobenius norm of the minibatch update  [app. C.3]
    bound = sum(lr * norm(x[i]) * norm(y[i]) for i in range(count))
    # Cap for the whole minibatch: N times the per-sample cap  [app. C.3]
    limit = count * max_change_per_sample
    return min(1.0, limit / bound)

Hyperparameters

Results

Author-reported. Setup: Fisher English, English conversational telephone speech at 8 kHz, 1600 hours of transcribed training data, and a held-out set of about 3.3 hours that the authors defined themselves §6. The GMM baseline uses 7 880 context-dependent states and 300 000 Gaussians §6.1. Table 1 is on the Fisher dev set tab. 1.

The set used for tab. 2 is not stated. Table 2 is the only place the simple NG method is measured, and only at N = 1 tab. 2.

Convergence, objective against data processed fig. 1a: natural gradient is always above plain SGD at the same job count. Online NG curves stay close to each other up to about 4 jobs, and the 8- and 16-job runs converge more slowly. The speed-up is linear for small N and sub-linear above about 4 §6.2. The paper also says linear speed-up holds only up to about 4 or 8 GPUs §1.

Simulated clock time fig. 1b, based on the time per outer iteration on a K20 GPU: 88 s for plain SGD, 93 s for online NG-SGD and 208 s for plain NG-SGD. Actual time depends on queue load §6.2.

Paper vs code

Official code: Kaldi, repository kaldi-asr/kaldi, commit e02e35f0254bb033fab73d1df99fc34123e31d56, Apache-2.0 COPYING. The paper does not link it. The repository was found by a GitHub search for "kaldi"; it is the toolkit the abstract names [abstract]. The method's preconditioner file carries the copyright of Daniel Povey, an author of this paper nnet-precondition-online.cc. The repository is owned by the kaldi-asr organisation, not a named author. The commit is the default branch at clone time, not the 2014 code, so the findings below describe the released code, not the code used for the paper. The recipe used is egs/wsj/s5/steps/nnet2/train_pnorm_fast.sh.

Differences found:

• Experiment settings are not the recipe defaults. The paper's DNN1 uses 12 epochs, 0.08 to 0.008 and 8 jobs §6.1. The recipe defaults are 15 epochs plus 5 extra, 0.04 to 0.004 and 16 jobs train_pnorm_fast.sh train_pnorm_fast.sh.
• Samples per outer iteration. The paper uses K = 400 000 §6.1. The recipe default is 200 000 train_pnorm_fast.sh.
• Minibatch size. The paper's typical GPU minibatch is 512 app. B.5. The recipe default is 128 train_pnorm_fast.sh.
• Rank, α, update period and max change match. Rank 20 and 80, α = 4, update period 4 and max change 0.075 match app. B.4 app. C.3 and the fast recipe train_pnorm_fast.sh.
• Class defaults are not the recipe values. The preconditioner class defaults are rank 40 and update period 1 nnet-precondition-online.cc. Anyone instantiating the class directly must pass the recipe values.
• Floor on ρ and D. The paper floors ρ and D at ε = 10^-10 app. B.2. The code floors both at max(ε, δ · max √c) with δ = 5·10^-4 nnet-precondition-online.cc nnet-precondition-online.cc.
• Floor on c. The paper floors each eigenvalue at (1 − η)² ρ² app. B.3. The code does the same with (ρ(1 − η))² nnet-precondition-online.cc.
• Cap on η. The paper gives η = 1 − exp(−N/S) without a cap app. B.4. The code caps η at 0.9 nnet-precondition-online.cc.
• Update schedule. The code updates on the first 10 minibatches and then skips until the period is reached nnet-precondition-online.cc. This matches the paper.
• Parameter initialisation. The paper uses std 1/√(fan-in) for hidden weights app. C.7. The recipe sets one std, 1/√(pnorm_input_dim), for every affine layer including the first train_pnorm_fast.sh train_pnorm_fast.sh train_pnorm_fast.sh. Its bias std is 0.5 train_pnorm_fast.sh, which the paper does not state. Softmax initialisation matches the paper: zero app. C.7 train_pnorm_fast.sh.
• Model combination regulariser. The paper uses about 1e-10 times the squared weights app. C.4. The generic recipe sets 1e-14 train_pnorm.sh. The fast recipe passes no regulariser to nnet-combine-fast and uses 40 L-BFGS iterations train_pnorm_fast.sh.
• Averaging. The fast recipe averages every iteration with nnet-am-average train_pnorm_fast.sh. This matches §3.1.
• Simple method. No implementation of the simple NG-SGD (appendix A) was found in src/nnet2/nnet-precondition-online.{h,cc}. Table 2's simple row cannot be reproduced from this commit.

Pitfalls

• Learning rate scales with jobs. The 0.01 to 0.001 rates belong to the 6-job case §6.2. For other job counts, scale the rates in proportion to N. The paper does not give a formula beyond that statement §3.1, §6.2.
• Averaging can hurt early. On the outer iterations right after random initialisation, including the first, averaging may make the objective worse. The paper then keeps the one job with the best objective on its own data instead of averaging app. C.7.
• Two ε values. ε = 1e-20 appears in the smoothing text §4.9, but the online floors use 1e-10 app. B.2. Use 1e-10 for the online method.
• Near-zero inputs. Early in training the inputs can be very small or zero, which makes the inverse huge. γ is set to 1 when its denominator is zero app. B.1. The factors are floored at ε app. B.2. The paper says single precision was used throughout app. B.3.
• Cancellation in tr(XXᵀ). The efficient form computes tr(XXᵀ) as tr(X̂X̂ᵀ) − tr(LE) + 2 tr(L), not directly app. B.3. This can lose precision when the terms nearly cancel. The paper does not discuss this.
• Orthogonality check. R Rᵀ = I can drift from roundoff. The check and fix run only when the condition number of C_t exceeds 10^6 or a floor was active, with a threshold of 10^-3 app. B.3. The paper says re-orthogonalisation is rare.
• Max-change division. The paper does not say what happens when the bound in the max-change scale is zero app. C.3. In the pseudocode above, a zero bound makes the ratio infinite and the factor 1, unless the limit is zero too. That is my reading of the formula, not a statement in the paper.
• Table 2 is ambiguous. Its evaluation set is not stated. DNN2's 23.79 in Table 1 matches no entry in Table 2, and the 6-job case, which is the default, is not in Table 2 tab. 1 tab. 2.
• Feature naming. §6.1 calls the DNN2 inputs log-mel filterbank features, but app. C.9 calls them MFCCs and says the two are equivalent for this purpose §6.1 app. C.9.
• Speed-up is limited. Gains flatten above about 4 jobs, and the authors expect the speed-up only to about 4 or 8 GPUs §1, §6.2.
• Cost of natural gradient. The online method adds about 6% to the time per outer iteration over plain SGD (93 s against 88 s). Plain NG-SGD takes 208 s §6.2. The simple method is more expensive on GPUs, where its matrix inversion takes most of the time app. A.3.
• Sub-classes. The authors later found that mixture components did not diverge enough to count as distinct. They replaced them with a fixed per-class scale proportional to count^−0.25, renormalised to average 1 app. C.5.
• Parameter-averaging of sequence training. Only the online method is used there, because the simple method assumes samples are independent app. C.8.
• Data access. Fisher English is an LDC corpus. The footnote lists the LDC catalogue numbers §6.

Lineage

• Builds on: the Kaldi toolkit app. C.2 [Povey et al. (2011)] and its nnet1 DNN setup app. C.7; p-norm networks §6.1 [Zhang et al. (2014)]; the exponential learning-rate schedule §3.2.1 [Senior et al. (2013)]; layer-wise backpropagation app. C.7 [Seide et al. (2011)]; natural gradient §4.1 [Amari (1998)] [Yang & Amari (1998)] and the approximate Fisher factorisation §4.3 [Roux et al. (2007)].
• Compared against: plain SGD §6.2; Downpour, cited as the model-parallel reference §1 [Dean et al. (2012)], with no experimental comparison; Hogwild for the CPU thread model app. C.1 [Niu et al. (2011)]; Adagrad, which the paper argues against §5.4 [Hazan et al. (2007)]; momentum, which the paper does not use §5.4.
• Superseded by: not stated in content already on this wiki.