L EARNING R EPRESENTATION FROM N EURAL F ISHER K ERNEL WITH L OW - RANK A PPROXIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, we study the representation of neural networks from the view of kernels.
We ﬁrst deﬁne the Neural Fisher Kernel (NFK), which is the Fisher Kernel (Jaakkola and Haussler, 1998) applied to neural networks.
We show that NFK can be computed for both supervised and unsupervised learning models, which can serve as a uniﬁed tool for representation extraction.
Furthermore, we show that practical NFKs exhibit low-rank structures.
We then propose an efﬁcient algorithm that computes a low rank approximation of NFK, which scales to large datasets and networks.
We show that the low-rank approximation of NFKs derived from unsupervised generative models and supervised learning models gives rise to high-quality compact representations of data, achieving competitive results on a variety of machine learning tasks.
I NTRODUCTION
Modern deep learning systems rely on ﬁnding good representations of data.
For supervised learning models with feed forward neural networks, representations can naturally be equated with the activations of each layer.
Empirically, the community has developed a set of effective heuristics for representation extraction given a trained network.
For example, ResNets (He et al., 2016) trained on
Imagenet classiﬁcation yield intermediate layer representations that can beneﬁt downstream tasks such as object detection and semantic segmentation.
The logits layer of a trained neural network also captures rich correlations across classes which can be distilled to a weaker model (Knowledge
Distillation) (Hinton et al., 2015).
Despite empirical prevalence of using intermediate layer activations as data representation, it is far from being the optimal approach to representation extraction.
For supervised learning models, it remains a manual procedure that relies on trial and error to select the optimal layer from a pre-trained model to facilitate transfer learning scenario.
Similar observations also apply to unsupervised learning models including GANs (Goodfellow et al., 2014), VAEs (Kingma and Welling, 2013), as evident from recent studies (Chen et al., 2020a) that the quality of representation in generative models heavily also depends on the choice of layer from which we extract activations as features.
Furthermore, although that GANs and VAEs are known to be able to generate high-quality samples of the training distribution, there is no strong evidence that they encode explicit layerwise representations to similar quality as in supervised learning models, which implies that there does not exist a natural way to explicitly extract a representation from intermediate layer activations in unsupervisedly pre-trained generative models.
Additionally, layer activations alone do not sufﬁce to reach the full power of learned representations hidden in neural network models, as shown in recent works (Mu et al., 2020) that incorporating additional gradients-based features into representation leads to substantial improvement over solely using activations-based features.
In light of these constraints, we are interested in the question: is there a principled method for representation extraction beyond layer activations?
In this work, we turn to the kernel view of neural networks.
Recently, initiated by the Neural Tangent Kernel (NTK)
(Jacot et al., 2018) work, there have been growing interests in the kernel interpretation of neural networks.
It was shown that neural networks in the inﬁnite width regime are reduced to kernel regression with the induced NTK.
Our key intuition is that, the kernel machine induced by the neural network provides a powerful and principled way of investigating the non-linear feature transformation in neural networks using the linear feature space of the kernel.
Kernel machines provide drastically different representations than layer activations, where the knowledge of a neural network is instantiated by the induced kernel function over data
points. In this work, we propose to make use of the linear feature space of the kernel, associated with the pre-trained neural network model, as the data representation of interest.
To this end, we made novel contributions on both theoretical and empirical side, as summarized below. • We propose Neural Fisher Kernel (NFK) as a uniﬁed and principled kernel formulation for neural networks models in both supervised learning and unsupervised learning settings.
• We introduce a highly efﬁcient and scalable algorithm for low-rank kernel approximation of NFK, which allows us to obtain a compact yet informative feature embedding as the data representation. • We validate the effectiveness of proposed approach from NFK in unsupervised learning, semisupervised learning and supervised learning settings, showing that our method enjoys superior sample efﬁciency and representation quality.
P RELIMINARY AND R ELATED W ORKS
In this section, we present technical background and formalize the motivation.
We start by introducing the notion of data representation from the perspective of kernel methods, then introduce the connections between neural network models and kernel methods.
Notations Throughout this paper, we consider dataset with N data examples D ≡ { ( x i , y i ) } , we use p ( x ) to denote the probability density function for the data distribution and use p data ( x ) to denote the empirical data distribution from D .
Kernel Methods Kernel methods (Hofmann et al., 2008) have long been a staple of practical machine learning.
At their core, a kernel method relies on a kernel function which acts as a similarity function between different data examples in some feature space.
Here we consider positive deﬁnite kernels
K : X × X
→ R over a metric space X which deﬁnes a reproducing kernel Hilbert space H of function from
X to R , along with a mapping function ϕ : X → H , such that the kernel function can be decomposed into the inner product
[Equation]
Kernel methods aim to ﬁnd a predictive linear function f ( x ) =
(cid:104) f, ϕ ( x ) (cid:105) H in H , which gives label output prediction for each data point x ∈ X .
The kernel maps each data example x ∈ X to a linear feature space ϕ ( x ) , which is the data representation of interest.
Given dataset
D , the predictive model function f is typically estimated via Kernel Ridge Regression (KRR), (cid:98) f
= argmin f ∈H 1 N (cid:80) Ni =
( f ( x i ) − y i )
+ λ (cid:107) f (cid:107) 2 H .
Neural Networks and Kernel Methods
A long line of works (Neal, 1996; Williams, 1996; Roux and Bengio, 2007; Hazan and Jaakkola, 2015; Lee et al., 2018; de G. Matthews et al., 2018; Jacot et al., 2018; Chen and Xu, 2021; Geifman et al., 2020; Belkin et al., 2018; Ghorbani et al., 2020), have studied that many kernel formulations can be associated to neural networks, while most of them correspond to neural network where being ﬁxed kernels (e.g. Laplace kernel, Gaussian kernel) or only the last layer is trained, e.g., Conjugate Kernel (CK) (Daniely et al., 2016), also called as
NNGP kernel (Lee et al., 2018).
On the other hand, Neural Tangent Kernel (NTK) (Jacot et al., 2018) is a fundamentally different formulation corresponding to training the entire inﬁnitely wide neural network models.
Let f ( θ ; x ) denote a neural network function with parameters θ , then the empirical NTK is deﬁned as K ntk ( x , z ) =
(cid:104)∇ θ f
( θ ; x ) , ∇ θ f
( θ ; z ) (cid:105) .
(Jacot et al., 2018; Lee et al., 2018) showed that under the so-called NTK parametrization and other proper assumptios, the function f ( x ; θ ) learned by training the neural network model with gradient descent is equivalent to the function estimated via ridgeless KRR using K ntk as the kernel.
For ﬁnite-width neural networks, by taking ﬁrst-order Taylor expansion of funnction f around the θ , kernel regression under K ntk can be seen as linearized neural network model at parameter θ , suggesting that pre-trained neural network models can also be studied and approximated from the perspective of kernel methods.
Fisher Kernel
The Fisher Kernel (FK) is ﬁrst introduced in the seminal work (Jaakkola and
Haussler, 1998). Given a probabilistic generative model p θ ( x ) , the Fisher kernel is deﬁned as:
K [Equation]
1 U z where U x = ∇ θ log p θ ( x ) is the so-called Fisher score and
I is the Fisher Information Matrix (FIM) deﬁned as the covariance of the Fisher score: E x ∼ p θ ( x ) ∇ θ log p θ ( x ) ∇ θ log p θ ( x ) (cid:62) .
Then the Fisher vector is deﬁned as
V x = I − 12 ∇ θ log p θ ( x ) =
I − 12 U x .
One can utilize the Fisher Score as the mapping from the data space
X to parameter space Θ , and obtain representations that are linearized.
As proposed [Figure] in (Jaakkola and Haussler, 1998; Perronnin and Dance, 2007), the Fisher vector V x can be used as the feature representation derived from probabilistic generative models, which was shown to be superior to hand-crafted visual descriptors in a variety of computer vision tasks. Generative Models
In this work, we consider a variety of representative deep generative models, including generative adversarial networks (GANs) (Goodfellow et al., 2014), variational autoencoders (VAEs) (Kingma and Welling, 2013), as well we normalizing ﬂow models (Dinh et al., 2015) and auto-regressive models (van den Oord et al., 2016).
Please refer to (Goodfellow et al., 2016) for more technical details on generative models.
L EARNING R EPRESENTATION FROM N EURAL F ISHER K ERNEL
We aim to propose a general and efﬁcient method for extracting high-quality representation from pre-trained neural network models.
As formalized in previous section, we can describe the outline of our proposed approach as: given a pre-trained neural network model f ( x ; θ ) (either unsupervised generative model p ( x ; θ ) or supervised learning model p ( y | x ; θ ) ), with pre-trained weights θ , we adopt the kernel formulation K f induced by model f ( x ; θ ) and make use of the associated linear feature embedding ϕ ( x ) of the kernel K f as the feature representation of data x . We present an overview introduction to illustrate our approach in Figure.
At this point, however, there exist both theoretical difﬁculties and practical challenges which impede a straightforward application of our proposed approach.
On theoretical side, the NTK theory is only developed in supervised learning setting, and its extension to unsupervised learning is not established yet.
Though Fisher kernel is immediately applicable in unsupervised learning setting, deriving Fisher vector from supervised learning model p ( y | x ; θ ) can be tricky, which needs the log-density estimation of marginal distribution p θ ( x ) from p ( y | x ; θ ) .
Note that it is a drastically different problem from previous works (Achille et al., 2019) where Fisher kernel is applied to the joint distribution over p ( x , y ) .
On practical efﬁciency side, the dimensionality of the feature space associated with NTK or FK is same as the number of model parameters | θ | , which poses unmanageably high time and space complexity when it comes to modern large-scale neural network models.
Additionally, the size of the NTK scales quadratically with the number of classes in multi-class supervised learning setting, which gives rise to more efﬁciency concerns.
To address the kernel formulation issue, we propose Neural Fisher Kernel (NFK) in Sec.
3.1 as a uniﬁed kernel for both supervised and unsupervised learning models.
To tackle the efﬁciency challenge, we investigate the structural properties of proposed NFK and propose a highly scalable low-rank kernel approximation algorithm in Sec.
3.2 to extract compact low-dimensional feature representation from NFK.
3. N EURAL F ISHER K ERNEL
In this section, we propose Neural Fisher Kernel (NFK) as a principled and general kernel formulation for neural network models.
The key intuition is that we can extend classical Fisher kernel theory to unify the procedure of deriving Fisher vector from supervised learning models and unsupervised learning models by using Energy-based Model (EBM) formulation.
3.1. U NSUPERVISED NFK
We consider unsupervised probabilistic generative models p θ ( x )
= p ( x ; θ ) here.
Our proposed NFK formulation can be applied to all generative models with tractable evaluation (or approximation) of ∇ θ log p θ ( x ) .
GANs We consider the EBM formulation of GANs (Dai et al., 2017; Zhai et al., 2019; Che et al., 2020).
Given pre-trained GAN model, we use D ( x ; θ ) to denote the output of the discriminator D , and use G ( h ) to denote the output of generator G given latent code h ∼ p ( h ) .
As a brief recap,
GANs can be interpreted as an implementation of EBM training with a variational distribution, where we have the energy-function E ( x ; θ ) =
− D ( x ; θ ) .
Please refer to (Zhai et al., 2019; Che et al., 2020) for more details.
Thus we have the unnormalized density function p θ ( x ) ∝ e − E ( x ; θ ) given by the GAN model.
Following (Zhai et al., 2019), we can then derive the Fisher kernel K nfk and Fisher vector from standard GANs as shown below:
[Equation]
where x , z ∈ X , I =
E
h ∼ p
( h ) (cid:104)
U G ( h ) U (cid:62) G ( h ) (cid:105) .
Note that we use diagonal approximation of FIM throughout this work for the consideration of scalability.
VAEs Given a VAE model pre-trained via maximizing the variational lower-bound ELBO
L ELBO ( x ) ≡ E q ( h | x ) (cid:104) log p ( x , h ) q ( h | x ) (cid:105) , we can approximate the marginal log-likelihood log p θ ( θ ) by evaluating L ELBO ( x ) via Monte-Carlo estimations or importance sampling techniques (Burda et al., 2016).
Thus we have our NFK formulation as
[Equation]
where x , z ∈ X , I =
E x ∼ p θ ( x ) (cid:2)
U x U (cid:62) x (cid:3) .
Flow-based Models, Auto-Regressive Models For generative models with explicit exact data density modeling p θ ( x )
, we can simply apply the classical Fisher kernel formulation in Sec. 2.
3.1. S UPERVISED NFK
In the supervised learning setting, we consider conditional probabilistic models p θ ( y | x )
=
p ( y | x ; θ ) .
In particular, we focus on classiﬁcation problems where the conditional probability is parameterized by a softmax function over the logits output f ( x ; θ ) :
p θ ( y | x ) = exp( f y ( x ; θ )) / (cid:80) y exp( f y ( x ; θ )) , where y is a discrete label and f y ( x ; θ ) denotes y -th logit output.
We then borrow the idea from JEM (Grathwohl et al., 2019) and write out a joint energy function term over ( x , y ) as E ( x , y ; θ ) =
− f y ( x ; θ ) .
It is easy to see that joint energy yields exactly the same conditional probability, at the same time leading to a free energy function:
E ( x ; θ ) = − log (cid:80) y exp( f y ( x ; θ )) .
It essentially reframes a conditional distribution over y given x to an induced unconditional distribution over x , while maintaining the same conditional probability p θ ( y | x ) .
This allows us to write out the NFK formulation as:
[Equation]
where I =
E x ∼ p θ ( x ) (cid:2)
U x U (cid:62) x (cid:3) , and p θ ( x ) is the normalized density corresponding to the free energy E θ , which could be sampled from via Markov chain Monte Carlo (MCMC) algorithm.
In this work, we use empirical data distribution as practical approximation.
3. NFK WITH L OW -R ANK A PPROXIMATION
Fisher vector V x is the linear feature embedding ϕ ( x ) given by NFK K nfk ( x , z )
=
(cid:104) V x , V z (cid:105) for neural network model f ( x ; θ ) .
However, straightforward application of NFK by using V x as feature representation suffers from scalability issue, since V x ∈ R | θ | shares same dimensionality as the number of parameters | θ | .
It is with that in mind that | θ | can be tremendously large considering the scale of modern neural networks, it is unfortunately infeasible to directly leverage V x as feature representation.
Low-Rank Structure of NFK Motivated by the Manifold Hypothesis of Data
that it is widely believed that real world high dimensional data lives in a low dimensional manifold (Roweis and Saul, 2000; Rifai et al., 2011a;b), we investigate the structure of NFKs and present empirical evidence that
NFKs of good models have low-rank spectral structure .
Firstly, we start by examining supervised learning models.
We study the spectrum structure of the empirical NFK of trained neural networks with different architectures.
We trained a LeNet-5 (LeCun et al., 1998) CNN and a 3-layer MLP network by minimizing binary cross entropy loss, and then compute the eigen-decomposition of the
NFK Gram matrix.
We show the explained ratio plot in Figure 2.
We see that the spectrum of CNN
NTK concentrates on fewer large eigenvalues, thus exhibits a lower effective-rank structure compared to the MLP, which can be explained by the fact that CNN has better model inductive bias for image data domain.
For unsupervised learning models, we trained a small unconditional DCGAN (Radford et al., 2015) model on MNIST dataset.
We compare the results of a fully trained model against a randomly initialized model in Fig.
2 (note the logarithm scale of the x axis).
Remarkably, the trained model demonstrates an extreme degree of low-rankness that top 100 principle components explain over 99 .
5% of the overall variance, where 100 is two orders of magnitude smaller than both number of examples and number of parameters in the discriminator.
We include more experimental results and discussions in appendix due to the space constraints.
The theoretical insights and empirical evidence presented above hint at a natural solution to address the challenge of high-dimensionality of
V x ∈ R | θ | : we can turn to seek a low-rank approximation to the NFK.
According to the Mercer’s theorem (Mercer, 1909), for positive deﬁnite kernel K ( x , z ) = (cid:104) ϕ
( x ) , ϕ ( z ) (cid:105) we have
K ( x , z ) = (cid:80) ∞
i =
λ i φ i ( x ) φ i ( z ) , x , z ∈ X , where { ( λ i , φ i ) } are the eigenvalues and eigenfunctions of the kernel K , with respect to the integral operator (cid:82) p ( z ) K ( x , z ) φ i ( z )
d z = λ
i φ i ( x ) .
The linear feature embedding representation ϕ
( x ) can thus be constructed from the orthonormal eigenbasis { ( λ i , φ i ) } as ϕ ( x ) ≡ (cid:2) √ λ 1 φ 1 ( x ) , √ λ 2 φ 2 ( x ) , . . .
(cid:3) ≡ (cid:2) √ λ
i φ i ( x ) (cid:3) , i = 1 , . . .
, ∞ .
To obtain a low-rank approximation, we only keep top- k largest eigen-basis { ( λ i , φ i ) } ordered by corresponding eigenvalues λ i to form the low-rank k -dimensional feature embedding e ( x ) ∈ R k , k (cid:28) N, k (cid:28) | θ | e ( x ) ≡ (cid:104)(cid:112) λ 1 φ 1 ( x ) , (cid:112) λ
2 φ 2 ( x ) , . . .
(cid:112) λ
k φ k ( x ) (cid:105) ≡ (cid:104)(cid:112) λ
i φ i ( x ) (cid:105) , i = 1 , . . .
, k (4)
By applying our proposed NFK formulation K nfk to pre-trained neural network model f ( x ; θ ) , we can obtain a compact low-dimensional feature representation e nfk ( x ) ∈ R k in this way.
We call it the low-rank NFK feature embedding .
We then illustrate how to estimate the eigenvalues and eigenfunctions of NFK K nfk from data.
Given dataset D , the Gram matrix K ∈ R N × N of kernel K is deﬁned as K ( x i , x j ) =
K ( x i , x j ) .
We use X ≡ [ x i ] Ni =1 to denote the matrix of all data examples, and use φ i ( X ) ∈
R N to denote the concatenated vector of evaluting i -th eigenfunction φ i at all data examples.
Then by performing eigen-decomposition of the Gram matrix K =
Φ diag (Λ) Φ (cid:62) , the i -th eigenvector Φ i ∈ R N and eigenvalue Λ i can be seen as unbiased estimation of the i -th eigenfunction φ i and eigenvalue
λ i of the kernel K , evaluated at training data examples X , φ i ( X ) ≈ √ N Φ i , λ i ≈ 1 N Λ i .
Based on these estimations, we can thus approximate the eigenfunction φ i via the integral operator by Monte-Carlo estimation with empirical data distribution,
[Equation]
Given new test data example x (cid:63) , we can now approximate the eigenfunction evluation φ i ( x (cid:63) ) by the projection of kernel function evaluation results centered on training data examples K ( x (cid:63) , X ) ≡ [ K ( x (cid:63) , x j )]
Nj =1 onto the i -th eigenvector Φ i of kernel Gram matrix K .
We adopt this method as the baseline approach for low-rank approximation, and present the baseline algorithm description in
Alg. 1.
However, due to the fact that it demands explicit computation and manipulation of the Fisher vector matrix V ∈ R N ×| θ | and the Gram matrix K ∈ R N × N in Alg. 1, straightforward application of the baseline approach, as well as other off-the-shelf classical kernel approximation (Williams and Seeger, 2000; Rahimi and Recht, 2007) and SVD methods (Halko et al., 2011), are practically infeasible to scale to larger-scale machine learning settings, where both the number of data examples N and the number of model parameters | θ | can be extremely large.
To tackle the posed scalability issue, we propose a novel highly efﬁcient and scalable algorithm for computing low-rank approximation of NFK.
Given dataset D and model f
( x ; θ ) , We aim to compute the truncated SVD of the Fisher vector matrix V =
Φ diag (Σ) P (cid:62) , P ∈ R | θ |× k .
Based on the idea of power methods (Golub and Van der Vorst, 2000; Bathe, 1971) for ﬁnding leading top eigenvectors, we starting from a random vector v 0 and iteratively construct the sequence v t +1 =
VV (cid:62) v t (cid:107) VV (cid:62) v t (cid:107) .
By leveraging the special structure of V that it can be obtained from the Jacobian matrix J θ ( X ) ∈
R N ×| θ | up to element-wise linear transformation under the NFK formulation in Sec.
3, we can decompose each iterative step into a Jacobian Vector Product (JVP) and a Vector Jacobian Product (VJP).
With modern automatic-differentiation techniques, we can evaluate both JVP and VJP efﬁciently, which only requires the same order of computational costs of one vanilla backward-pass and forward-pass of neural networks respectively.
With computed truncated SVD results, we can approximate the projection term in Eq.
5 by K ( x , X ) (cid:62) Φ i
=
V x V (cid:62) Φ
i ≈ V x diag (Σ i ) P
i , which is again in the JVP form so that we can pre-compute and store the truncated SVD results and evaluate the eigenfunction of any test data example via one efﬁcient JVP forward-pass.
We describe our proposed algorithm brieﬂy in Alg. 2.
E XPERIMENTS
In this section we evaluate NFK in the following settings.
We ﬁrst evaluate the proposed low-rank kernel approximation algorithm (Sec. 3.2), in terms of both approximation accuracy and running time efﬁciency.
Next, we evaluate NFK on various representation learning tasks in both supervised, semi-supervised and unsupervised learning settings. [Figure]
Best viewed in color.
4. Q UALITY AND E FFICIENCY OF L OW - RANK NFK A PPROXIMATIONS
We implement our proposed low-rank kernel approximation algorithm in Jax (Bradbury et al., 2018) with distributed multi-GPU parallel computation support.
For the baseline methods for comparison, we ﬁrst compute the full kernel Gram matrix using the neural-tangets (Novak et al., 2019) library, and then use sklearn.decomposition.
TruncatedSVD to obtain the truncated SVD results.
All model and algorithm hyperparmeters are included in the Appendix.
Computational Costs We start by comparing running time costs of computing top NFK eigen-vectors via truncated SVD.
We use two models for the comparison, a DCGAN-like GAN model in (Zhai et al., 2019) and a Wide ResNet (WRN) with 40 layers and 2 times wider than original network (denoted as WRN- 40 - 2 ).
Please see appendix for the detailed description of hyper-parameters.
We observed that our proposed algorithm could achieve nearly linear time scaling, while the baseline method would not be able to handle more than 2 14 data examples as the memory usage and time complexity are too high to afford.
We also see in Fig.
3 that by utilizing multi-GPU parallelism, we achieved further speed-up which scales almost linearly with the number of GPUs.
We emphasize that given the number of desired eigenvectors, the time complexity of our method scales linearly with the number of data examples and the demanded memory usage remains constant with adequate data batch size, since explicit computation and storage of the full kernel matrix is never needed.
Approximation accuracy We investigate the approximation error of our proposed low-rank approximation method.
Since we didn’t introduce any additional approximations, our method shares the same approximation error bound with the existing randomized SVD algorithm (Martinsson and Tropp, 2020;
Halko et al., 2011) and would only expect differences compared to the baseline randomized
SVD algorithm up to numerical errors.
To evaluate the quality of the low-rank kernel approximation, we use LeNet-5 and compute its full NFK Gram matrix on MNIST dataset.
Please see appendix for detailed hyper-parameter setups.
We show in Fig.
3 the approximation errors of top- 128 eigenvalues along with corresponding explained variances.
We obtain less than 1 e − 8 absolute error and less than 1 e − 7 relative error in top eigen-modes which explains most of the data. [Table]
4. L OW - RANK NFK E MBEDDING AS D ATA R EPRESENTATIONS
In this section we evaluate NFK to answer the following: Q1 .
In line with the question raised in Sec. 1, how does our proposed low-rank NFK embedding differ from the intermediate layer activations for data representation?
Q2 .
How does the low-rank NFK embedding compare to simply using gradients (Jacobians) as data representation?
Q3 .
To what extent can the low-rank NFK embedding preserve the information in full Fisher vector?
Q4 .
Does the NFK embedding representation lead to better generalization performance in terms of better sample efﬁciency and faster adaptation?
We conduct comparative studies on different tasks to understand the NFK embedding and present empirical observations to answer these questions in following sections.
NFK Representations from Unsupervised Generative Models
In order to examine the effectiveness of the low-rank NFK embeddings as data representations in unsupervised learning setting, we consider GANs and VAEs as representative generative models and compute the low-rank NFK embeddings.
Then we adopt the linear probing protocol by training a linear classiﬁer on top of the obtained embeddings and report the classiﬁcation performance to quantify the quality of NFK data representation.
For GANs, we use the same pretrained GAN model from (Zhai et al., 2019) and reimplemented the AFV baseline.
For VAEs, we follow the same model architecture proposed in (Child, 2020).
We then apply the proposed truncated SVD algorithm with 10 power iterations to obtain the 256 dimensional embedding via projection.
We present our results on CIFAR-10 (Krizhevsky et al., 2009a) in
Table.
We use GAN-NFK-128d ( GAN-NFK-256d ) to denote the NFK embedding obtained from using top- 128 (top- 256 ) eigenvectors in our GAN model.
Our VAE models ( VAE-NFK-128d and
VAE-NFK-256d ) follow the same notations.
For VAE baselines, the method proposed in (Mu et al., 2020) combines both gradients features and activations-based features into one linear model, denoted as VAE in the table.
For GAN baselines, we ﬁrst consider using intermediate layer activations only as data representation, referred to as the GAN-Activations model.
We then consider using full
Fisher vector as representation, namely using the normalized gradients w.r.t all model parameters as features, denoted as the GAN-AFV model as proposed in (Zhai et al., 2019).
Moreover, we also compare our results against training whole neural network using data labels in a supervised learning way, denoted as GAN-Supervised model.
As shown in Table.
1, by contrasting against the baseline GAN-AFV from GAN-Activations , as well as validation in recent works (Zhai et al., 2019; Mu et al., 2020), gradients provide additional useful information beyond layer activations based features.
However, it would be impractical to use all gradients or full Fisher vector as representation when scaling up to large-scale neural network models.
For example, VAE (Child, 2020) has ∼ 40 M parameters, it would not be possible to apply the baseline methods directly.
Our proposed low-rank NFK embedding approach addressed this challenge by building low-dim vector representation from efﬁcient kernel approximation algorithm, making it possible to utilize all model parameters’ gradients information by embedding it into a lowdim vector, e.g. 256 -dimensional embedding in VAE-NFK-256d from ∼ 40 M parameters.
As our low-rank NFK embedding is obtained by linear projections of full Fisher vectors, it naturally provides answers for Q2 that the NFK embedding can be viewed as a compact yet informative representation containing information from all gradients.
We see from Table.
1 that by using top- 128 eigenvectors,
[Table]
(Ahn et al., 2019), numbers are from (Ahn et al., 2019).
[Table]
the low-rank NFK embedding is able to recover the performance of full ∼ 5 .
9 M -dimension Fisher vector, which provides positive evidence for Q3 that we can preserve most of the useful information in Fisher vector by taking the advantage of the low-rank structure of NFK spectrum.
NFK Representations for Semi-Supervised Learning We then test the low-rank NFK embeddings in the semi-supervised learning setting.
Following the standard semi-supervised learning benchmark settings (Berthelot et al., 2019;
Miyato et al., 2019; Laine and Aila, 2017; Sajjadi et al., 2016;
Tarvainen and Valpola, 2017), we evaluate our method on CIFAR-10 (Krizhevsky et al., 2009a) and
SVHN datasets (Krizhevsky et al., 2009b).
We treat most of the dataset as unlabeled data and use few examples as labeled data.
We use the same GAN model as the unsupervised learning setting above, and compute top-
128 eigenvectors using training dataset (labeled and unlabeled) to derive the 128 -dimensional NFK embedding.
Then we only use the labeled data to train a linear classiﬁer on top of the NFK embedding features, denoted as the NFK-128d model.
We vary the number of labeled training examples and report the results in Table.
2, in comparison with other baseline methods.
We see that NFK-128d achieves very competitive performance.
On CIFAR-10, NFK-128d is only outperformed by the state-of-the-art semi-supervised learning algorithm MixMatch (Berthelot et al., 2019), which also uses a stronger architecture than ours.
The results on SVHN are mixed though NFK-128d is competitive with the top performing approaches.
The results demonstrated the effectiveness of NFK embeddings from unsupervised generative models in semi-supervised learning, showing promising sample efﬁciency for Q4 .
NFK Representations for Knowledge Distillation
We next test the effectiveness of using low-rank
NFK embedding for knowledge distillation in the supervised learning setting.
We include more details of the distillation method in the Appendix.
Our experiments are conducted on CIFAR10, with a teacher set as the WRN-40-2 model and student being WRN-16-1 .
After training the teacher, we compute the low-rank approximation of NFK of the teacher model, using top- 20 eigenvectors.
We include more details about the distillation method setup in the Appendix.
Our results are reported in
Table.
We see that our method achieves superior results compared to other competitive baseline knowledge distillation methods, which mainly use the logits and activations from teacher network as distillation target.
C ONCLUSIONS
In this work, we propose a novel principled approach to representation extraction from pre-trained neural network models. We introduce NFK by extending the Fisher kernel to neural networks in both unsupervised learning and superevised learning settings, and propose a novel low-rank kernel approximation algorithm, which allows us to obtain a compact feature representation in a highly efﬁcient and scalable way.
A E XTENDED P RELIMINARIES
We extend Sec.
2 to introduce additional technical background and related work.
Kernel methods in Deep Learning.
Popularized by the NTK work Jacot et al.
(2018), there has been great interests in the deep learning community around the kernel view of neural networks.
In particular, several works have studied the low-rank structure of the NTK, including (Baratin et al., 2020; Papyan, 2020; Canatar et al., 2020), which demonstrate that empirical NTK demonstrates low-rankness and that encourages better generalization theoretically.
Our low-rank analysis of NFK shares a similar ﬂavor, but generalizes across supervised and unsupervised learning settings.
Besides, we make an explicit effort in proposing an efﬁcient implementation of the low-rank approximation, and demonstrate strong empirical performances.
Unsupervised/self supervised representation learning.
Unsupervised representation learning is an old idea in deep learning.
A large body of work is dedicated to designing better learning objectives (self supervised learning), including denoising (Vincent et al., 2010), contrastive learning (Oord et al., 2018;
Chen et al., 2020b; He et al., 2020), mutual information maximization (Hjelm et al., 2018) and other “pretext tasks" Jing and Tian (2020).
Our attempt falls into the same category of unsupervised representation learning, but differs in that we instead focus on effectively extracting information from a standard probabilistic model.
This makes our effort orthogonal to many of the related works, and can be easily plugged into different family of models.
Knowledge Distillation.
Knowledge distillation (KD) is generally concerned about the problem of supervising a student model with a teacher model (Hinton et al., 2015; Ba and Caruana, 2013).
The general form of KD is to directly match the statistics of one or a few layers (default is the logits).
Various works have studied the layer selection (Romero et al., 2014) or loss function design aspects (Ahn et al., 2019).
More closely related to our work is efforts that consider the second order statistics between examples, including (Tung and Mori, 2019; Tian et al., 2019).
NFKD differs in that we represent the teacher’s knowledge in the kernel space, which is directly tied to the kernel interpretation of neural networks which introduces different inductive biases than layerwise representations.
Neural Tangent Kernel
Recent advancements in the understanding of neural networks have shed light on the connection between neural network training and kernel methods.
In (Jacot et al., 2018), it is shown that one can use the Neural Tangent Kernel (NTK) to characterize the full training of a neural network using a kernel.
Let f ( θ ; x ) denote a neural network function with parameters θ .
The
NTK is deﬁned as follows:
[Equation]
where P θ is the probability distribution of the initialization of θ .
(Jacot et al., 2018) further demonstrates that in the large width regime, a neural network undergoing training under gradient descent essentially evolves as a linear model.
Let θ
denote the parameter values at initialization.
To determine how the function f t ( θ t ; x ) evolves, we may naively taylor expand the output around θ 0 : f t +1 ( θ t +1 ; x ) ≈ f t ( θ t ; x ) − η ∇ θ t f t ( θ t ; x ) (cid:62) ( θ t +1 − θ t ) .
As the weight updates are given by θ
[Equation]
mi =1 ∇ θ t L t ( x i ) , hence we have f t +1 ( θ t +1 ; x ) ≈ f t ( θ t ; x ) − η 1 N (cid:80) Ni =
1 K ntk ( x , x i ) ∇ f L t ( x i ) .
The signiﬁcance of the NTK stems from two ob- servations.
1) When suitably initialized, the NTK converges to a limit kernel when the width tends to inﬁnity lim
[Equation]
( x , z ) .
In that limit, the NTK remains frozen in its limit state throughout training.
B O N THE C ONNECTIONS BETWEEN NFK AND NTK
In Sec 3.1, we showed that our deﬁnition of NFK in the supervised learning setting bares great similarity to the NTK.
We provide more discussion here on the connections between NFK and NTK.
For the L 2 regression loss function, the empirical ﬁsher information reduces to I =
N (cid:80) Ni =1 ∇ θ f θ ( x ) ∇ θ f θ ( x ) (cid:62) .
Note that the ﬁsher information matrix I is give by a covariance matrix of J , while the NTK matrix is deﬁned as the Gram matrix of J , where J is the Jacobian matrix, implying they share the same spectrum, and that the NTK and the NFK share the same eigenvectors.
The addition of I − 1 in the deﬁnition of K nfk can be seen as a form of conditioning, facilitating fast convergence in all directions spanned by J . {+Algorithm 
+}
Our proposed method: compute low-rank NFK feature embedding
[Equation]
Equation 3 also has immediate connections to NTK.
In NTK, the kernel K ntk ( x , ¯ x ) ∈ R N × N is a matrix which measures the dot product of Jacobian for every pair of logits.
The NFK, on the other hand, reduces the Jacobian ∇ θ f θ ( x ) for each example x to a single vector of dimension n (i.e., size of θ ), weighted by the predicted probability of each class p θ ( y | x ) .
The other notable difference between NFK and NTK is the subtractive and normalize factors, represented by E x
[Equation]
(cid:80) y
p θ ( y | x ) ∇
[Equation]
f y θ ( x (cid:48) ) and I , respectively.
This distinction is related to the difference between Natural Gradient Descent (Amari, 1998) and gradient descent.
In a nutshell, our deﬁnition of
NFK in the supervised learning setting can be considered as a reduced version of NTK, with proper normalization.
These properties make NFK much more scalable w.r.t.
the number of classes, and also less sensitive to the scale of model’s parameters.
To better see this, we can deﬁne an “unnormalized" version of NFK as K u ( x , ¯ x ) =
[ (cid:80) y
p θ ( y | x ) ∇ θ
f y
θ x )] (cid:62) (cid:80) y
p θ ( y | ¯ x ) ∇ θ f y θ (¯ x ) .
It is easy to see that K u has the same rank as the original
NFK K , as I − 1 is full rank by deﬁnition.
We can then further rewrite it as
[Equation]
In words, the unnormalized version of NFK can be considered as a reduction of NTK, where the weights of each element is weighte by the predicted probability for the respective class.
If we further assume that the model of interest is well trained, as is often the case in knowledge distillation, we can approximate the K u as K y ∗ , ¯ y ∗ ntk ( x , ¯ x ) , where y ∗ = arg max y p
θ ( y | x ) and likewise for ¯ y ∗ .
This suggests that the unnormalized NFK can roughly viewd as a downsampled version of NTK.
As a result, we expect the unnormalized NFK (and hence the NFK) to exhibit similar low rank properties as demonstrated in the NTK literature.
On the low-rank structure of NTK Consider the NTK Gram matrix
K ntk ∈ R N × N of some network Given the dataset { x i } Ni =1 (for simplicity we assume a scalar output) and its eigen decomposition K ntk = (cid:80) mj
=
λ j u j u (cid:62) j .
Let f ∈ R N denote the concatenated outputs.
Under GD in the linear regime, the outputs f t evolves according to:
[Equation]
The updates f t +1 − f t projected onto the bases of the kernel therefore converge at different speeds, determined by the eigenvalues { λ j } .
Intuitively, a good kernel-data alignment means that the ∇ f L is spanned by a few eigenvectors with large corresponding eigenvalues, speeding up convergence and promoting generalization.
C N EURAL F ISHER K ERNEL WITH L OW -R ANK A PPROXIMATION C. N EURAL F ISHER K ERNEL F ORMULATION
We provide detailed derivations of the various NFK formulations presented in Section. NFK for Energy-based Models Consider an Energy-based Model (EBM) p θ ( x ) = exp( − E ( x ; θ )) Z ( θ ) , where E ( x ) is the energy function parametrized by θ and Z ( θ ) =
(cid:82) exp( − E ( x ; θ ))
d x is the
1: U
U set
2: Draw
Ω ∈
U 3: Ω =
matrix_jacobian product ( f θ
X , Ω , kernel ) (cid:46) Ω
=
J θ ( X ) Ω ∈ R M
×
U 4: for do 5: Ω
_matrix_product ( f θ
X
Ω kernel )
(cid:46) Ω =
J (cid:62) θ
( X ) Ω ∈ R N × U 6: Ω = matrix_jacobian_product
( f θ , X , Ω , kernel ) (cid:46) Ω
= J θ ( X ) Ω ∈ R M
×
U 7: Ω = qr_decomposition ( Ω ) 8: end for 9: B = jacobian_matrix_product ( f θ , X , Ω , kernel ) (cid:46) B
= J (cid:62) θ
( X ) Ω ∈ R N × U 10: P , Σ , Q (cid:62) svd ( B (cid:62)
11: P
12: Keep
K to obtain
P k , k , Q
13: Return P partition function, we could apply the Fisher kernel formulation to derive the Fisher score U x as
[Equation]
Then we can obtain the FIM I and the Fisher vector V x from above results, shown as below
[Equation]
NFK for GANs As introduced in Section 3, we consider the EBM formulation of GANs.
Given pre-trained GAN model, we use D ( x ; θ ) to denote the output of the discriminator D , and use G ( h ) to denote the output of generator G given latent code h ∼ p ( h ) .
Then we have the energy-function deﬁned as E ( x ; θ ) = − D ( x ; θ ) .
Based on the NFK formulation for EBMs, we can simply substitute
Input Neural network model f
Input Input
Input
E ( x ; θ ) = − D ( x ; θ ) into Eq.
9 and Eq.
10 and derive the NFK formulation for GANs as below
[Equation]
Note that we use diagonal approximation of FIM throughout this work for the consideration of scalability.
Also, since the generator of GANs is trained to match the distribution induced by the discriminator’s EBM from the perspective of variational training for GANs, we could use the samples generated by the generator to approximate x ∈ p θ ( x ) , which is reﬂected in above formulation.
NFK for VAEs, Flow-based Models, Auto-Regressive Models For models including VAEs, Flowbased Models, Auto-Regressive Models, where explicit or approximate density estimation is available
, we can simply apply the classical Fisher kernel formulation as introduced in the main text.
NFK for Supervised Learning Models
In the supervised learning setting, we consider conditional probabilistic models p θ ( y | x ) =
p ( y | x ; θ ) .
In particular, we focus on classiﬁcation problems where the conditional probability is parameterized by a softmax function over the logits output f ( x ; θ ) :
p θ ( y | x ) = exp( f y ( x ; θ )) / (cid:80) y exp( f y ( x ; θ )) , where y is a discrete label and f y ( x ; θ ) denotes y -th logit output.
We then borrow the idea from JEM (Grathwohl et al., 2019) and write out a joint energy function term over ( x , y ) as E ( x , y ; θ ) =
− f y ( x ; θ ) .
It is easy to see that joint energy yields exactly the same conditional probability, at the same time leading to a free energy function:
[Equation]
Based on the NFK formulation for EBMs, we can simply substitute above results into Eq.
9 and
Eq.
10 and derive the NFK formulation for GANs as below [Equation] C. E FFICIENT LOW - RANK NFK/NTK APPROXIMATION VIA TRUNCATED SVD We provide mode details on experimental observations on the low-rank structure of NFK and the low-rank kernel approximation algorithm here.
[Figure]
Low-Rank Structure of NFK For supervised learning models
, we trained a LeNet-5 (LeCun et al., 1998) CNN and a 3-layer MLP network by minimizing binary cross entropy loss, and then compute the eigen-decomposition of the NFK Gram matrix.
For unsupervised learning models, we trained a small unconditional DCGAN (Radford et al., 2015) model on MNIST dataset.
We deliberately selected a small discriminator, which consists of 17 K parameters.
Because of the relatively lowdimensionality of θ in the discriminator, we were able to directly compute the Fisher Vectors for a random subset of the training dataset.
We then performed standard SVD on the gathered Fisher
Vector matrix, and examined the spectrum statistics.
In particular, we plot the explained variance ration quantity, deﬁned as r k = (cid:80) ki
=1 λ 2 i (cid:80) i =
λ 2 i where λ i is the i -th singular value.
In addition, we have also visualized the top 5 principle components, by showing example images which have the largest projections on each component in Fig.
Furthermore, we conducted a GAN inversion experiment.
We start by sampling a set of latent variables from the generator’s prior h ∈ p ( h ) , and get a set of generated example { x i } , x i = G ( h i ) , i = 1 , ..., n . We then apply Algorithm 2 on the generated example { x i } to obtain their NFK embeddings { e ( x i ) } , and we set the dimension of both h and e to 100.
We now have a compositional mapping that reads as h → x → e .
We then learn a linear mapping W ∈ R 100 × 100 from { e ( G ( h i )) } to { h i } by minimizing (cid:80) ni =1 (cid:107) h i − We ( G ( h i )) (cid:107) 2 . In doing so, we have constructed an auto encoder from a regular GAN, with the compositional mapping of x → e → h → ˜x , where ˜x is the reconstruction of an input x .
The reconstructions are shown in Figure 4 (a).
Interestingly, the 100d SVD embedding gives rise to a qualitatively faithful reconstruction on real images.
n contrast, a PCA embedding with the same dimension gives much more blurry reconstructions (eg., noise in the background), as shown in Figure 4 (b).
This is a good indication that the 100d embedding captures most of the information about an input example.
Power iteration of NFK as JVP/VJP evaluations Our proposed algorithm is based on the Power method Golub and Van der Vorst (2000); Bathe (1971) for ﬁnding the leading top eigenvectors of the real symmetric matrix.
[-Algorithm 
-]
Starting from a random vector v 0 drawn from a rotationally invariant distribution and normalize it to unit norm (cid:107) v 0 (cid:107)
= 1 , the power method iteratively constructs the sequence v t +1 =
Kv t (cid:107) Kv t (cid:107) up to q power iterations.
Given the special structure of K that it’s a Gram matrix of the Jacobian matrix J θ ( X ) ∈ R D × N , to evaluate Kv t in each power iteration step we need to evaluate J θ ( X ) (cid:62) J θ ( X ) v t , which can be decomposed as: (i) evaluating z t =
J θ ( X ) v t , and then (ii) Kv t =
J θ ( X ) (cid:62) z t .
Note that when K is in the form of NTK of neural networks, step (i) of evaluating z t is a Vector-Jacobian-Product (VJP) and step (ii) is a Jacobian-Vector-Product (JVP).
With the help of automatic-differentiation techniques, we can evaluate both JVP and VJP efﬁciently, which only requires the same order of computational costs of one backward-pass and forward-pass of neural networks respectively.
In this way, we can reduce the Kernel matrix vector product operation in each power iteration step to one VJP evaluation and one JVP evaluation, without the need to computing and storing the Jacobian matrix and kernel matrix explicitly.
As introduced in Section.
3.2, we include detailed algorithm description here, from Algorithm.to Algorithm.
In Algorithm.
3, we show the algorithm to compute the low-rank NFK embedding, which can be used as data representations.
In Algorithm.
4, we present our proposed automatic-differentiation based truncated SVD algorithm for kernel approximation.
Note that in
Algorithm.
5 and 6, we only need to follow Equation.
3 to pre-compute the model distribution statistics, Z θ = E x (cid:48) ∼ p θ ( x (cid:48) )
(cid:80) y
p θ ( y | x ) ∇ θ
f
y
θ
( x (cid:48) ) , and FIM I =
E x (cid:48) ∼ p θ ( x (cid:48) )
[ U x (cid:48) U (cid:62) x (cid:48) ] .
We adopt the EBM formulation of classiﬁer f θ ( x ) then replace the Jacobian matrix J θ ( X ) with the Fisher [Figure] vector matrix V θ ( X ) =
diag( I ) −
( J θ ( X ) − Z θ ) .
Note that our proposed algorithm is also readily applicable to empirical NTK via replacing the FIM by the identity matrix.
D E XPERIMENTS SETUP D. QUALITY ANDEFFICIENCY OFLOW-RANKNFK APPROXIMATIONS
Experiments on Computational Cost We randomly sample N ∈ (cid:8) 2 k : 7 ≤ k ≤ 16 (cid:9) data examples from CIFAR-10 dataset, and compute top-
32 eigenvectors of the NFK Gram matrix ( R N × N ) by truncated SVD.
We use same number of power iterations ( 10 ) in baseline method and our algorithm.
We show in Fig.
3 the running time of SVD for both methods in terms of number of data examples N .
Experiments on Approximation Accuracy We randomly sample 10000 examples and compute top-
128 eigenvalues using both baseline methods and our proposed algorithm.
Speciﬁcally, we compute the full Gram matrix and perform eigen-decomposition to obtain baseline results.
For our implementation, we run 10 power iterations in randomized SVD.
D. N EURAL F ISHER K ERNEL D ISTILLATION
With the efﬁcient low-rank approximation of NFK, one can immediately obtain a compact representation of the kernel.
Namely, each example can be represented as a k dimension vector.
Essentially, we have achieved a form of kernel distillation, which is a useful technique on its own.
Furthermore, we can use Q as an generalized form for teacher student styled knowledge distillation (KD), as in (Hinton et al., 2015).
In standard KD, one obtain a teacher network (e.g., deep model) and use it to train a student network (e.g., a shallow model) with a distillation loss in the following format:
[Equation]
where L cls is a standard classiﬁcation loss (e.g., cross entropy) and L t is a teacher loss which forces the student network’s output f s to match that of the teacher f t .
We propose a straightforward extension of KD with NFK, where we modify the loss function to be:
[Equation]
where Q t ( x ) denotes the k dimensional embedding from the SVD of teacher NFK, for example x .
h s is a prediction head from the student, and L t is overloaded to denote a suitable loss (e.g., (cid:96) 2 distance or cosine distance).
Equation 16 essentially uses the low dimension embedding of the teacher’s NFK as supervision, inplace of the teacher’s logits.
There are arguable beneﬁts of using L nfkd over L kd .
For example, when the number of classes is small, the logit layer contains very little extra information (measured in number of bits) than the label alone, whereas Q t can still provide dense supervision to the student.
For the Neural Fisher Kernel Distillation (NFKD) experiments, we adopt the WideResNet40x2 (Zagoruyko and Komodakis, 2016) neural network as the teacher model.
We train another
WideResnet with 16 layers as the student model, and keep the width unchanged.
We run 10 power iterations to compute the SVD approximation of the NFK of the teacher model, to obtain the topeigenvectors and eigenvalues.
Then we train the student model with the additional NFKD distillation loss using mini-batch stochastic gradient descent, with 0 .
9 momentum, for 250 epochs.
The initial learning rate begins at 0 .
and we decay the learning rate by 0 .
1 at 150 -th epoch and decay again by 0 .
1 at 200 -th epoch.
We also show the linear probing accuracies on CIFAR10 by using different number of embedding dimensions in Figure.