Tuning Pre-trained Model via Moment Probing

Transfer learning aims to reuse the pre-trained models and re-adapting them to new tasks by fine-tuning them on a specific dataset. In both of natural language processing and computer vision communities, transferring pre-trained large models to downstream tasks has long been a popular paradigm [9, 2, 50, 4, 15, 39, 5, 7], which always provides rewarding performance with the total amount of pre-trained data and the scale of the model itself increase. During the transfer process, tuning strategies are limited to basic full fine-tuning and linear probing, suffering from low parameter efficiency and inferior performance. To solve above problem, recent works attempt to explore parameter-efficient methods by tuning a few of learnable parameters in the frozen backbone and sharing most of the parameters in the backbone for each downstream task [22, 21, 2, 32, 23, 6]. For example, some works selectively fine-tune the original parameters in the backbone where SpotTune [14] investigates which layers need to tune and Bitfit [2] finds that fine-tuning bias term is enough in some cases. Other works focus on inserting additional modules to adapt feature transfer. Among them, Adaptor-based methods [21, 24] insert an additional non-linear MLP block between frozen layers. VPT [23] adds learnable prompt tokens into the input space of frozen backbone. SSF [32] introduces both the scale and shift factors after each operation in backbone. AdaptFormer [6] proposes a parallel branch adding to the feed-forward network (FFN) of the pre-trained model. Different from aforementioned works, our MP makes the first attempt to explore the potential of LP, while achieving comparable or better performance than above counterparts at much lower training cost.

Second-order moment involved in our MP is closely related to global covariance pooling (GCP) [34], which has been studied to improve the representation and generalization abilities of deep architectures. Previous works [34, 33, 28, 29] have shown that GCP is an effective alternative to global average pooling (GAP) for visual tasks. For example, B-CNN [34] and MPN-COV [28] insert GCP with different post-normalization methods into deep convolutional neural networks (CNNs) for fine-grained and large-scale visual, respectively. Recently, DropCov [46] proposes an adaptive channel dropout method for GCP normalization, and shows the effectiveness in both deep CNNs and ViTs. Besides, some researches focus on reducing the dimension of GCP representations, resulting in several compact models [11, 26]. Different from the above GCP methods, our MP presents a multi-head convolutional cross-covariance (MHC³) to efficiently compute the second-order moment, which is designed for parameter-efficient tuning of pre-trained models and outperforms existing counterparts in terms of efficiency and effectiveness.

To this end, we propose a Moment Probing (MP) to generate stronger representations by exploiting probability distribution $P(\mathbf{X})$ of features $\mathbf{X}$ instead of simple mean point in $g(\mathbf{X})$, because $P(\mathbf{X})$ can characterize the full picture of $\mathbf{X}$. As such, our MP aims to perform a linear classifier on feature distribution $P(\mathbf{X})$:

$$\mathbf{y}_{pred}=\widehat{\mathbf{W}}P(\mathbf{X}).$$

(2)

However, $P(\mathbf{X})$ is usually unknown. Based on the classical probability theory [3], $P(\mathbf{X})$ can be defined by its characteristic function $\varphi_{\mathbf{X}}(t)$, according to Fourier transforms between $\varphi_{\mathbf{X}}(t)$ and probability density function $p(\mathbf{X})$ as

$$\varphi_{\mathbf{X}}(t)=E\left[e^{it\mathbf{X}}\right]=\int_{\mathbb{R}}e^{it\mathbf{x}}dP_{\mathbf{X}}(\mathbf{x})=\int_{\mathbb{R}}e^{it\mathbf{x}}p_{\mathbf{X}}(\mathbf{x})d\mathbf{x},$$

(3)

where $i$ is the imaginary unit, and $t\in\mathbb{R}$ is the argument of the characteristic function. Furthermore, we can rewrite $\varphi_{\mathbf{X}}(t)$ with the Taylor series of $e^{it\mathbf{X}}$ as:

	$\displaystyle\varphi_{\mathbf{X}}(t)$	$\displaystyle=E\left[\sum_{k=0}^{\infty}\frac{(it\mathbf{X})^{k}}{k!}\right]=\sum_{k=0}^{\infty}\frac{(it)^{k}}{k!}E\left[\mathbf{X}^{k}\right]$		(4)
		$\displaystyle=1+(it)E\left[\mathbf{X}\right]+\frac{(it)^{2}}{2!}E\left[\mathbf{X}^{2}\right]+\cdots,$

where $E\left[\mathbf{X}^{k}\right]$ indicates $k^{\text{th}}$-order moment ($\mathbf{M}_{k}$) of $\mathbf{X}$. Let the coefficient of $\mathbf{M}_{k}$ be $\omega_{k}$, we can rewrite Eqn. (4) as

$$\varphi_{\mathbf{X}}(t)=1+{\omega}_{1}\mathbf{M}_{1}+{\omega}_{2}{\mathbf{M}}_{2}+\cdots=1+\sum_{k=1}^{\infty}{\omega}_{k}\mathbf{M}_{k},$$

(5)

So far, we can represent feature distribution $P(\mathbf{X})$ by using its characteristic function as for Eqn. (5). By omitting the constant term, we reformulate Eqn. (2) as:

$$\mathbf{y}_{pred}=\widehat{\mathbf{W}}\left(\sum_{k=1}^{\infty}{\omega}_{k}\mathbf{M}_{k}\right)=\sum_{k=1}^{\infty}\left(\widehat{\mathbf{W}}_{k}\mathbf{M}_{k}\right),$$

(6)

where we can see that $\mathbf{y}_{pred}$ in Eqn. (6) fuses prediction results from different order moments, which are superior to single first-order statistics in the original LP. Although combination of more statistics generally leads to a more precise approximation on feature distribution, it brings much more computational cost and is sensitive to noise [48]. Therefore, our MP exploits the first- and second-order moments to approximate feature distribution for final prediction:

$$\mathbf{y}_{pred}=\sum_{k=1}^{2}\left(\widehat{\mathbf{W}}_{k}\mathbf{M}_{k}\right)=\widehat{\mathbf{W}}_{1}\mathbf{M}_{1}+\widehat{\mathbf{W}}_{2}\mathbf{M}_{2},$$

(7)

where $\mathbf{M}_{1}$ can be computed as classification token or average pooling of word tokens as for LP, while $\mathbf{M}_{2}$ is calculated by $\mathbf{X}^{T}\mathbf{X}$.

It can be seen that $\mathbf{Z}=\{\mathbf{Z}_{1};\mathbf{Z}_{2};\cdots;\mathbf{Z}_{h-1}\}$ only capture second-order statistics between adjacent split features. To perform interaction among all features and further reduce representation size, we introduce a parameter-efficient Local Representation Aggregation (LRA) block for $\mathbf{Z}$:

$$\rm LRA(\mathbf{Z})=\rm Concat[\rm Conv_{3\times 3}^{2}(\sigma(\rm Conv_{3\times 3}^{2}(\mathbf{Z})))],$$

(10)

where $\rm Conv_{3\times 3}^{2}$ indicates a $3\times 3$ convolution with the stride of 2, while input channel and output one of convolution are $h-1$. $\sigma$ is GELU non-linearity, and $\rm Concat$ is a concatenated operation. By using Eqn. (8)$\sim$Eqn. (10), we can compute the proposed ${\rm MHC^{3}}(\mathbf{X})$, which results in a $(h-1)(\hat{d}/4h)^{2}$-dimensional cross-covariance representation. Compared with $d^{2}$-dimensional $\mathbf{M}_{2}=\mathbf{X}^{T}\mathbf{X}$, size of our ${\rm MHC^{3}}(\mathbf{X})$ is at least $16h$ times less than the original second-order moment. Based on the proposed MHC³, our MP is finally computed by

$$\mathbf{y}_{pred}=\widehat{\mathbf{W}}_{1}\mathbf{M}_{1}+\widehat{\mathbf{W}}_{2}\rm MHC^{3}(\mathbf{X}).$$

(11)

To further explore the potential of MP, we investigate the effect of MP on learning intermediate features. To avoid tuning a number of parameters in large-scale pre-trained models, recently proposed parameter-efficient methods learn a few of additional parameters to recalibrate features from frozen backbones. Particularly, as shown in Figure 3 (a), SSF [32] introduces learnable parameters $\boldsymbol{\gamma}_{i}\in\mathbb{R}^{d_{i}}$ and $\boldsymbol{\beta}_{i}\in\mathbb{R}^{d_{i}}$ for scaling and shifting the output from each operation OP${}_{i}(\mathbf{X}_{i})$ (e.g., multi-head self-attention or feed forward network (FFN) ) of pre-trained model:

$$\mathbf{Y}_{i}=\boldsymbol{\gamma}_{i}\odot{\rm OP}_{i}(\mathbf{X}_{i})\oplus\boldsymbol{\beta}_{i},$$

(12)

where $d_{i}$ dimension $\mathbf{X}_{i}$ are inputs of $i$-operation OP_i and $\mathbf{Y}_{i}$ are recalibrated features. $\odot$ and $\oplus$ indicate element-wise multiplication and addition along dimension, respectively. As for Eqn. (12), the parameters $\boldsymbol{\gamma}_{i}$ and $\boldsymbol{\beta}_{i}$ are irrelevant to $\mathbf{X}_{i}$, neglecting effect of input. As shown in Figure 3 (b), AdaptFormer [6] designs an input-based tiny network to learn shifting parameters for feature recalibration in FFN, but it ignores the effect of scaling parameters. As suggested in SSF [32], both scaling and shifting operations help model transfer in downstream tasks.

Based on above analysis, we introduce two input-based tiny networks to learn both scaling and shifting parameters for feature recalibration. By considering parameter efficiency, we develop a partially shared module to learn two recalibrating parameters (PSRP). As shown in Figure 3 (c), our PSRP learns scaling ($\mathbf{w}_{l}$) and shifting ($\mathbf{b}_{l}$) parameters of $l^{\text{th}}$-layer outputs as

$$\begin{split}{\mathbf{w}_{l}}={\mathbf{W}}_{up}^{1}({\rm ReLU}({\mathbf{W}}_{down}(\mathbf{X}_{l}))),\\ {\mathbf{b}_{l}}={\mathbf{W}}_{up}^{2}({\rm ReLU}({\mathbf{W}}_{down}(\mathbf{X}_{l}))),\end{split}$$

(13)

where $\mathbf{X}_{l}\in\mathbb{R}^{d_{l}}$ are inputs of $l^{\text{th}}$-layer, while $\mathbf{W}_{up}^{1}\in\mathbb{R}^{d_{l}\times d_{h}}$, $\mathbf{W}_{up}^{2}\in\mathbb{R}^{d_{l}\times d_{h}}$ and $\mathbf{W}_{down}\in\mathbb{R}^{d_{h}\times d_{l}}$ are learnable parameters with hidden dimension of $d_{h}$. Finally, we obtain the recalibrated features by using $\mathbf{w}_{l}$ and $\mathbf{b}_{l}$ as

$$\mathbf{Y}_{l}=(\mathbf{w}_{l}+\mathbf{1})\odot{\rm FFN(LN}(\mathbf{X}_{l}))\oplus\mathbf{b}_{l},$$

(14)

where $\rm LN$ indicates layer normalization [1], and $\mathbf{1}$ is a vector of all ones.

To verify the effectiveness of our MP, we compare with several SOTA fine-tuning methods on seven downstream classification tasks. Specifically, we compare with SOTA methods on CIFAR-100 and FGVC datasets, where ViT-B/16 model pre-trained on ImageNet-21K is used as a basic backbone. Besides, we compare with SOTA methods using various backbone models on ImageNet-1K and CIFAR-100. Top-1 accuracy is used as evaluation metric while the best and second-best results for all methods are highlighted in red and blue, respectively.
Downstream Tasks. First, we compare our MP with SOTA methods on CIFAR-100 and FGVC datasets by using ViT-B/16 model pre-trained on ImageNet-21K. From Table 2 we can observe that (1) our MP consistently outperforms the original LP on all downstream tasks by a large margin, achieving 9.24% gains on average. We owe these performance gains to exploration of the stronger representations for linear classifier inherent in our MP. (2) Compared to full fine-tuning and other parameter-efficient methods, MP achieves comparable or even better performance at much lower training cost (see Figure 4), showing the potential of MP in tuning pre-trained models. (3) By considering the effect of MP on feature learning, MP₊ achieves further improvement and surpasses all compared tuning methods, verifying the effectiveness of MP on feature learning.

In this subsection, we make ablation studies to evaluate the effect of key components, including probing representations, feature dimension $d$ in MP and hidden dimension $d_{h}$ in PSRP. Besides, we compare MHC³ with state-of-the-art second-order representations. Here, we use ViT-B/16 pre-trained on ImageNet-21K as backbone and fine-tune it on ImageNet-1K.
Comparison of Different Probing Representations. To assess effect of representations on linear probing, we compare with several representations, including classification token (CLS_token), global average pooling (GAP) of word tokens, the original global covariance pooling (GCP) of word tokens and their combinations. As shown in Table 4(a), we can see that CLS_token (weighted combination of word tokens) is superior to GAP. Particularly, the original GCP brings no gain over CLS_token. In contrast, our MHC³ is slightly superior to CLS_token, while outperforming GCP by $\sim$2% with fewer parameters. It indicates that our MHC³ explores second-order moments for tuning pre-trained models in a more effective and efficient way. By combining CLS_token with GCP or MHC³, the performance will further increase. However, combination of CLS_token with GAP leads to inferior performance, which may be caused by both CLS_token and GAP are first-order statistics, suffering from weak complementary. Our MP explores both first- and second-order moment to achieve the best results, performing better than second-best method (CLS_token + GCP) by 0.6% with fewer parameters. Above results clearly demonstrate the effectiveness of representations in our MP for linear probing.
Effect of Dimension on MP and PSRP. For computing our MHC³, we first exploit a $1\times 1$ convolution layer to reduce dimension of features $\mathbf{X}$ from $d$ to $\hat{d}$ ($\hat{d}<d$). To evaluate effect of dimension $\hat{d}$, we experiment with MP by changing $\hat{d}$ from 128 to 768. As listed in Table 4(b), performance of MP consistently increases when $\hat{d}$ becomes larger. However, larger $\hat{d}$ will involve more parameters. To balance efficiency and effectiveness, we set $\hat{d}$ to 512 throughout all experiments. Furthermore, we assess effect of hidden dimension $d_{h}$ in Eqn. (13) on our PSRP, and compare with AdaptMLP [6]. As compared in Table 4(c), larger $d_{h}$ leads to better performance but more parameters for both AdaptMLP and PSRP. Meanwhile, our PSRP consistently outperforms AdaptMLP with few additional parameters, verifying the effectiveness of PSRP. In our work, $d_{h}$ is set to 16 for efficiency and effectiveness trade-off.
Comparison of Second-order Moment To evaluate the efficiency and effectiveness of MHC³ for computing second-order moment, we compare with two state-of-the-art second-order representations (i.e., B-CNN [34] and iSQRT-COV [28]) on ImageNet-1k, where ViT-B/16 pretrained on ImageNet-21K is used as the backbone. To achieve a good efficiency and effectiveness trade-off, we set feature dimension $\hat{d}$ to 128 for all compared second-order representations. As shown in Table 5, our MHC³ improves existing second-order representations more than 1.71% in Top-1 accuracy, while having fewer parameters. Furthermore, our MP is superior to all combinations of classification token (CLS_token) with other second-order representations in terms of both efficiency and effectiveness. These results above clearly demonstrate that MHC³ involved in our MP is more suitable for tuning pre-trained models.

To verify the robustness of our MP, we conduct experiments on three out-of-distribution (OOD) datasets, including ImageNet-A (IN-A) [20], ImageNet-R (IN-R) [18] and ImageNet-C (IN-C) [19]. Specifically, we first fine-tune ViT-Base/16 model pre-trained on ImageNet-21K by using ImageNet-1K (IN-1K), and directly perform inference on three OOD datasets without any training. Here we compare our MP/MP₊ with several SOTA methods in Table 6, where we can observe that our MP improves LP over about 3%$\sim$5% on three OOD datasets. Meanwhile, MP is very competitive to existing parameter-efficient methods at lower training cost. Particularly, full fine-tuning is high-performance on IN-1K, but it shows weak generalization to OOD datasets. On the contrary, our MP₊ shows high-performance on both IN-1K and OOD datasets. MP₊ achieves the best results on IN-A and IN-C, while being comparable to the recently proposed SSF on IN-R. Above results demonstrate our MP and MP₊ have the ability to equip pre-trained models with stronger robustness to OOD.

To verify the generalization of our MP, we further evaluate the effect of our MP by combining it with other parameter-efficient methods (ie, VPT [23], AdaptFormer [6], SSF [32]). Specifically, we keep the experimental settings in Sec. 4.4. As shown in Table 6, we can see that MP brings 0.86%, 0.58% and 0.46% improvement gains for VPT-Deep, AdaptFormer, and SSF in IN-1K, respectively. Besides, MP has an ability to improve the robustness of existing parameter-efficient methods to OOD settings. These results above verify the effectiveness and complementarity of our MP to existing parameter-efficient methods. Additionally, our MP₊ outperforms all combinations of MP with existing parameter-efficient methods, demonstrating the superiority of our partially shared module to learn two recalibrating parameters (PSRP) over existing parameter-efficient methods.

To explore the effect of different pre-training strategies, we conduct experiments by using various backbones under supervised (SUP [10]) and self-supervised (MAE [15], CLIP [39]) settings. For SUP setting, we evaluate our methods using five ViT models (i.e., ViT-Small/16, ViT-Small/32, ViT-Base/16, ViT-Base/32, ViT-Large/16) those are pre-trained on ImageNet-21K, and fine-tune them on ImageNet-1K. As shown in the left two columns of Figure 5, our MP consistently outperforms LP by a large margin. Besides, our MP obtains comparable or even better results with full fine-tuning, when model sizes increase. Particularly, MP (85.25%, 4.4M) and MP₊ (85.95%, 5.6M) outperform full fine-tuning (85.04%, 304.3M) by 0.21% and 0.71% with tuning much fewer parameters for ViT-L/16. For self-supervised objectives (i.e., MAE and CLIP), we use ViT-Base/16, ViT-Large/14, ViT-Large/16 and ViT-Huge/14 as basic backbones. As shown in the right two columns of Figure 5, MP still improves LP by a clear margin, while MP₊ brings further performance gains over MP. Since MAE aims to optimize parameters for image reconstruction, which suffers from a clear gap for classification task and may require to tune amount of parameters for performing model transfer. As such, full fine-tuning achieves the best results. For CLIP, our MP₊ achieves comparable results with full fine-tuning especially on large-size model, but it is more parameter-efficient. These results show our MP methods can generalize well to different pre-training strategies.

Finally, we evaluate the performance of our MP and MP₊ under few-shot setting. Following the setting in [39, 47], we conduct experiments on ImageNet-1K by selecting $\{1,2,4,8,16\}$ samples for each class as the training set, where ViT-B/16 model pre-trained on ImageNet-21K is used as the backbone. For each setting, we make three trials and report the average results in Figure 6. It can be seen that full fine-tuning is inferior to LP, which may be caused by that full fine-tuning requires to learn amount of parameters and is difficult to optimize on tiny/small training sets. Our MP is superior to LP for all cases, while MP₊ achieves further performance improvement by considering parameter-efficient feature learning. The comparisons above demonstrate that parameter-efficient designs for powerful representations and feature learning encourage our MP methods show good generalization to few-shot samples.