Year: 2022

Abstract:
In this paper, we study the representation of neural networks from the view of kernels. We first define 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 unified tool for representation extraction. Furthermore, we show that practical NFKs exhibit low-rank structures. We then propose an efficient 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.

heading: 1 INTRODUCTION
text: Modern deep learning systems rely on finding 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 classification yield intermediate layer representations that can benefit 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. Similar observations also apply to unsupervised learning models including GANs (Goodfellow et al., 2014), VAEs (Kingma and Welling, 2014), as evident from recent studies (Chen et al., 2020a) that the quality of representation in generative models heavily 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 from the data 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 suffice 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 infinite 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 unified and principled kernel formulation for neural networks models in both supervised learning and unsupervised learning settings.
• We introduce a highly efficient 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 efficiency and representation quality.

heading: 2 PRELIMINARY AND RELATED WORKS
text: 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 ≡ {(xi, yi)}, we use p(x) to denote the probability density function for the data distribution and use pdata(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 definite kernels K : X × X → R over a metric space X which defines 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 K(x, z) = ⟨φ(x), φ(z)⟩. Kernel methods aim to find a predictive linear function f(x) = ⟨f, φ(x)⟩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), f̂ = argminf∈H 1 N ∑N i=1 (f (xi)− yi) 2 + λ∥f∥2H.
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 fixed 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 infinitely wide neural network models. Let f(θ;x) denote a neural network function with parameters θ, then the empirical NTK is defined as Kntk(x, z) = ⟨∇θf(θ;x),∇θf(θ; z)⟩. (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 Kntk as the kernel. For finite-width neural networks, by taking first-order Taylor expansion of funnction f around the θ, kernel regression under Kntk 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 first introduced in the seminal work (Jaakkola and Haussler, 1998). Given a probabilistic generative model pθ(x), the Fisher kernel is defined as: Kfisher(x, z) = ∇θ log pθ(x)⊤I−1∇θ log pθ(z) = U⊤x I−1Uz where Ux = ∇θ log pθ(x) is the so-called Fisher score and I is the Fisher Information Matrix (FIM) defined as the covariance of the Fisher score: I = Ex∼pθ(x)∇θ log pθ(x)∇θ log pθ(x)⊤. Then the Fisher vector is defined as Vx = I− 1 2∇θ log pθ(x) = I− 1 2Ux. One can utilize the Fisher Score as a mapping from the
data space X to parameter space Θ, and obtain representations that are linearized. As proposed in (Jaakkola and Haussler, 1998; Perronnin and Dance, 2007), the Fisher vector Vx 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, 2014), as well we normalizing flow models (Dinh et al., 2015) and auto-regressive models (van den Oord et al., 2016). Please refer to (Salakhutdinov, 2014) for more technical details on generative models.

heading: 3 LEARNING REPRESENTATION FROM NEURAL FISHER KERNEL
text: We aim to propose a general and efficient 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 Kf induced by model f(x;θ) and make use of the associated linear feature embedding φ(x) of the kernel Kf as the feature representation of data x. We present an overview introduction to illustrate our approach in Figure. 1.
At this point, however, there exist both theoretical difficulties and practical challenges which impede a straightforward application of our proposed approach. On the 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 the practical efficiency 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 efficiency concerns.
To address the kernel formulation issue, we propose Neural Fisher Kernel (NFK) in Sec. 3.1 as a unified kernel for both supervised and unsupervised learning models. To tackle the efficiency challenge, we investigate the structural properties of the 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.

heading: 3.1 NEURAL FISHER KERNEL
text: 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.

heading: 3.1.1 UNSUPERVISED NFK
text: 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 Knfk and Fisher vector from standard GANs as shown below:
Knfk(x, z) = ⟨Vx, Vz⟩ Vx = (diag(I)− 1 2 )Ux Ux = ∇θD(x;θ)− Eh∼p(h)∇θD(G(h);θ) (1)
where x, z ∈ X , I = Eh∼p(h) [ UG(h)U ⊤ G(h) ] . 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 LELBO(x) ≡ Eq(h|x) [ log p(x,h)q(h|x) ] , we can approximate the marginal log-likelihood log pθ(θ) by evaluating LELBO(x) via Monte-Carlo estimations or importance sampling techniques (Burda et al., 2016). Thus we have our NFK formulation as
Knfk(x, z) = ⟨Vx, Vz⟩ Vx = (diag(I)− 1 2 )Ux Ux ≈ ∇θLELBO(x) (2)
where x, z ∈ X , I = Ex∼pθ(x) [ UxU ⊤ x ] .
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.

heading: 3.1.2 SUPERVISED NFK
text: In the supervised learning setting, we consider conditional probabilistic models pθ(y | x) = p(y | x;θ). In particular, we focus on classification problems where the conditional probability is parameterized by a softmax function over the logits output f(x;θ): pθ(y | x) = exp(fy(x;θ))/ ∑ y exp(f
y(x;θ)), where y is a discrete label and fy(x;θ) denotes y-th logit output. We then borrow the idea from JEM (Grathwohl et al., 2020) and write out a joint energy function term over (x,y) as E(x,y;θ) = −fy(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:
Algorithm 1 Baseline method: compute low-rank NFK feature embedding Input dataset D; pre-train NN model f(x;θ); NFK feature dimensionality k; test data input x⋆ Output low-rank NFK feature embedding enfk(x∗)
1: compute Fisher vector for all data examples V = [Vxi ] ∈ RN×|θ| 2: compute kernel Gram matrix K = VV⊤ ∈ RN×N 3: compute truncated eigen-decomposition K = Φdiag(Λ)Φ⊤,Φ ∈ RN×k 4: kernel function evaluations between x∗ and all data examples K(x⋆,X) ≡ [K(x⋆,xj)]Nj=1 5: obtain enfk(x∗) ∈ Rk via Eq. 5 and Eq. 4
E(x;θ) = − log∑y exp(fy(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:
Knfk(x, z) = ⟨Vx, Vz⟩ Vx = (diag(I)− 1 2 )Ux Ux = ∑ y pθ(y | x)∇θfy(x;θ)− Ex′∼pθ(x′) ∑ y pθ(y | x)∇θfy(x′;θ) (3)
where I = Ex∼pθ(x) [ UxU ⊤ x ] , 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.

heading: 3.2 NFK WITH LOW-RANK APPROXIMATION
text: Fisher vector Vx is the linear feature embedding φ(x) given by NFK Knfk(x, z) = ⟨Vx, Vz⟩ for neural network model f(x;θ). However, straightforward application of NFK by using Vx as feature representation suffers from scalability issue, since Vx ∈ 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 Vx 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 exhibiting 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., 2016) 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.
Efficient Low-Rank Approximation of NFK. The theoretical insights and empirical evidence presented above hint at a natural solution to address the challenge of high-dimensionality of Vx ∈ R|θ|: we can turn to seek a low-rank approximation to the NFK. According to the Mercer’s theorem (Mercer, 1909), for positive definite kernel K(x, z) = ⟨φ(x), φ(z)⟩ we have K(x, z) = ∑∞i=1 λ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 ∫ p(z)K(x, z)ϕi(z) dz = λiϕi(x). The linear feature embedding representation φ(x) can thus be constructed from the orthonormal eigenbasis {(λi, ϕi)} as φ(x) ≡ [√ λ1ϕ1(x), √ λ2ϕ2(x), . . . ] ≡ [√ λiϕi(x) ] , 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) ∈ Rk, k ≪ N, k ≪ |θ|
e(x) ≡ [√ λ1ϕ1(x), √ λ2ϕ2(x), . . . √ λkϕk(x) ] ≡ [√ λiϕi(x) ] , i = 1, . . . , k (4)
Algorithm 2 Our proposed method: compute low-rank NFK feature embedding 1: V = Φdiag(Σ)P⊤,P ∈ R|θ|×k using power iteration methods via JVP/VJP evaluations 2: compute K (x,X)⊤ Φi ≈ Vxdiag(Σi)Pi via JVP evaluation 3: obtain enfk(x∗) ∈ Rk via Eq. 5 and Eq. 4
By applying our proposed NFK formulation Knfk to pre-trained neural network model f(x;θ), we can obtain a compact low-dimensional feature representation enfk(x) ∈ Rk in this way. We call it the low-rank NFK feature embedding.
We then illustrate how to estimate the eigenvalues and eigenfunctions of NFK Knfk from data. Given dataset D, the Gram matrix K ∈ RN×N of kernel K is defined as K(xi,xj) = K(xi,xj). We use X ≡ [xi]Ni=1 to denote the matrix of all data examples, and use ϕi(X) ∈ RN to denote the concatenated vector of evaluating i-th eigenfunction ϕi at all data examples. Then by performing eigen-decomposition of the Gram matrix K = Φdiag(Λ)Φ⊤, the i-th eigenvector Φi ∈ RN 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 ≈ 1NΛi. Based on these estimations, we can thus approximate the eigenfunction ϕi via the integral operator by Monte-Carlo estimation with empirical data distribution,
λiϕi(x) = ∫ p(z)K(x, z)ϕi(z) dz ≈ Exj∈pdataK(x,xj)ϕj(xj) ≈ 1
N N∑ j=1 K(x,xj)Φji (5)
Given new test data example x⋆, we can now approximate the eigenfunction evaluation ϕi(x⋆) by the projection of kernel function evaluation results centered on training data examples K(x⋆,X) ≡ [K(x⋆,xj)]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 ∈ RN×|θ| and the Gram matrix K ∈ RN×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 efficient 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⊤,P ∈ R|θ|×k. Based on the idea of power methods (Golub and Van der Vorst, 2000; Bathe, 1971) for finding leading top eigenvectors, we start from a random vector v0 and iteratively construct the sequence vt+1 = VV
⊤vt ∥VV⊤vt∥ . By leveraging
the special structure of V that it can be obtained from the Jacobian matrix Jθ(X) ∈ RN×|θ| 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 efficiently, 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)⊤ Φi = VxV⊤Φi ≈ Vxdiag(Σi)Pi, 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 efficient JVP forward-pass. We describe our proposed algorithm briefly in Alg. 2.

heading: 4 EXPERIMENTS
text: In this section, we evaluate NFK in the following settings. We first evaluate the proposed low-rank kernel approximation algorithm (Sec. 3.2), in terms of both approximation accuracy and running time efficiency. Next, we evaluate NFK on various representation learning tasks in both supervised, semi-supervised and unsupervised learning settings.

heading: 4.1 QUALITY AND EFFICIENCY OF LOW-RANK NFK APPROXIMATIONS
text: 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 first compute the full kernel Gram matrix using the neural-tangets (Novak et al., 2020) library, and then use sklearn.decomposition.TruncatedSVD to obtain the truncated SVD results. All model and algorithm hyper-parameters are included in the Appendix.
Computational Costs. We start by comparing running time costs of computing top NFK eigenvectors 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 214 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 did not 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 1e− 8 absolute error and less than 1e− 7 relative error in top eigen-modes which explains most of the data.

heading: 4.2 LOW-RANK NFK EMBEDDING AS DATA REPRESENTATIONS
text: 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 efficiency 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 classifier on top of the obtained embeddings and report the classification 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. 1. 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 first 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 ∼ 40M 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 efficient 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 ∼ 40M 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,
the low-rank NFK embedding is able to recover the performance of full ∼ 5.9M -dimension Fisher vector, which provides positive evidence for Q3 that we can preserve most of the useful information in Fisher vector by taking 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 classifier 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 efficiency 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. 3. 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.

heading: 5 CONCLUSIONS
text: 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 efficient and scalable way.

heading: A EXTENDED PRELIMINARIES
text: 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., 2021; 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 flavor, but generalizes across supervised and unsupervised learning settings. Besides, we make an explicit effort in proposing an efficient 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 based methods (Hjelm et al., 2019; Poole et al., 2019; Zhang et al., 2020) 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, 2014). 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., 2015) 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., 2020). 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 defined as follows:
Kntk(x, z) = Eθ∼Pθ ⟨∇θf(θ;x),∇θf(θ; z)⟩ . (6) 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 θ0 denote the parameter values at initialization. To determine how the function ft(θt;x) evolves, we may naively taylor expand the output around θ0: ft+1(θt+1;x) ≈ ft(θt;x) − η∇θtft(θt;x)⊤(θt+1 − θt). As the weight updates are given by θt−1 − θt = − 1N η ∑m i=1 ∇θtLt(xi), hence we have ft+1(θt+1;x) ≈
ft(θt;x)− η 1N ∑N
i=1 Kntk(x,xi)∇fLt(xi). The significance of the NTK stems from two observations. 1) When suitably initialized, the NTK converges to a limit kernel when the width tends to infinity limwidth→∞ Kntk(x, z;θ0) = K̊ntk(x, z). 2) In that limit, the NTK remains frozen in its limit state throughout training.

heading: B ON THE CONNECTIONS BETWEEN NFK AND NTK
text: In Sec 3.1, we showed that our definition 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 L2 regression loss function, the empirical fisher information reduces to I = 1 N ∑N i=1 ∇θfθ(x)∇θfθ(x)⊤. Note that the fisher information matrix I is give by a covariance matrix of J , while the NTK matrix is defined 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 definition of Knfk can be seen as a form of conditioning, facilitating fast convergence in all directions spanned by J .
Algorithm 3 Our proposed method: compute low-rank NFK feature embedding 1: V = Φdiag(Σ)P⊤,P ∈ R|θ|×k via truncated_svd(X, fθ, topk=K,
kernel_type="NFK") 2: 3: compute K (x,X)⊤ Φi ≈ Vxdiag(Σi)Pi via JVP evaluation 4: obtain enfk(x∗) ∈ Rk via Eq. 5 and Eq. 4
Equation 3 also has immediate connections to NTK. In NTK, the kernel Kntk(x, x̄) ∈ RN×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 Ex′∼pθ(x′) ∑ y pθ(y|x)∇θf y θ (x
′) and I, respectively. This distinction is related to the difference between Natural Gradient Descent (Amari, 1998; Karakida and Osawa, 2020) and gradient descent. In a nutshell, our definition 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 define an “unnormalized" version of NFK as Ku(x, x̄) = [ ∑
y pθ(y | x)∇θfyθ (x)]⊤ ∑ y pθ(y | x̄)∇θf y θ (x̄). It is easy to see that Ku has the same rank as the original NFK K, as I−1 is full rank by definition. We can then further rewrite it as Ku(x, x̄) = ∑ y ∑ ȳ pθ(y | x)pθ(ȳ|x̄)∇θfyθ (x)⊤∇θf ȳ θ (x̄) = ∑ y ∑ ȳ pθ(y | x)pθ(ȳ | x̄)Ky,ȳntk (x, x̄) (7) 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 Ku as Ky ∗,ȳ∗ ntk (x, x̄), where y ∗ = argmaxy 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 Kntk ∈ RN×N of some network Given the dataset {xi}Ni=1 (for simplicity we assume a scalar output) and its eigen decomposition Kntk = ∑m j=1 λjuju ⊤ j . Let f ∈ RN denote the concatenated outputs. Under GD in the linear regime, the outputs ft evolves according to:
∀j , u⊤j (ft+1 − ft) ≈ −ηλju⊤j ∇fL. (8)
The updates ft+1 − ft 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 ∇fL is spanned by a few eigenvectors with large corresponding eigenvalues, speeding up convergence and promoting generalization.

heading: C NEURAL FISHER KERNEL WITH LOW-RANK APPROXIMATION
text: C.1 NEURAL FISHER KERNEL FORMULATION
We provide detailed derivations of the various NFK formulations presented in Section. 3.
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(θ) = ∫ exp(−E(x;θ)) dx is the
Algorithm 4 truncated_svd, Truncated SVD Algorithm for Low-rank Kernel Approximation. Comments are based on NTK for simplicity.
Input Dataset X ≡ {xi}Ni=1 Input Neural network model fθ Input Kernel type kernel, NFK or NTK Input Low-rank embedding size K Input Number of power iterations L = 10 Input Number of over samples U = 10 Output Truncated SVD of Jacobian Jθ(X) ≈ PkΣkQ⊤k
1: U = K + U ▷ Size of augmented set of vectors in power iterations 2: Draw random matrix Ω ∈ RN×U 3: Ω =matrix_jacobian_product(fθ,X,Ω,kernel) ▷ Ω = Jθ(X)Ω ∈ RM×U 4: for step = 1 to L do 5: Ω =jacobian_matrix_product(fθ,X,Ω,kernel) ▷ Ω = J⊤θ (X)Ω ∈ RN×U 6: Ω =matrix_jacobian_product(fθ,X,Ω,kernel) ▷ Ω = Jθ(X)Ω ∈ RM×U 7: Ω =qr_decomposition(Ω) 8: end for 9: B =jacobian_matrix_product(fθ,X,Ω,kernel) ▷ B = J⊤θ (X)Ω ∈ RN×U
10: P,Σ,Q⊤ = svd(B⊤)

heading: 11: P = ΩP
text: 12: Keep top rank-K vectors to obtain the truncated results Pk,Σk,Q⊤k 13: Return Pk,Σk,Q⊤k
Algorithm 5 jacobian_matrix_product Input Neural network model fθ Input Input data X ∈ RB×D, where B is batch size Input Input matrix M Input Kernel type kernel, NFK or NTK Output J⊤θ (X)M for NTK, Fisher-vector-matrix-product V ⊤θ (X)M for NFK
1: jmp_fn = jax.vmap(jax.jvp) 2: P =jmp_fn(fθ,X,M) 3: if kernel = "NFK" then 4: P = diag(I)− 12 (P− Z⊤θ M) 5: end if 6: Return P
partition function, we could apply the Fisher kernel formulation to derive the Fisher score Ux as
Ux = ∇θ log pθ(x) = ∇θ log [exp(−E(x;θ))]−∇θ logZ(θ) = −∇θE(x;θ)−∇θ logZ(θ) = −∇θE(x;θ)− Ex∼pθ(x)∇θ log [exp(−E(x;θ))] = Ex∼pθ(x)∇θE(x;θ)−∇θE(x;θ)
(9)
Then we can obtain the FIM I and the Fisher vector Vx from above results, shown as below
I = Ex∼pθ(x) [ UxU ⊤ x ] Vx = I− 1 2Ux
(10)
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 defined as E(x;θ) = −D(x;θ). Based on the NFK formulation for EBMs, we can simply substitute
Algorithm 6 matrix_jacobian_product Input Neural network model fθ Input Input data X ∈ RB×D, where B is batch size Input Input matrix M Input Kernel type kernel, NFK or NTK Output Jθ(X)M for NTK, Fisher-vector-matrix-product Vθ(X)M for NFK
1: mjp_fn = jax.vmap(jax.vjp) 2: P =mjp_fn(fθ,X,M) 3: if kernel = "NFK" then 4: P = diag(I)− 12 (P− ZθM) 5: end if 6: Return P
E(x;θ) = −D(x;θ) into Eq. 9 and Eq. 10 and derive the NFK formulation for GANs as below Ux = ∇θD(x;θ)− Eh∼p(h)∇θD(G(h);θ) I = Eh∼p(h) [ UG(h)U ⊤ G(h)
] Vx = (diag(I)− 1 2 )Ux Knfk(x, z) = ⟨Vx, Vz⟩
(11)
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 reflected 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 classification problems where the conditional probability is parameterized by a softmax function over the logits output f(x;θ): pθ(y | x) = exp(fy(x;θ))/ ∑ y exp(f
y(x;θ)), where y is a discrete label and fy(x;θ) denotes y-th logit output. We then borrow the idea from JEM (Grathwohl et al., 2020) and write out a joint energy function term over (x,y) as E(x,y;θ) = −fy(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 ∑ y exp(fy(x;θ))
∇θE(x;θ) = − ∑ y pθ(y | x)∇θfy(x;θ) (12)
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
Ux = ∑ y pθ(y | x)∇θfy(x;θ)− Ex′∼pθ(x′) ∑ y pθ(y | x)∇θfy(x′;θ) (13)
I = Ex∼pθ(x) [ UxU ⊤ x ] Vx = (diag(I)− 1 2 )Ux
Knfk(x, z) = ⟨Vx, Vz⟩ (14)
C.2 EFFICIENT 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.
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., 2016) model on MNIST dataset. We deliberately selected a small discriminator, which consists of 17K parameters. Because of the relatively low-dimensionality 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, defined as rk = ∑k i=1 λ 2 i∑
i=1 λ 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. 6.
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 {xi},xi = G(hi), i = 1, ..., n. We then apply Algorithm 2 on the generated example {xi} to obtain their NFK embeddings {e(xi)}, 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 ∈ R100×100 from {e(G(hi))} to {hi} by minimizing ∑n i=1 ∥hi −We(G(hi))∥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 finding the leading top eigenvectors of
the real symmetric matrix. Starting from a random vector v0 drawn from a rotationally invariant distribution and normalize it to unit norm ∥v0∥ = 1, the power method iteratively constructs the sequence vt+1 = Kvt∥Kvt∥ up to q power iterations. Given the special structure of K that it’s a Gram matrix of the Jacobian matrix Jθ(X) ∈ RD×N , to evaluate Kvt in each power iteration step we need to evaluate Jθ(X)⊤Jθ(X)vt, which can be decomposed as: (i) evaluating zt = Jθ(X)vt, and then (ii) Kvt = Jθ(X)⊤zt. Note that when K is in the form of NTK of neural networks, step (i) of evaluating zt 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 efficiently, 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. 3 to Algorithm. 6. 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θ = Ex′∼pθ(x′) ∑ y pθ(y|x)∇θf y θ (x ′), and FIM I = Ex′∼pθ(x′)[Ux′U⊤x′ ]. We adopt
the EBM formulation of classifier fθ(x) then replace the Jacobian matrix Jθ(X) with the Fisher vector matrix Vθ(X) = diag(I)− 1 2 (Jθ(X)− Zθ). Note that our proposed algorithm is also readily applicable to empirical NTK via replacing the FIM by the identity matrix.

heading: D EXPERIMENTS SETUP
text: D.1 QUALITY ANDEFFICIENCY OFLOW-RANKNFK APPROXIMATIONS Experiments on Computational Cost. We randomly sample N ∈ { 2k : 7 ≤ k ≤ 16 } data examples from CIFAR-10 dataset, and compute top-32 eigenvectors of the NFK Gram matrix (RN×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. Specifically, 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.2 NEURAL FISHER KERNEL DISTILLATION
With the efficient 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:
Lkd(x,y) = α ∗ Lcls(fs(x),y) + (1− α) ∗ Lt(fs(x), ft(x)), (15) where Lcls is a standard classification loss (e.g., cross entropy) and Lt is a teacher loss which forces the student network’s output fs to match that of the teacher ft. We propose a straightforward extension of KD with NFK, where we modify the loss function to be:
Lnfkd(x,y) = α ∗ Lcls(fs(x),y) + (1− α) ∗ Lt(hs(x), Qt(x)), (16) where Qt(x) denotes the k dimensional embedding from the SVD of teacher NFK, for example x. hs is a prediction head from the student, and Lt is overloaded to denote a suitable loss (e.g., ℓ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 benefits of using Lnfkd over Lkd.
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 Qt 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 top-20 eigenvectors 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.1 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. 7.

Review summaries:
The paper investigates the representation of modern neural networks from the perspective of kernels by extracting features from pre-trained network models. The authors show the effectiveness of the proposed neural fisher kernel (NFK) for both unsupervised and supervised learning tasks. A low-rank approximation strategy of NFK is adopted to reduce computational burdens.

The paper proposes a new strategy for extracting compact and intuitive representations of the data from a neural network. The approach exploits the feature map associated with the kernel representation of a neural network. In particular, the authors consider the Neural Fisher Kernel and propose a series of linear approximations to make their approach scalable.