Fishy : Layerwise Fisher Approximation for Higher-order Neural Network Optimization
Abstract
We introduce Fishy , a local approximation of the Fisher information matrix at each layer for natural gradient descent training of deep neural networks.
The true
Fisher approximation for deep networks involves sampling labels from the model’s predictive distribution at the output layer and performing a full backward pass –
Fishy deﬁnes a Bregman exponential family distribution at each layer, performing the sampling locally.
Local sampling allows for model parallelism when forming the preconditioner, removing the need for the extra backward pass.
We demonstrate our approach through the Shampoo optimizer, replacing its preconditioner gradients with our locally sampled gradients.
Our training results on deep autoencoder and
VGG16 image classiﬁcation models indicate the efﬁcacy of our construction.
Introduction
Natural gradient descent (NGD) [1] is a second-order update rule that preconditions the gradient direction with the inverse of the model’s
Fisher information matrix (FIM).
NGD corresponds to the steepest descent direction in the Riemannian space associated with the FIM and is equivariant to any differentiable reparameterization of the model weights.
However, this property only holds in the small step limit and whether the beneﬁts extend to ﬁnite steps remains an open question [16].
Nevertheless, NGD has proven to be remarkably effective in training deep neural networks and many standard optimizers in deep learning can be seen as approximations of NGD [17, 11].
NGD is challenging to implement for training deep neural networks.
Firstly, the FIM requires samples from the model’s predictive distribution, with an additional backward pass to calculate the gradients.
Secondly, NGD requires calculating the inverse FIM of the whole network, which immediately becomes infeasible even for medium-sized models.
K-FAC [17] uses a block diagonal
Kronecker factor approximation to address the latter obstacle, considering the FIM for each layer as a Kronecker product of two matrices which can more easily be stored and inverted.
However, the Kronecker approximation used in K-FAC becomes challenging to generalize for layers other than fully-connected layers.
Moreover, the sampling and additional backward passes keep K-FAC expensive.
The Shampoo optimizer [8, 5] approximates the full-matrix AdaGrad preconditioner using the same minibatch gradients of the training examples.
Shampoo preconditioning avoids additional sampling or extra backward passes but is no longer related to true Fisher [14, 16] and, as we show experimentally, is outperformed when replaced with a true Fisher approximation.
Using a Shampoo approximation of the true Fisher is related to [19]; primary differences are the choice of exponents and dampening term used for per-dimension preconditioners.
Our work follows Shampoo’s Kronecker product approximation of the preconditioner, which has the advantage of being agnostic to the type of layers.
Sun and Nielsen [22] recently proposed a layerwise relative FIM, but do not consider the full form of layers’ predictive distributions nor demonstrate scaling to practical architectures.
This paper presents a new layerwise FIM approximation approach, extending the popular Shampoo optimizer to NGD.
First, we update Shampoo to an approximate NGD method, sampling from the output layer distribution to construct the preconditioners.
Second, inspired by the matching loss [9, 4] construction for LocoProp and layerwise representation learning [3, 2], we deﬁne Bregman exponential family distributions for the outputs of each layer based on their activation functions.
Sampling locally from these distributions allows preconditioner calculation to be performed in parallel across layers.
Third, we compare our approach to standard Shampoo and its Fisher counterpart on a benchmark MNIST deep autoencoder problem and CIFAR100 classiﬁcation with a VGG16 [20].
Methodology
Fisher Information Matrix and Natural Gradient Descent Given a probabilistic model
P ( x , y | θ ) with input random variable x and target random variable y , the Fisher Information
Metric (FIM) is deﬁned in terms of a local Riemannian metric as an approximation to the KL divergence [18],
[Equation]
[Table]
[Equation]
where the expectation is with respect to the model distribution P ( x , y | θ ) .
The result is identical when ﬂipping the order of the arguments in the KL divergence.
The FIM corresponds to the Riemannian metric that locally explains the model in the parameter space [15].
This probabilistic model approximates an underlying joint data distribution P d ( x , y ) from which the data is sampled.
Thus, training the model corresponds to ﬁnding optimal parameters θ ∗ that minimize the KL divergence between the data and model distributions.
For a discriminative model that generates P ( y | x , θ ) , the marginal distribution P ( x | θ ) is assumed to be the same as the underlying input data distribution
P d ( x ) .
Thus, the loss can be written as
[Equation]
where the expectations are applied with respect to P d ( x , y ) .
In practice, the empirical loss is formed by a Monte Carlo approximation of the expectation,
[Equation]
using training examples
X = { ( x ( i ) , y ( i ) ) } ni =1 that are sampled from P d ( x , y ) .
Given the loss L ( θ |X ) on the batch of examples X , the natural gradient step [1] is deﬁned as the minimizer of the following linearized objective,
[Equation]
in which the loss at L ( θ + d θ |X ) is approximated by its ﬁrst-order Taylor expansion.
The nonnegative multiplier γ > 0 in the regularizer term controls the step size.
Speciﬁcally, Equation (2) corresponds to minimizing a linear approximation of the loss at the current parameter θ while the amount of deviation is measured in terms of the local Mahalanobis distance induced by the FIM.
Fishy:
Layerwise Fisher Construction Let ˆ a (cid:96) denote the pre-activation at layer (
cid:96) ∈
[ L ] of an L -layer network.
For a fully-connected layer with weights W (cid:96) , we can write ˆ a (cid:96) = W (cid:96) ˆ y (cid:96) − 1 and
[Figure]
( y l |
[Figure]
l )
Tanh a
[Figure] [Figure]
[Figure]
10.0 y l 0.
[Figure]
Leaky Relu 0.
a l =0.
[Figure]
l
[Figure]
the post-activation ˆ y (cid:96) = f (cid:96) (ˆ a (cid:96) ) . Here, f (cid:96) denotes the (elementwise) strictly-increasing activation function in layer (
cid:96) ∈ [ L ] .
(Note that using this notation, ˆ y {+
+} = x .)
Fishy aims to approximate the Fisher F for each layer independently by deﬁning a predictive distribution for y (cid:96) , the layer’s output random variable, as
[Equation]
where ˆ y (cid:96) is the post-activation of the layer given the input ˆ y (cid:96) − 1 .
The term Z (ˆ y (cid:96) ) is the normalizing partition function that ensures the distribution sums to one.
The distribution in Equation (3) belongs to a Bregman exponential family [6] in which the Bregman divergence is induced by the convex function F ∗ (cid:96) .
We apply the layerwise loss construction in [3] to obtain D F ∗ (cid:96) , which corresponds to the matching loss [9, 13, 4] of the activation function f (cid:96) induced by its integral function.
Intuitively, D F ∗ (cid:96) encodes the layer’s local geometry; thus, correcting for the gradient direction using a local
Fisher approximation can stabilize training and improve convergence [22].
For the linear, sigmoid, and softmax activations, Equation (3) yields the standard Gaussian, Bernoulli, and categorical distributions, respectively.
Figure 1 gives some example distributions for tanh and leaky ReLU activations.
To calculate the layer’s FIM, we need to calculate the gradient of the score function log p (cid:96) ( y (cid:96) | ˆ y (cid:96) ) .
Taking the gradient with respect to the pre-activation ˆ a (cid:96) , we have
[Equation]
The last term in Equation (4) is non-zero in general for asymmetric distributions induced by activation functions such as tanh and leaky ReLU (see Figure 1).
Nevertheless, in our approximation, we discard the last term in Equation (4) to form an upper bound of the local FIM approximation.
For a fully-connected layer with θ (cid:96) =
vec( W (cid:96) ) , we can write
[Equation] [Equation]
where P ( x , y (cid:96) )
=
P d ( x ) P (cid:96) ( y (cid:96) | ˆ y (cid:96) ( x | θ (cid:96) )) .
The FIM approximation in Equation (5) corresponds to a block-diagonal approximation of the full FIM as the layers are treated independently.
This is a common procedure for calculating FIM for large neural networks [17, 22].
With Fishy, we have the ﬂexibility of sampling at every layer or dividing the network into multiple disjoint sections and performing the sampling at the output layer of each section.
The sampled gradients are then backpropagated through the corresponding section.
In the extreme case where the entire network is treated as a single section, Equation (5) reduces to a block diagonal approximation of the true
Fisher by sampling from the model’s output predictive distribution [17].
Next, we discuss using the
Shampoo approximation to apply the inverse of Equation (5) as the preconditioner at each layer.
[Figure]
Fishy via Shampoo Approximation For a fully-connected layer with gradient G ∈ R d i × d o ,
Shampoo approximates the outer-product of the vectorized gradients
g = vec( G ) ∈ R d i d
o as a
Kronecker product of two factors
[Equation]
(6) where r is an upper-bound on the rank of the gradients (see [8] for further details)
[Figure]
and L and R correspond to the left and right preconditioners.
The construction in Equation (6) is more general and can be applied to many types of layers.
However, notice that g in Equation (6) corresponds to the average gradient of a batch (and not the per example gradient).
To approximate Equation (5) via Shampoo, we ﬁrst sample y
( i ) (cid:96) from the predictive distribution in Equation (3) given the input sample x
( i ) in the batch { x ( i ) } ni =1 and the corresponding postactivations ˆ y (cid:96) − 1 and ˆ y (cid:96) at layer (cid:96) − 1 and (cid:96) , respectively.
Next, instead of the per example sampled gradient g ( i ) s
=
( y ( i ) (cid:96) − ˆ y
( i ) (cid:96) ) ⊗ ˆ y ( i ) (cid:96) − 1 , we pass the average sampled gradient g s = 1 / n (cid:80) i g
( i ) s into
Shampoo to form the preconditioners.
A similar averaging is performed in K-FAC [17] on the sampled gradients before forming the preconditioners.
Shampoo then applies the (EMA of the) left
L s and right R s Fishy preconditioners on the gradient of the training examples G ,
[Equation]
Notice that the exponents in Equation (7) are different than the original Shampoo formulation ( − 1 / 4 ) as we are interested in the inverse of the local FIM (and not its inverse square root).
In practice, the exponent in Shampoo is treated as a tunable hyperparameter.
Experiments
We perform experiments on two architectures, and datasets, a deep MNIST autoencoder [10] with fully connected hidden layers [1000, 500 × 8, 250, 30, 250, 500 × 8, 1000], a standard benchmark for second-order methods, and a VGG16 CNN architecture [20] applied to CIFAR100.
We compare
Adam [12], standard Shampoo [5], FIM Shampoo (true Fisher), and our method Fishy.
For VGG16, we include a variant of Fishy that only approximates the local FIM for the ﬁnal fully-connected layers (called Fishy classiﬁer), letting earlier convolutional layers inherit gradients from the ﬁrst fully-connected layer’s predictive distribution.
Hyperparameters for each method are tuned over 400 trials with the Vizier Bayesian optimization toolbox [21, 7], and the trial with the best ﬁnal objective value is displayed.
All Shampoo and Fishy variants share the same tunable hyperparameter ranges.
We ﬁnd signiﬁcant improvements over standard Shampoo with Fishy or its variants on most of the tasks.
Fishy even outperforms FIM Shampoo on some problems.
FIM Shampoo can be considered Fishy applied to the output layer only.
We plan to open-source the code to reproduce these experiments.
Conclusion
We introduce Fishy, a layerwise Fisher construction method that allows model parallelism for approximating the FIM.
We apply our construction to Shampoo, which only requires simple adjustments but can improve performance on several problems.
In the future, we would like to perform more extensive experiments comparing K-FAC, derive a full expectation version of Fishy without requiring sampling and develop Fishy for the K-FAC approximation.
Further treatment of layerwise predictive distributions on their own can also yield interesting research directions