PaperlayerView exact revision on arXiv

Source-linked AI summary

A Kronecker-factored approximate Fisher matrix for convolution layers

Roger Grosse, James Martens

arXiv:1602.01407v2stat.MLcs.LG

TL;DR

Large neural networks make exact natural-gradient computation impractical, while existing approximations trade curvature fidelity against computational cost. The paper introduces KFC, a Kronecker-factored Fisher approximation for convolutional networks, and reports optimization several times faster than carefully tuned SGD, with 10-20 times fewer iterations.

  • Problem

    Exact natural-gradient computation is impractical for large networks, while existing Fisher approximations either impose restrictive structure or require expensive iterative procedures.

  • Method

    KFC models backpropagated derivatives probabilistically and uses blockwise Kronecker factorization to obtain a tractable Fisher approximation for convolutional layers.

  • Results

    KFC optimized convolutional networks several times faster than carefully tuned SGD with momentum and required 10-20 times fewer iterations.

  • Takeaways & Limitations

    KFC provides approximate natural-gradient updates for convolutional networks while retaining tractable computation and invariance to common activation reparameterizations.

  • Takeaways & Limitations

    Iterative alternatives may require tens or hundreds of conjugate-gradient iterations per update, while empirical autocovariance estimates can produce non-positive-semidefinite factors and inaccurate covariances.

Abstract

from arXiv · show

Second-order optimization methods such as natural gradient descent have the potential to speed up training of neural networks by correcting for the curvature of the loss function. Unfortunately, the exact natural gradient is impractical to compute for large models, and most approximations either require an expensive iterative procedure or make crude approximations to the curvature. We present Kronecker Factors for Convolution (KFC), a tractable approximation to the Fisher matrix for convolutional networks based on a structured probabilistic model for the distribution over backpropagated derivatives. Similarly to the recently proposed Kronecker-Factored Approximate Curvature (K-FAC), each block of the approximate Fisher matrix decomposes as the Kronecker product of small matrices, allowing for efficient inversion. KFC captures important curvature information while still yielding comparably efficient updates to stochastic gradient descent (SGD). We show that the updates are invariant to commonly used reparameterizations, such as centering of the activations. In our experiments, approximate natural gradient descent with KFC was able to train convolutional networks several times faster than carefully tuned SGD. Furthermore, it was able to train the networks in 10-20 times fewer iterations than SGD, suggesting its potential applicability in a distributed setting.

1 Introduction

The paper develops KFC as a tractable Fisher-matrix approximation for convolutional networks, extending Kronecker-factored curvature methods to convolution layers. Experiments show faster optimization than carefully tuned SGD and substantially fewer iterations.

  • The approach targets the gap between exact natural-gradient computation, which is impractical for large models, and cheaper approximations that may be restrictive or iterative.The paper seeks curvature information with minimal overhead beyond standard gradient computations.
  • KFC approximates each convolutional-network Fisher block with a tractable Kronecker-factored structure for efficient inversion.Its construction is based on a structured probabilistic model of backpropagated derivatives.
  • KFC updates are invariant to widely used network reparameterizations, including whitening or centering of activations.
  • KFC optimized convolutional networks several times faster than carefully tuned SGD with momentum on object-recognition benchmarks.The comparison covered both training and test error.
  • KFC trained the networks in 10-20 times fewer iterations than SGD, suggesting usefulness for highly distributed training algorithms.

2 Background

The background frames natural-gradient optimization as solving a Fisher-based curvature system and reviews K-FAC's tractable block and Kronecker approximations. These approximations reduce inversion cost but rely on structural assumptions and may require damping or iterative procedures.

  • Second-order optimization of neural networks: Natural-gradient descent uses the Fisher information matrix as the curvature matrix in a second-order optimization update.For exponential-family models with natural parameters, the Fisher is equivalent to the generalized Gauss-Newton matrix.
  • Second-order optimization of neural networks: The Fisher matrix has dimension n × n, where n can reach tens of millions of parameters in modern deep architectures.Explicitly representing or exactly solving the associated system is therefore impractical.
  • Second-order optimization of neural networks: Iterative curvature methods can require tens or hundreds of conjugate-gradient iterations per update, each comparable in cost to an SGD update.
  • Kronecker-factored approximate curvature: K-FAC assumes weight derivatives in different layers are uncorrelated, making the Fisher approximation block diagonal with one block per layer.
  • Kronecker-factored approximate curvature: K-FAC further assumes activations and pre-activation derivatives are independent, decomposing each layer block into a Kronecker product of smaller matrices.The factors are estimated online from empirical moments using exponential moving averages.
  • Kronecker-factored approximate curvature: Adding regularization or damping to the approximate curvature can break the Kronecker factorization structure.The paper notes that exact solution remains possible, but approximation may be preferred.

3 Kronecker factorization for convolution layers

KFC extends Kronecker-factor Fisher approximations to convolution layers by combining structured assumptions about activations and backpropagated derivatives. The resulting factors are tractable to estimate, invert, and use for approximate natural-gradient updates.

  • KFC for convolution layers: Convolutional Fisher blocks are approximated using a block-diagonal structure and a new Kronecker factorization called Kronecker Factors for Convolution (KFC).Pooling and response-normalization layers are excluded because they have no trainable weights.
  • Approximation assumptions: The approximation assumes independent activations and derivatives, spatial homogeneity, and spatially uncorrelated pre-activation derivatives.Spatial homogeneity makes statistics depend on relative spatial offsets, while SUD sets covariance between distinct spatial locations to zero.
  • Kronecker factorization: Combining IAD, SH, and SUD yields a Kronecker factorization of each approximate Fisher block, enabling efficient computation of its inverse.The factorization is applied to each layer block in vectorized form and supports approximate natural-gradient computation.
  • Computational tractability: For an AlexNet example, the factors are 1201 × 1201 and 128 × 128, making exact representation and inversion practical.The corresponding convolutional weight-and-bias matrix is 128 × 1201.
  • Estimating and using the factors: Empirical factor estimates are positive semidefinite and reuse convolution-operation matrices, while boundary-aware covariance estimation can otherwise produce instability or inaccurate covariances.The estimates are obtained from model-distribution samples and maintained with exponential moving averages; KFC can then provide approximate natural gradients or preconditioning.

4 Theoretical analysis

KFC provides invariance guarantees for convolutional networks under a restricted class of activation transformations and relates its Fisher approximation to other curvature methods. The analysis also identifies spatial-independence assumptions as an important source of lost curvature information.

  • Invariance: KFC updates remain equivalent under homogeneous, pointwise affine transformations of convolutional activations before and after the nonlinearity.These transformations include normalization to zero mean and unit variance.
  • Invariance: Arbitrary affine transformations are excluded for convolutional layers because transformations such as spatial permutations can change the set of representable functions.The guarantee is therefore narrower than full parameterization invariance.
  • Invariance: The invariance theorem states that a transformed network can compute the same function as the original while producing equivalent KFC updates.The proof is deferred to Appendix B.
  • Relationship with other algorithms: Combining IAD, SH, SUA, and WD yields an approximation to Fisher-matrix entries based on structured activation and derivative assumptions.The approximation uses an uncentered autocovariance function Ω̃(j, j′, δ).
  • Relationship with other algorithms: Dropping the β(δ, δ′) term makes the approximate natural-gradient update equivalent to idealized PRONG up to rescaling.This connects KFC to activation-whitening approaches.
  • Relationship with other algorithms: WD appears empirically adequate, whereas SUA loses substantial information because image activations exhibit a 1/f^p power spectrum with p ≈ 2.The resulting curvature can differ substantially between low- and high-frequency Fourier directions.

5 Experiments

Experiments evaluated KFC on CIFAR-10 and SVHN convolutional architectures, testing its probabilistic assumptions and optimization performance against tuned SGD and adaptive baselines.

  • Probabilistic modeling assumptions: KFC assumed independent activations and derivatives, spatial homogeneity, and spatially uncorrelated derivatives when modeling convolutional Fisher structure.The experiments specifically examined whether these assumptions matched observed correlations.
  • Probabilistic modeling assumptions: Derivatives were very weakly correlated spatially and across feature maps, supporting the spatially uncorrelated derivatives assumption for the tested networks.Modest exceptions occurred in early cross-map correlations and the top CIFAR-10 convolution layer.
  • Probabilistic modeling assumptions: Replacing max-pooling with average pooling tested the hypothesis that max-pooling causes weak derivative correlations through its sparsifying effect.Max-pooling makes derivatives zero unless the corresponding activation is maximal within its pooling group.
  • Probabilistic modeling assumptions: Activations had much stronger spatial and cross-map correlations than derivatives, challenging spatially uncorrelated activations assumptions.This contrast was observed in both CIFAR-10 and SVHN convolutional networks.
  • Optimization performance: Adagrad, RMSProp, and Adam did not outperform carefully tuned SGD with momentum.The comparison used these methods as additional optimization baselines.
  • Optimization performance: KFC-pre reached 19% CIFAR-10 test error in 3 minutes versus 9 minutes for SGD, and 6% training error in 4 versus 30 minutes.Both methods approximately reached the previously published 18% CIFAR-10 test error before overfitting.
  • Optimization performance: KFC-pre achieved a 20-fold reduction in updates for 10% CIFAR-10 training error: 300 updates versus 6000 for SGD.Similar iteration speedups were reported for test error and on SVHN.

A.1 KFC as a preconditioner for SGD

KFC-pre uses the approximate Fisher inverse as a preconditioner within an SGD-like procedure, with clipping, momentum, averaging, and damping-related controls.

  • Preconditioned optimization: KFC-pre substitutes the approximate natural gradient for the Euclidean gradient in an SGD-with-momentum-like optimization procedure.Equivalently, the approximate Fisher matrix serves as the preconditioner.
  • Hyperparameters: Suggested settings include momentum µ = 0.9, averaging timescale τ equal to the dataset’s mini-batch count, damping γ = 10^-3, and clipping parameter C = 0.3.The damping value may require tuning.
  • Update construction: The procedure computes the approximate natural gradient by applying the inverse damped approximate Fisher matrix to the backpropagated gradient.The update is then scaled by the learning rate before momentum and parameter updates.
  • Update control: KFC-pre clips updates when ν = v⊤Fv + λv⊤v exceeds the clipping threshold C.The estimate uses a subset of the current mini-batch and constrains the update’s curvature-related size.
  • Update control: The algorithm applies momentum and exponential parameter averaging after computing the clipped update.The averaging step smooths the parameter trajectory over training iterations.

A.2 Kronecker-factored approximate curvature

The curvature framework combines quadratic models, damping, and Kronecker-factored Fisher approximations to obtain practical second-order updates without explicitly forming the full Fisher matrix.

  • Quadratic optimization: K-FAC combines Fisher approximations so that approximate natural-gradient updates can be computed within a second-order optimization framework.The framework uses quadratic models to guide updates and parameter choices.
  • Quadratic optimization: The update v = α ˆ∇h + µv_prev is chosen by minimizing a quadratic model over learning rate α and momentum parameter µ.The coefficients can be obtained by solving a low-dimensional linear system involving curvature matrix-vector products.
  • Damping: Adaptive Tikhonov-style damping limits the risk of poor steps caused by inaccurate quadratic models or approximate curvature.Damping is applied to the quadratic models and to the approximate Fisher matrix before inversion.
  • Adaptive control: K-FAC dynamically adapts damping according to how accurately the quadratic model predicts the objective decrease.The reduction ratio determines whether damping is loosened or tightened.
  • Damping: Adding γI before inversion compensates for approximate Fisher inaccuracies and makes the resulting natural-gradient direction more conservative.The damping can be implemented using the Kronecker factors or exact eigendecomposition identities.

A.3 Efficient implementation

The implementation places most KFC computation on the GPU but incurs additional per-iteration costs for factor estimation, inverses, and update clipping.

  • Implementation: KFC required additional forward-mode automatic-differentiation kernels for max-pooling and response-normalization layers.These kernels were used for update clipping.
  • Implementation: Most KFC operations ran on the GPU, while computing the Kronecker-factor inverses was performed on the CPU.The inverse computation was the main exception to GPU-bound heavy computation.
  • Computational overhead: KFC introduces per-iteration overhead from estimating the clipped-update quantity ν and from recomputing curvature factors and inverses.These costs are in addition to the approximate natural-gradient computation.
  • Computational overhead: Stale curvature information and subset-based estimates can reduce overhead without substantially harming per-epoch progress.Inverse computations could also be performed asynchronously, although this was not used in the experiments.

B.1 Proofs for Section 3

The proofs combine the IAD, SH, and SUD approximations to derive a factorization of convolutional Fisher-matrix entries.

  • Lemma 1 supplies the first equality used in the derivation, with the remaining equalities obtained analogously.
  • Assumption IAD is used in the final line of the derivation.
  • Theorem 1 combines approximations IAD, SH, and SUD to obtain the stated factorization.
  • The factorization includes a spatial-overlap count, a derivative covariance term, and an indicator enforcing matching output coordinates.The indicator arises because activations are defined as zero outside the spatial-location set.
  • Under SH, the corresponding covariance expression reduces to a spatial-overlap term multiplied by Ω(j, j′, δ′−δ).

B.2 Proofs for Section 4

The proofs establish that KFC updates remain functionally equivalent under supported affine reparameterizations and derive an approximate Fisher block using several structural assumptions.

  • Reparameterization invariance: The invariance proof uses homogeneous coordinates, affine transformations, and induction to show that both networks compute the same function.
  • Reparameterization invariance: Theorem 3 shows that reparameterized networks can compute the same function and receive equivalent KFC updates.The equivalence follows by relating transformed parameters, Kronecker factors, and updates between the two networks.
  • Fisher approximation: Theorem 4 combines IAD, SH, SUA, and WD to approximate entries of the convolutional Fisher matrix.
  • Relation to PRONG: Dropping the β and |T| terms makes the approximate natural-gradient rule equivalent to idealized PRONG, up to rescaling.
  • Fisher approximation: The derivation computes second moments under SH, SUA, and WD before expressing the result in matrix form and using block Cholesky decompositions.
  • Fisher approximation: The derived approximation has an approximate Fisher block of ˜Ω⊗I under the theorem’s assumptions.
Loading 1602.01407v2…