Title: GRADIENTS AS FEATURES FOR DEEP REPRESENTATION LEARNING

Abstract: We address the challenging problem of deep representation learning -the efficient adaption of a pre-trained deep network to different tasks. Specifically, we propose to explore gradient-based features. These features are gradients of the model parameters with respect to a task-specific loss given an input sample. Our key innovation is the design of a linear model that incorporates both gradient and activation of the pre-trained network. We demonstrate that our model provides a local linear approximation to an underlying deep model, and discuss important theoretical insights. Moreover, we present an efficient algorithm for the training and inference of our model without computing the actual gradients. Our method is evaluated across a number of representation-learning tasks on several datasets and using different network architectures. Strong results are obtained in all settings, and are well-aligned with our theoretical insights 1 .

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 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, 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) ∈ R d . ω ∈ R d×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

EQUATION

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 ∇θ 2 Fθ ,ω (x) = ωT Jθ 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), ωT Jθ 2 (x)) for x, our linear model ĝ, hereby considered as a classifier for concreteness, takes the form

EQUATION

where w 1 ∈ R d×c are linear classifiers initialized from ω, w 2 ∈ R |θ2| are shared linear weights for gradient features, and |θ 2 | is the size of the parameter θ 2 . Both w 1 and w 2 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 w 2 is multiplied to each row of the feature matrix. Therefore, w 2 is shared for all output dimensions. Similar to standard linear classifiers, the output of ĝ 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 ĝ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 θ.

3.2. THEORETICAL INSIGHT:

The key insight is that our model provides a local linear approximation to F θ2,ω (x). This approximation comes from Eq. ( 2) -the crux of our approach. Importantly, our linear model is mathematically well motivated -it can be interpreted as the first-order Taylor expansion of F θ,ω with respect to its parameters (θ 2 , ω) around the point of ( θ2 , ω). More formally, we note that

F θ,ω (x) ≈ ωT fθ(x) + ωT Jθ 2 (x)(θ 2 -θ2 ) + (ω -ω) T fθ(x) = ω T fθ(x) + ωT Jθ 2 (x)(θ 2 -θ2 ) = ĝω,θ2-θ2 (x).

(3)

With ω = w 1 and θ 2 -θ2 = w 2 , Eq. ( 2) provides a linear approximation of the deep model F θ2,ω (x) around the initialization ( θ2 , ω). Our key intuition is that given a sufficiently strong base network, training our model approximates fine-tuning F θ2,ω .

The quality of the linear approximation can be theoretically analyzed via the recent neural tangent kernel approach (Jacot et al., 2018; Lee et al., 2019; Arora et al., 2019b) or some related ideas (Arora et al., 2019a; Li & Liang, 2018; Allen-Zhu et al., 2019a; Du et al., 2019; Allen-Zhu et al., 2019b; Cao & Gu, 2020; Mei et al., 2019) when the base network F θ,ω is sufficiently wide and at random initialization. Unlike prior works, we apply the linear approximation on pre-trained networks of practical sizes. We argue that such an approximation is useful in practice for the following reasons:

• The pre-trained network provides a strong starting point. Thus, the pre-trained network parameter θ is close to a good solution for the downstream task. The key to good linear approximation is that the network output is stable with respect to small changes in the network parameters. The pre-trained base network also has such stability properties, which are supported by empirical observations. For example, the pre-trained network has similar predictions for a significant fraction of data in the downstream task as a fine-tuned network. • The network width required for the linearization to hold decreases as data becomes more structured. An assumption made in existing analysis is that the network is sufficiently or even infinitely wide compared to the size of the dataset, so that the approximation can hold for any dataset. We argue that this is not necessary in practice, since the practical datasets are wellstructured, and theoretically it has been shown that as long as the trained network is sufficiently wide compared to the effective complexity determined by the structure of the data, then the approximation can hold (Li & Liang, 2018; Allen-Zhu et al., 2019a) . Our approach thus takes advantage of the bottom layers to reduce data complexity in the hope that linearization of the top (and often the widest) layers can be sufficiently accurate.

3.3. SCALABLE TRAINING:

Moving beyond the theoretical aspects, a practical challenge of our method is the high cost of computing ĝ during training and inference. A naïve approach requires evaluating and storing ωT Jθ 2 (x) for all x. This is computationally expensive and can become infeasible as the number of output dimensions c and the parameter size |θ 2 | grow. Inspired by Pearlmutter (1994), we design an efficient training and inference scheme for ĝ. Thanks to this scheme, the complexity of training our model using gradient features is of the same magnitude as training a linear classifier on network activation.

Central to our scalable approach is the inexpensive evaluation of the Jacobian-vector product (JVP) ωT Jθ 2 (x)w 2 , whose size is the same as c. First, we note that Fθ +rw2,ω (x) = Fθ ,ω (x) + r ωT Jθ 2 (x)w 2 + o(r 2 ) (4) by first-order Taylor expansion around a scalar r = 0. Rearrange and take the limit of r to 0, we get

EQUATION

which can be conveniently evaluated via forward-mode automatic differentiation. 

where ∂zc ∂r is the JVP coming from the upstream layer. When a nonlinearity is encountered, we have, using the ReLU function as an example,

EQUATION

where is the element-wise product, 1 is the element-wise indicator function and z c is the layer input. Other nonlinearities as well as pooling layers can be handled in the same spirit. For batch normalization, we fold them into their preceding convolutions.

Importantly, Eqs. ( 6) and ( 7) provide an efficient approach to evaluating the desired JVP in Eq. ( 5) by successively evaluating a set of JVPs on the fly. This chain of evaluations starts with the seed ∂z0 ∂r = 0, where z 0 is the output of the section of f parameterized by θ 1 and can be pre-computed. ωT Jθ 2 (x)w 2 can be computed along with the standard forward propagation through f . Moreover, during the training of ĝ, its parameters w 1 and w 2 can be learned via standard back-propagation. In summary, our approach only requires a single forward pass through the fixed f for evaluating ĝ, and a single backward pass for updating the parameters w 1 and w 2 .

Complexity Analysis. We further discuss the complexity of our method in training and inference, and contrast our method to the fine-tuning of network parameters θ 2 . Our forward pass, as demonstrated by Eqs. ( 6) and ( 7), is a chain of linear operations intertwined by element-wise multiplications. The second term in Eq. ( 6) forms the "main stream" of computation, while the first term merges into the main stream at every layer of the ConvNet f . The same reasoning holds for the backward pass. Overall, our method requires twice as many linear operations as fine-tuning θ 2 of the ConvNet. Note, however, that half of the linear operations by our method are slightly cheaper due to the absence of the bias term. Moreover, in the special case where θ 2 only includes the topmost layer of f , our method carries out the same number of operations as fine-tuning since the second term in Eq. ( 6) can be dropped. For memory consumption, our method requires to store an additional "copy" (linear weights) of the model parameters compared to fine-tuning. As the size of θ 2 is typically small, this minor increase of computing and memory cost puts our method on the same page as fine-tuning.

4. EXPERIMENTS:

Our experimental results are organized into two parts. We first perform ablation studies to understand the representation power of the gradient features. Next, we evaluate our method on three representation-learning tasks: learning deep generative models, self-supervised learning using a pretext task, and transfer learning from ImageNet.

We report results on several datasets and using different base networks to demonstrate the strength of our method. In all cases, our method achieves significant better performance than the standard logistic regressor on network activation. In the most interesting scenario where the pre-training task is similar to the target task, our method achieves comparable or even better performance than fine-tuning. Before we present our results, we summarize our implementation details.

Implementation Details. We adopt the NTK parametrization (Jacot et al., 2018) for θ 2 and fold batch normalization into their preceding convolutional layers prior to training our model. For the SVHN and CIFAR-10/100 experiments, We train the models for 80K iterations with initial learning rate 1e-3, halved every 20K iterations. For the VOC07 and COCO2014 experiments, we train the models for 50 epochs with initial learning rate 1e-3, halved every 20 epochs. All models are trained with the Adam optimizer (Kingma & Ba, 2015) with batch size 64, β 1 = 0.5, β 2 = 0.999 and weight decay 1e-6. In addition to Adam, we also use the SGD optimizer with weight decay 5e-5, momentum 0.9 and the same learning rate schedule for fine-tuning, and we report the better result between the two runs.

4.1. ABLATION STUDY:

We conduct ablation studies to address the following two questions: Does pre-training encourage more powerful gradient features? What is the optimal size of the gradient features? We answer the first question by comparing gradient features derived from various combinations of random and pre-trained network parameters, and the second by varying the number of layers that contribute to the gradient features.

Baselines. We compare our linear model in Eq. ( 2), referred to as the full model in the sequel, against two baselines. The first baseline g w1 (x), or the activation model, is the first term in the full model. It is a logistic regressor on network activation and is widely used for representation learning in practice. The second baseline ωT Jθ 2 (x)w 2 , or the gradient model, is the second term in the full model. It is linear in the gradient features and can serve as a linear classifier on its own. Moreover, we compare our method against fine-tuning in order to assess the validity of our theoretical insight.

Settings. For ablation study, we consider both unsupervised-and supervised-learning settings and report classification accuracy. In the unsupervised setting (Table 1 ), we transfer features learned from a deep generative model for image classification. Specifically, we train a BiGAN on CIFAR-10 ( Krizhevsky et al., 2009) and use its encoder as the base network for classification of the same dataset. In the supervised setting (Table 2 ), we consider fine-tuning ImageNet pre-trained models on a different classification tasks. Concretely, we use the PyTorch (Paszke et al., 2017) distribution of ImageNet pre-trained ResNet18 (He et al., 2016) as the base network for VOC07 (Everingham et al., 2010) object classification. For both settings, activation features are always output from the pre-trained base network. In contrast, gradient features are evaluated with respect to all possible combinations of random and pre-trained network parameters. In this way, we ensure that variation in performance of the full model can be solely attributed to the gradient features.

Results and Conclusions. Our results are shown in Table 1 (unsupervised) and Table 2 (supervised).:

We summarize our main conclusions from the ablation study.

• The discriminative power of the gradient features is a consequence of pre-training. In particular, the full model's performance gain is not a consequence of an increase in parameter size. The full model supplied with random gradients is no better than the activation baseline. • Gradient from the topmost layer of a pre-trained network suffices to ensure a reasonable performance gain. Further inflating the gradient features from bottom layers has diminishing returns, introducing little accuracy gain and sometimes hurting the performance. The results in Zinkevich et al. ( 2017) also suggest that gradient features from top layers are the most effective in terms of approximating the underlying network.

Table 1 : Ablation study in the unsupervised setting. We train a BiGAN on CIFAR-10 and use its encoder as the base network for classification of the same dataset. All results are reported using classification accuracy. We obtain the pre-trained θ1 and θ2 from the pre-trained base network, and ω by first training the activation baseline for the target task. The specifications of θ 1 , θ 2 and ω only affects gradient evaluation. Activation features are always output from the pre-trained base network. • Our results improve as the dataset and base network grow in complexity. Our method always outperforms the activation baseline by a significant margin when the gradient features are derived from pre-trained parameters. Moreover, our method consistently beats the gradient baseline and even fine-tuning on the more challenging VOC07 dataset when using a more complicated residual network.

θ 2 : conv5 θ 2 : conv4-5 θ 2 : conv3-

Remarks. We obtain the pre-trained θ = ( θ1 , θ2 ) from representation learning, and the pre-trained ω by first training the standard logistic regressor for the target task. We follow the same training procedure in later experiments. According to the ablation study, a good heuristic for our method is to use the gradients from the topmost convolutional layer (or residual block) of a pre-trained network.

A more principled strategy for selecting gradient-contributing layers is left for future work.

4.2. RESULTS ON REPRESENTATION LEARNING:

We now present results on three different representation-learning tasks: unsupervised learning using generative modeling, self-supervised learning using pretext tasks and transfer learning from pretrained models. For all experiments in this section, we contrast our proposed linear classifier (i.e., the full model in Eq. ( 2)) with the activation and gradient baselines as well as fine-tuning θ 2 for the target task. Following our conclusion from the ablation study, we make use of gradient features only from the topmost layer (or residual block) of a pre-trained base network.

Unsupervised Learning. Similar to our ablation study, we consider unsupervised representation learning using deep generative models. We train both BiGAN and VAE, and use their encoders as the pre-trained ConvNet f . Our BiGAN and VAE models follow the architecture and training setup from Dumoulin et al. ( 2016) and Berthelot et al. (2019) . We average-pool the ConvNet output for activation features, and obtain gradient features from the topmost convolutional layer. We train on the train split of CIFAR-10/100 and the extra split of SVHN, and report classification accuracy on their test splits.

Unsupervised Learning Results. Our results are summarized in Table 3 . Our model consistently improves the activation baseline by over 10% on all three datasets. Moreover, we observe good agreement of performance between our model and fine-tuning. Finally, we notice that our model does not have significant advantage over the gradient baseline. We think this is because the datasets and base networks used here are so simple that the gradient features alone can be equally effective.

Self-supervised Learning. Moving beyond unsupervised learning, we also consider self-supervised setting, where the representation is learned via a pretext task and pseudo-labels. In this case, we consider a ResNet50 pre-trained on the jigsaw self-supervising task on ImageNet provided by Goyal et al. (2019) . See Noroozi & Favaro (2016) for details on the jigsaw task. We average-pool the ConvNet output for activation features, and use gradient features from the last residual block. We train on the trainval split of VOC07 and the train split of COCO2014 for object classification, and report the mean average precision (mAP) scores on their respective test and val splits.

Self-supervised Learning Results. Our results are summarized in the top part of Table 4 . We again observe a significant improvement of up to 4.5% over the two baselines on both datasets. Our method also outperforms a competitive self-supervised learning method (Colorization) (Zhang et al., 2016) using the activation from the same backbone network (ResNet50) and pre-trained on the same dataset (ImageNet). We note that our method still lags behind fine-tuning in this setting with a large gap (-6% on VOC07 and -7% on COCO2014).

Transfer Learning. Finally, we consider the most widely used transfer learning setting, where the representation is learned by pre-training on a large scale labeled dataset, e.g., ImageNet. We experiment with an ImageNet pre-trained ResNet18 from PyTorch. We follow the same setting as self-supervised learning and report results on VOC07 and COCO2014 datasets.

Transfer Learning Results. Our results are presented in the bottom part of Table 4 . Not surprisingly, our results consistently outperforms the two baselines across both datasets. However, the gap between our method and gradient baseline is much smaller. To our surprise, our method outperforms fine-tuning on both datasets by a large margin. In comparison to fine-tuning, our method has a gain of +0.8% and +3.0% on VOC07 and the larger COCO2014 dataset.

4.3. FURTHER DISCUSSIONS:

We provide further discussions of our results. Specifically, we analyze the performance gap between out method and fine-tuning under three different settings. We then compare our method with learning a linear classifier only using gradient features. Moreover, we contrast our method with the recent effort of learning infinitely wide networks using exact NTK (Arora et al., 2019b) .

The performance gap between our method and fine-tuning. We argue that this gap is controlled by the degree of semantic overlap between the representation-learning and target tasks. The jigsaw pretext task, though proven effective for learning features, is still semantically distant to an image classification task, and hence results in the large performance gap we observe in the self-supervised setting. On the other hand, it is reasonable to expect classifiers for datasets with overlapping categories to share common semantics. This may explain why our method outperforms fine-tuning in the transfer-learning experiments, where the source and target datasets are conceptually close.

Comparing our method with training a linear classifier on gradient features alone. In general, both methods are considerably more powerful than the standard activation-based logistic regressor.

In the unsupervised setting, the gradient model performs on par with, or even better than the full model, which consumes both activation and gradient features (see Table 3 ). This raises the question of whether network activations are redundant in the presence of gradient features. We argue that this is not necessarily the case. In fact, as we demonstrate in the self-supervised and transfer-learning experiments (see Table 4 ), combining activation and gradient features is a formula for delivering the best performance on challenging representation-learning tasks (jigsaw pretext task) and datasets (VOC07 and COCO2014). It would be interesting to further characterize the "representation bias" induced by the two types of features.

Contrasting our method with learning infinitely wide networks from scratch. Arora et al. (2019b) recently proposed exact NTK computation algorithms for learning infinitely wide Con-vNets, and proved that such networks are equivalent to kernel regression using NTK. They obtained 66% accuracy on CIFAR-10 classification using a vanilla network and achieved an extra 11% boost after including customized network modules. Their approach, at its current stage, only supports full-batch learning and is limited to single output dimension. Our method, though not directly comparable with theirs, does not suffer from these limitations and delivers competitive results (e.g., we obtain 70.5% accuracy on CIFAR-10 using a four-layer BiGAN encoder as the base network, see Table 3 ). Hence, out work highlights another promising direction to capitalize on the insights generated by the theoretical analysis of NTK.

5. CONCLUSION:

We presented a novel method for deep representation learning. Specifically, given a pre-trained deep network, we explored as features the per-sample gradients of the network parameters relative to a task-specific loss, and constructed a linear model that combines network activation with the gradient features. We showed that our model can be very efficient in training and inference, and may be understood as a local linear approximation to an underlying deep model by an appeal to the neuraltangent-kernel (NTK) theory. We empirically demonstrated that the gradient features are highly discriminative for downstream tasks, and our method can significantly improve over the baseline method of representation learning across pre-training tasks, network architectures and datasets. We believe that our work offers a novel perspective to improving deep representation learning. Future research directions include using NTK as a distance metric for measuring sample similarity in lowshot classification and assessing model similarity in knowledge distillation.

:

Project webpage at http://pages.cs.wisc.edu/ ˜fmu/gradfeat20

Figures and Tables:

Figure 1: (a) An illustration of our parameterization. We consider a deep network F (x; θ, ω) ω T f θ (x) that consists of a ConvNet f θ with its parameters θ (θ 1 , θ 2 ) and linear weights ω. (b) An overview of our proposed model. Our model takes the activation f θ (x) and the gradient Jθ 2 (x) as input (see illustration in (a)), and learns linear weights w 1 (matrix) and w 2 (vector) for prediction. 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.

More precisely, let us consider the basic building block of f -convolutional layers. These layers are defined as a linear function h(z c ; w c , b c ) = w T c z c + b c , where w c and b c , which are part of θ, are the weight and bias respectively, and z c is the layer input. We denote the counterparts of w c and b c in w 2 as wc and bc , i.e., wc and bc are the linear weights applied to w c and b c . It can be shown that ∂h(z c ; w c + r wc , b c + r bc ) ∂r = h(z c ; wc , bc ) + h( ∂z c ∂r ; w c , 0),

Ablation study in the supervised setting. We use the PyTorch distribution of ImageNet pre-trained ResNet18 as the base network for VOC07 object classification. All results are reported using mean average precision (mAP). Predictions are based on a single center crop at test time.

Table content: <table><tr><td>5</td></tr></table>

Unsupervised Learning Results. We consider two base networks, namely the encoder of a BiGAN and that of a VAE, trained on three datasets. Our target task is image classification on the same datasets. All results are reported using classification accuracy.

Table content: <table><tr><td/><td/><td>SVHN</td><td>CIFAR-10</td><td>CIFAR-100</td></tr><tr><td/><td>activation</td><td>82.04</td><td>62.87</td><td>34.30</td></tr><tr><td>BiGAN</td><td>gradient</td><td>89.91</td><td>70.14</td><td>38.99</td></tr><tr><td>(θ 2 : conv5)</td><td>full</td><td>89.96</td><td>70.51</td><td>39.37</td></tr><tr><td/><td>fine-tuning</td><td>91.15</td><td>71.78</td><td>41.05</td></tr><tr><td/><td>activation</td><td>81.95</td><td>52.05</td><td>29.20</td></tr><tr><td>VAE</td><td>gradient</td><td>91.44</td><td>61.47</td><td>35.10</td></tr><tr><td>(θ 2 : conv8)</td><td>full</td><td>90.90</td><td>61.16</td><td>34.83</td></tr><tr><td/><td>fine-tuning</td><td>93.61</td><td>65.16</td><td>37.86</td></tr></table>

Self-supervised and Transfer Learning Results. We consider a ResNet50 pre-trained on the jigsaw pretext task, and a ResNet18 pre-trained on the ImageNet classification task. Our target task is VOC07 and COCO2014 object classification. All results are reported using mean average precision (mAP). Predictions are averaged over ten random crops at test time.

Table content: <table><tr><td/><td/><td>VOC07</td><td>COCO2014</td></tr><tr><td/><td>activation</td><td>57.83</td><td>42.10</td></tr><tr><td>Self-supervised</td><td>gradient</td><td>57.73</td><td>41.02</td></tr><tr><td>(Jigsaw)</td><td>full</td><td>61.70</td><td>45.62</td></tr><tr><td/><td>fine-tuning</td><td>67.88</td><td>52.66</td></tr><tr><td>Self-supervised (Colorization)</td><td>activation</td><td>52.30</td><td>n/a</td></tr><tr><td/><td>activation</td><td>83.59</td><td>62.98</td></tr><tr><td>Transfer</td><td>gradient</td><td>84.63</td><td>66.33</td></tr><tr><td>(ImageNet)</td><td>full</td><td>84.95</td><td>66.45</td></tr><tr><td/><td>fine-tuning</td><td>84.14</td><td>63.43</td></tr></table>

Peer review (iclr2020):
Summary: This paper proposes to use the gradients of specific layers of convolutional networks as features in a linearized model for transfer learning and fast adaptation. The method is theoretically backed by an appeal to the recently proposed neural tangent kernel and seems like it could be practically useful. Edit: Post rebuttal, I am somewhat satisfied by the authors* response and find their ablation study somewhat compelling hence am increasing my score to a weak accept. However, I*m not completely convinced that their choice of layers to use as gradient features is reasonable. If this paper does end up getting accepted, I*d very strongly advocate for being extremely clear as to 1) why the gradient features used are used, 2) a good choice of gradient features in practice 3) inclusion of all of the relevant ablation studies in the final version. I tend to weakly reject this paper currently despite liking the simplicitly and timeliness of the approach. Specifically, I would like further discussion of the choice of the layers to use as gradient features and an ablation study on supervised trained networks. Originality: I believe that this is one of the first papers to explicitly use the Taylor approximation (neural tangent kernel) in a transfer learning setting, making the approach timely and potentially practically useful. Significance: The approach is a quite nice merging of theoretical insights with a neat practical implementation. Although the main methodological advance is a straightforward application of the thinking behind Jacobian vector products, the method is well described and ought to be practically useful. However, I have a bit of a concern as to its practical necessity in comparison to simple fine-tuning of the final softmax layer (referred to in Table 1 as 	heta_2) on a new dataset. Clarity: While the approach described in Section 3 is quite generic – theoretically, the method should simply consist of training the final layer of the neural network (w_1) and weights from the Jacobian features around the Taylor expansion. However, the experimental approaches suggest that different layers were used in each experiment – see e.g. “[we] … compute our gradient based features from one, two, or all of the top-3 conv layers” (Section 4.2 for the BiGAN architectures) versus “our gradient based features are from the last 2 conv layers” (Table 2 for the AlexNet architectures). Why was the entire network (modulo the last feed forward layer) not simply used for the Jacobian features? Similarly, why was the entirety of the model (again modulo the last feedforward layer) not used for forwards model in the AlexNet experiments? Based on the ablation study in Section 4.1, the authors find that “it suffices to set the very top layer as 	heta_2 to enjoy a reasonably large performance gain.” When this is the case, it is a bit tough to distinguish the approach from standard final layer transferring approaches (Sharif Razavian et al, 2014) and indeed Bayesian final layer approaches that simply retrain an output layer (e.g. Perrone et al, 2018). Indeed this seems to be the case in Table 1, if that is so, then why not just re-train 	heta_2 as well throughout the experiments as it will typically just be another stochastic convex optimization problem? Could the authors quickly describe an implementation of the scalable Jacobian inner products in Section 3.3? It seems like outputs will have to be cached during the forwards passes, thereby requiring a somewhat significant amount of software engineering (and memory overhead) to be have to do this process for each new layer. Is this the correct understanding of how one would implement the procedure? A quick paste of PyTorch pseudo-code would be sufficient here. Quality: The experiments seem to be relatively convincing – it seems exciting that a linear model + the network’s features itself can typically perform as well as fine-tuning and occasionally even better than fine-tuning itself. However, I am a bit concerned by the fact that the ablation studies themselves only utilize models trained in an un-supervised fashion. I’d instead like the authors to run an ablation study on supervised trained models (perhaps 8 layer conv nets or VGG16) in the same manner as in Table 1 and Section 4.1. Specifically, I’d like to see this done to see whether the gradients in fact are as interesting as features when they have been trained in a supervised fashion. Similarly, I’d like to see the features themselves used as a linear classifier (no network forwards passed) in the same two ablation studies. That is, could the authors use w_1^T J w_2 as the features for their linear classifier. If they have already done so and I’ve missed that in the tables somehow, I apologize. This experiment should help to test out how _useful_ the features defined by the Jacobian matrix are in comparison to the network’s forward pass itself. Minor Comments: Introduction: “After learning, the the (sic) activation of the deep network are considered as generic features.” Not only is there a small typo, but you should either include a citation here or be more specific as to what “the activation” of the deep network is here. Introduction: “And the accuracy of the rthe (sic) …” typo + please do not start sentences with and if at all possible. Section 3.1 (beneath eq. 2): “liner” should be linear. Table 3: “Self-supervise” should be “Self-supervised.” References: Perrone et al, Scalable Hyperparameter Transfer Learning, NeurIPS, 2018. http://papers.nips.cc/paper/7917-scalable-hyperparameter-transfer-learning.pdf Sharif Razavian et al, CNN Features off-the-Shelf: an Astounding Baseline for Image Recognition, CVPRW, 2014. https://arxiv.org/abs/1403.6382