Learning Representation from Neural Fisher Kernel with Low-rank Approximation

1 INTRODUCTION . 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 . 2 PRELIMINARY AND RELATED WORKS . 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 . 3 LEARNING REPRESENTATION FROM NEURAL FISHER KERNEL . 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 .

Summary: 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.