Gradients as Features for Deep Representation Learning

1 INTRODUCTION . Despite tremendous success of deep learning , training deep neural networks requires a massive amount of labeled data and computing resources . The recent development of representation learning holds great promises for improving data efficiency of training , and enables an easy adaption to different tasks using the same feature representation . These features can be learned via either unsupervised learning using deep generative models ( Kingma & Welling , 2014 ; Dumoulin et al. , 2016 ) , or self-supervised learning with “ pretext ” tasks and pseudo-labels ( Noroozi & Favaro , 2016 ; Zhang et al. , 2016 ; Gidaris et al. , 2018 ) , or transfer learning from another large-scale dataset ( Yosinski et al. , 2014 ; Oquab et al. , 2014 ; Girshick et al. , 2014 ) . After learning , the per-sample activation of the network is considered as generic features . By leveraging these features , simple classifiers such as linear models can be learned for different tasks . However , given sufficient amount of training data , the performance of representation-learning methods lags behind fully supervised deep models . As a step towards bridging this gap , we propose to make use of gradient-based features from a pretrained network , i.e. , gradients of the model parameters with respect to a task-specific loss given an input sample . Our key intuition is that these per-sample gradients contain task-relevant discriminative information . More importantly , we design a novel linear model that accounts for both gradientand activation-based features . The design of our linear model stems from the recent advances in the theoretical analysis of deep models . Specifically , our gradient-based features are inspired by the neural tangent kernel ( NTK ) ( Jacot et al. , 2018 ; Lee et al. , 2019 ; Arora et al. , 2019b ) , and adapt NTK in the setting of finite-width networks . Therefore , our model provides a local approximation of fine-tuning an underlying deep model , and the accuracy of the approximation is controlled by the semantic gap between the representation-learning and the target tasks . Finally , the structure of the gradient features and the linear model allows us to derive an efficient and scalable algorithm for training and inference . To evaluate our method , we focus on visual representation learning in this paper , although our model can be easily modified for natural language or speech data . To this end , we consider a number of learning tasks in vision , including unsupervised , self-supervised and transfer learning . Our method 1Project webpage at http : //pages.cs.wisc.edu/˜fmu/gradfeat20 is evaluated across tasks , datasets and network architectures and compared against a set of baseline methods . We observe empirically that our model with the gradient features outperforms the traditional activation-based logistic regressor by a significant margin in all settings . Moreover , our model compares favorably against fine-tuning of network parameters . Our main contributions are thus summarized as follows . • We propose a novel representation-learning method . At the core of our method lies in a linear model that builds on gradients of model parameters as the feature representation . • From a theoretical perspective , we claim that our linear model provides a local approximation of fine-tuning an underlying deep model . From a practical perspective , we devise an efficient and scalable algorithm for the training and inference of our method . • We demonstrate strong results of our method across various representation-learning tasks , different network architectures and several datasets . Furthermore , these empirical results are wellaligned with our theoretical insight . 2 RELATED WORK . Representation Learning . Learning good representation of data without expensive supervision remains a challenging problem . Representation learning using deep models has been recently explored . For example , different types of deep latent variable models ( Kingma & Welling , 2014 ; Higgins et al. , 2017 ; Berthelot et al. , 2019 ; Dumoulin et al. , 2016 ; Donahue et al. , 2016 ; Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ; Grathwohl et al. , 2019 ) were considered for representation learning . These models were designed to fit to the distribution of data , yet their intermediate responses were found useful for discriminative tasks . Another example is self-supervised learning . This paradigm seeks to learn from a discriminative pretext task whose supervision comes almost for free . These pretext tasks for images include predicting rotation angles ( Gidaris et al. , 2018 ) , solving jigsaw puzzles ( Noroozi & Favaro , 2016 ) and colorizing grayscale images ( Zhang et al. , 2016 ) . Finally , the idea of transfer learning hinges on the assumption that features learned from a large and generic dataset can be shared across closely related tasks and datasets ( Girshick et al. , 2014 ; Sharif Razavian et al. , 2014 ; Oquab et al. , 2014 ) . The most successful models for transfer learning so far are those pre-trained on the ImageNet classification task ( Yosinski et al. , 2014 ) . As opposed to proposing new representation-learning tasks , our work studies how to get the most out of the existing tasks . Hence , our method is broadly applicable – it offers a generic framework that can be readily combined with any representation-learning paradigm . Gradient of Deep Networks . Our method makes use of the Jacobian matrix of a deep network as feature representation for a downstream task . Gradient information is traditionally employed for visualizing and interpreting convolutional networks ( Simonyan et al. , 2013 ) , and more recently for generating adversarial samples ( Szegedy et al. , 2013 ) , crafting defense strategies ( Goodfellow et al. , 2015 ) , facilitating knowledge distillation ( Sinha et al. , 2018 ; Srinivas & Fleuret , 2018 ) , and boosting multi-task and meta learning ( Sinha et al. , 2018 ; Achille et al. , 2019 ) . Our work draws inspiration from Fisher vectors ( FVs ) ( Jaakkola & Haussler , 1999 ) – gradient-based features from a probabilistic model ( e.g. , mixture of Gaussians ) . The FV pipeline has demonstrated its success for visual recognition based on hand-crafted local descriptors ( Perronnin & Dance , 2007 ) . More recently , FVs have shown promising results with deep models , first as an ingredient of a hybrid system ( Perronnin & Larlus , 2015 ) , and then as task embeddings for meta-learning ( Achille et al. , 2019 ) . Our method differs from the FV approaches in two folds . First , it is not built around a probabilistic model , hence has distinct theoretical motivation as we describe later . Second , our method enjoys exact gradient computation with respect to network parameters and allows scalable training , whereas Achille et al . ( 2019 ) employs heuristics in their method to aggressively approximate the computation of FVs . Another line of relevant research is from Zinkevich et al . ( 2017 ) . They also explored similar gradient features of deep networks , which they call holographic features , by drawing inspiration from generalized linear models ( GLM ) . They showed that a linear model trained on the gradient features can faithfully recover the output of the original network for the same task , and also showed some other desired properties of gradient features . In contrast , our work is motivated by the NTK theory , and we study the more challenging problem of adapting the features for a different task . It would be interesting to investigate the connection between GLM and NTK . Neural Tangent Kernel ( NTK ) for Wide Networks . Jacot et al . ( 2018 ) established the connection between deep networks and kernel methods by introducing the neural tangent kernel ( NTK ) . Lee et al . ( 2019 ) further showed that a network evolves as a linear model in the infinite width limit when trained on certain losses under gradient descent . Similar ideas have been used to analyze wide networks ( see , e.g. , Arora et al . ( 2019b ; a ) ; Li & Liang ( 2018 ) ; Allen-Zhu et al . ( 2019a ) ; Du et al . ( 2019 ) ; Allen-Zhu et al . ( 2019b ) ; Cao & Gu ( 2020 ) ; Mei et al . ( 2019 ) ) . Our method is , to our best knowledge , the first attempt to materialize the theory in the regime of practical networks . Instead of assuming random initialization of network parameters as all the prior works do , we for the first time empirically evaluate the implication of pre-training on the linear approximation theory . 3 GRADIENT-BASED FEATURES FOR REPRESENTATION LEARNING . We start by introducing the setting of representation learning using deep models . Consider a feedforward deep neural network Fθ , ω ( x ) , ωT fθ ( x ) that consists of a backbone fθ ( x ) with its vectorized parameters θ and a linear model defined by ω ( italic for vectors and bold for matrix ) . Specifically , fθ encodes the input x into a vector representation fθ ( x ) ∈ Rd . ω ∈ Rd×c are thus linear weights that map a feature vector into c output dimensions . For this work , we focus on convolutional networks ( ConvNets ) for classification tasks . With trivial modifications , our method can easily extend beyond ConvNets and classification , e.g. , for a recurrent network as the backbone and/or for a regression task . Following the setting of representation learning , we assume that a pre-trained fθ̄ is given with θ̄ as the learned parameters . The term representation learning refers to a set of learning methods that do not make use of discriminative signals from the task of interest . For example , f can be the encoder of a deep generative model ( Kingma & Welling , 2014 ; Dumoulin et al. , 2016 ; Donahue et al. , 2016 ) , or a ConvNet learned by using proxy tasks ( self-supervised learning ) ( Goyal et al. , 2019 ; Kolesnikov et al. , 2019 ) or from another large-scale labeled dataset such as ImageNet ( Deng et al. , 2009 ) . Given a target task , it is common practice to regard fθ̄ as a fixed feature extractor ( activation-based features ) and train a set of linear weights , given by gω̄ ( x ) = ω̄ T fθ̄ ( x ) , ( 1 ) We omit the bias term for clarity . Note that ω̄ and θ̄ are instantiations of ω and θ , where ω̄ is the solution of the linear model and θ̄ is given by representation learning . Based on this setup , we now describe our method , discuss the theoretical implications and present an efficient training scheme . 3.1 METHOD OUTLINE . Our method assumes a partition of θ , ( θ1 , θ2 ) , where θ1 and θ2 parameterize the bottom and top layers of the ConvNet f ( see Figure 1 ( a ) for an illustration ) . Importantly , we propose to use gradient-based features ∇θ̄2Fθ̄ , ω̄ ( x ) = ω̄ TJθ̄2 ( x ) in addition to activation-based features fθ̄ ( x ) . Specifically , Jθ̄2 ( x ) ∈ R d×|θ2| is the Jacobian matrix of fθ̄ with respect to the pre-trained parameter θ̄2 from the top layers of f . Given the features ( fθ̄ ( x ) , ω̄ TJθ̄2 ( x ) ) for x , our linear model ĝ , hereby considered as a classifier for concreteness , takes the form ĝw1 , w2 ( x ) = w T 1 fθ̄ ( x ) + ω̄ TJθ̄2 ( x ) w2 = gw1 ( x ) + ω̄ TJθ̄2 ( x ) w2 , ( 2 ) where w1 ∈ Rd×c are linear classifiers initialized from ω̄ , w2 ∈ R|θ2| are shared linear weights for gradient features , and |θ2| is the size of the parameter θ2 . Both w1 and w2 are our model parameters that need to be learned from a target task . An overview of the model is shown in Figure 1 ( b ) . Our model subsumes the linear model in Eq . ( 1 ) as the first term , and includes a second term that is linear in the gradient-based features . We note that this extra linear term is different from standard linear classifiers as in Eq . ( 1 ) . In this case , the gradient-based features form a matrix and the weight vector w2 is multiplied to each row of the feature matrix . Therefore , w2 is shared for all output dimensions . Similar to standard linear classifiers , the output of ĝ is further normalized by the softmax function and trained with the cross-entropy loss using labeled data from the target dataset . Conceptually , our method can be summarized into three steps . • Pre-train the ConvNet fθ̄ . This is accomplished by substituting in any existing representationlearning algorithm . Examples include unsupervised , self-supervised and transfer learning . • Train linear classifiers ω̄ using fθ̄ ( x ) . This is a standard step in representation learning by fitting a linear model using “ generic features ” . • Learn the linear model ĝw1 , w2 ( x ) . A linear model in the form of Eq . ( 2 ) is learned using gradient- and activation-based features . Note that the features are obtained when θ = θ̄ is kept fixed , hence our method requires no extra tuning of the parameter θ̄ .

Summary: This paper considers the use of a neural network's Jacobian as additional features for semi-supervised, unsupervised, and transfer learning of representations. The idea is simple and the authors motivate this choice by connecting it to the literature on the neural tangent kernel (although nothing is proven in this paper). Some experiments are performed in which the authors demonstrate some improvements over the simple baseline of using the neural network's final layer alone.