ZooD: Exploiting Model Zoo for Out-of-Distribution Generalization
Abstract
Recent advances on large-scale pre-training have shown great potentials of leveraging a large set of Pre-Trained Models (PTMs) for improving Out-of-Distribution (OoD) generalization, for which the goal is to perform well on possible unseen domains after fine-tuning on multiple training domains.
However, maximally exploiting a zoo of PTMs is challenging since fine-tuning all possible combinations of PTMs is computationally prohibitive while accurate selection of PTMs requires tackling the possible data distribution shift for OoD tasks.
In this work, we propose ZooD, a paradigm for PTMs ranking and ensemble with feature selection.
Our proposed metric ranks PTMs by quantifying inter-class discriminability and inter-domain stability of the features extracted by the PTMs in a leave-one-domain-out cross-validation manner.
The top-K ranked models are then aggregated for the target OoD task.
To avoid accumulating noise induced by model ensemble, we propose an efficient variational EM algorithm to select informative features.
We evaluate our paradigm on a diverse model zoo consisting of 35 models for various OoD tasks and demonstrate: (i) model ranking is better correlated with fine-tuning ranking than previous methods and up to 9859x faster than brute-force fine-tuning; (ii) OoD generalization after model ensemble with feature selection outperforms the state-of-the-art methods and the accuracy on most challenging task DomainNet is improved from 46.5% to 50.6%.
Furthermore, we provide the fine-tuning results of 35 PTMs on 7 OoD datasets, hoping to help the research of model zoo and OoD generalization.
Code will be available at https://gitee.com/mindspore/models/tree/master/research/cv/ zood.
Introduction
Although data Independent and Identically Distributed (IID) is a primary assumption behind most machine learning systems, it does not hold in many real-world scenarios due to continuous distribution shifts [39, 88].
Machine learning models encounter serious performance degradation [9, 32, 34] in such Out-of-Distribution (OoD) scenarios.
To alleviate the accuracy degradation caused by distribution shifts, numerous algorithms have been proposed [3, 1, 40, 44, 5, 41, 70, 31, 20, 47, 6].
Recently,
Gulrajani and Lopez-Paz [29] have argued for the systematic comparisons of OoD algorithms and introduced a standard and rigorous test bed called DomainBed.
Their experimental comparison has raised concerns about the effectiveness of OoD algorithms since they often fail to outperform the simple Empirical Risk Minimization (ERM).
On the other hand, recent works [33, 2, 89, 64] have shown the advantages of pre-training for improving
OoD generalization, i.e., learning from multiple training domains in order to generalize to an unseen domain.
The availability of a large set of Pre-Trained Models (PTMs) provides a huge potential for solving various OoD tasks.
However, it is challenging to sufficiently exploit the power of a model zoo (a large set of PTMs).
One naive approach could be fine-tuning all possible combinations of PTMs on the target dataset and choosing the best-performing one, which is computationally expensive especially when the number of PTMs and the data size are large.
Besides, fine-tuning also over
Recently, many ranking metrics have been proposed to estimate the transferability of models under IID assumption [8, 76, 58, 91, 90].
However, ranking a zoo of models for generalization on unseen distribution shifts is more challenging compared with the IID setting.
Moreover, even if a metric can correctly evaluate the transferability of each PTM, simply using the best model will not fully utilize rich knowledge present in a zoo of models.
But the problem is even more serious that the most transferable model will include some noise, because noise and invariant features are undistinguishable in the sense that they are all stable across domains.
Previous study [87] also pointed this out and emphasized the necessity of feature denoising.
Therefore, if we leverage the model zoo by assembling relatively transferable models, the accumulation of noise features may increase memory use and hurt the predictive performance.
To solve the aforementioned problems, we propose ZooD, a paradigm to rank and aggregate a Zoo of
PTMs for OoD generalization.
An overview of our method is shown in Figure 1 .
Given a classification task with multiple training domains, to evaluate the generalization capability of each model, we quantify both the inter-class discriminability and inter-domain stability of the features extracted from each PTM in a leave-one-domain-out cross-validation manner, i.e., choosing one domain as the validation domain and each domain rotating as the validation domain, which is critical for identifying models that can extract domain-invariant features.
Each PTM in the zoo is ranked by this quantification.
ZooD then continues with model aggregation consisting of model ensemble and feature selection.
By introducing latent masks over candidate features, an efficient EM algorithm is proposed to select informative features.
To tackle the intractability of the posterior, variational approximation to the true posterior using a factorizable distribution is derived.
We further extend it to large-scale datasets by building a local estimator under the stochastic approximation [65].
To demonstrate the efficacy of our method, we have performed extensive experiments with 35 diverse
PTMs and 7 OoD datasets.
First, we show that our ranking metric is strongly correlated with the fine-tuning performance of PTMs compared with existing IID metrics.
Second, we illustrate the outstanding performance of ZooD on OoD datasets.
For instance, on Office-Home, we get 85.1% average accuracy compared with the previous SOTA of 70.6%.
Lastly, we show the speedup of our method compared with brute-force fine-tuning.
ZooD gives a maximum speedup of ≈ 10000 × (0.
GPU hours vs 2662.27 GPU hours), making it practical and scalable.
Finally, to speed up research and make our work more reproducible, we have devised a test bench consisting of extracted features, fine-tuning accuracy results, and ranking scores for all 35 PTMs in our model zoo.
This testbed can help future research as the process of getting fine-tuning accuracy results based on DomainBed [29] for a zoo of models is computationally expensive.
For instance, fine-tuning 35 models on all 7 OoD datasets costs approximately 35140 GPU hours (equivalent to 1464 GPU days or 4 GPU years).
Concisely, our contributions are as follows:
• We propose an efficient and scalable ranking metric to gauge the generalization-ability of PTMs for unseen domains.
• Using EM, we propose a method for selecting informative features and discarding invariant but noisy features in an ensemble of models.
• We have established a test bed for PTMs on 7 OoD datasets, including features extracted by 35 PTMs in our model zoo, fine-tuning accuracy results, and model ranking scores by different methods.
Related Work
Pre-training for OoD generalization.
To tackle the problem of distribution shifts between training and test data, various OoD methods [3, 1, 40, 44, 22, 13, 5, 41, 70, 20, 47, 6] have been proposed with the aim to learn invariant representations across different environments.
However, a standard evaluation [29] of many OoD algorithms shows that they do not significantly outperform simple ERM.
On the other hands, recent works have shown the effectiveness of pre-trained models for OoD generalization.
Yi et al.
[89] theoretically showed that adversarially pre-trained models also perform better for OoD generalization.
Yu et al.
[92] showed that the right choice of pre-trained models can achieve SOTA results.
They also showed that IID performance is not a good indicator of OoD performance and emphasized on the importanceofmodelselection.
Albuquerqueetal.[2]showedtheimportanceoffeatureextractorbyproposing a new OoD-based pretext task for self-supervised pre-training.
CLIP [64] demonstrated that large-scale pre-training on a dataset of image-text pairs results in much more robust models for downstream tasks with various distribution shifts.
Our work is based on these observations and we aim to facilitate utilization of PTMs by proposing an efficient metric as well as efficient feature ensemble and selection method.
Ranking pre-trained models by metric design.
Recently, a number of metrics have been introduced to estimate transferability of source-task-learned representations for target task under IID conditions.
H-score [8] estimates the transferability by finding the relationship between extracted features and target class labels.
NCE [76] proposes to estimate transferability via measuring conditional entropy between source and target labels.
LEEP [58] simplifies NCE by using the joint distribution of source and target labels to estimate log expected empirical prediction.
LogME [91, 90] estimates the maximum value of label evidence given features from pre-trained models.
These transferability metrics focus on determining the compatibility of source-task-learned representations for the target task.
We, on the other hand, aim to compute stability of these features across domains in addition to source-target transferability.
Ensemble and feature selection.
Early works have shown that model ensemble can significantly improve predictive performance [21].
In the age of deep learning, Lakshminarayanan et al.
[42] propose deep ensemble to measure predictive uncertainty.
Similar works [60, 62] on uncertainty estimation focus on the context of outlier detection and reinforcement learning.
When facing a zoo of PTMs, it’s natural to leverage the rich knowledge by assembling multiple PTMs.
In prior works,
Liu et al.
[49] propose using PTMs as teacher models that distill knowledge to a target model for downstream tasks.
Shu et al.
[71] propose Zoo-Tuning that learns to aggregate the parameters of multiple PTMs to a target model.
However, these methods require the target model must have the identical architecture as the PTMs, thus sacrificing flexibility.
Our proposed paradigm involves selecting informative features from assembled feature extractors.
In the framework of Bayesian variable selection, it is common practice to identify promising features by estimating the posterior probabilities over all potential feature subsets.
Here we mainly focus on Stochastic Search Variable Selection (SSVS) [59] that involves specifying priors over regression coefficients such that higher posterior probabilities will be allocated to coefficients substantially different from zero.
Then the features whose coefficients have higher posterior probabilities will be selected.
George and McCulloch [25] first propose SSVS for the linear model and conduct the posterior inference using Gibbs sampling [57].
Li and Zhang [45] consider SSVS for regression modeling in high-dimensional spaces incorporating structural information.
Roˇcková and George [66] propose
EMVS for efficient SSVS in high-dimentional cases with sparse estimations of posterior probabilities.
Note that all aforementioned feature selection methods have inherent assumptions that observed datasets must be IID, which makes these methods difficult to use in our scenarios.
ZooD for OoD Generalization 3.1 Model Ranking
Here α and β are hyper-parameters.
The goal is to inference the conditional distribution of y ′
j given ϕ ( x ′ j ) , y i , and ϕ ( x i ) .
Assume that we have a domain distribution D from which we observe m domains: (cid:8) D 1 , D 2 , ··· , D m (cid:9) .
Each domain
D i is a set of label and datum pairs, i.e. D i = (cid:8) ( y ij ,x ij ) , 1 ≤ j ≤ n i (cid:9) .
Meanwhile, we have a zoo of pre-trained feature extractors: M = { ϕ 1 ,ϕ 2 , ··· ,ϕ k , ···} .
Our objective is to select a feature extractor from M (e.g., ϕ k ), such that when we train a predictor f on top of it, the composed model f ◦ ϕ k can have the best performance on both the m observed domains and unseen domains from D .
In this section, we propose a method that facilitates model selection without carrying out the fine-tuning step.
For every model in M , our method produces an associated score, by which we can rank the models, such that the higher-ranked ones have a better chance to deliver stronger results after fine-tuning.
The proposed method is a combination of 1) a model transferability metric and 2) a leave-one-domain-out crossvalidation scheme.
More specifically, we evaluate each feature extractor m times, and each time we treat the data from the held-out domain as validation data { ( y ′ j ,x ′ j ) } n ′ j =1 , while aggregating all remaining ( m − 1) domains’ data as the training data { ( y i ,x i ) } ni =1 .
In the end, we average the m values of the model transferability metric.
Finally, we rank all feature extractors in descending order of the average.
To simplify the notation, we denote the aggregated domain’s label and feature as y =( y 1 ,...,y n ) ⊤ ∈ R n and Φ= (cid:0) ϕ
( x 1 ) ,...,ϕ ( x n ) (cid:1) ⊤ ∈ R n × d , respectively.
We use y ′ ∈
R n ′ and Φ ′ ∈ R n ′
× d for the held-out domain.
The main idea of the designed metric is to evaluate whether the classifier fitted on ( y , Φ) also performs well on ( y ′ , Φ ′ ) .
Hence, we formulate the problem as estimating the likelihood function of ( y ′ , Φ ′ ) given ( y , Φ) :
[Equation]
where p ( y ′ | Φ ′ , y , Φ) measures inter-class discriminability between features Φ ′ and labels y ′ , given the aggregated training data.
Meanwhile, p (Φ ′ | Φ) measures covariate shift between features Φ and Φ ′ , which quantify the inter-domain stability .
Given a hypothetical space F of classifiers, we can write p ( y | Φ)= (cid:82) f ∈F p ( y | Φ ,f ) p ( f )d f .
We consider a linear classifier 3 , i.e. f ◦ ϕ
( x )= w ⊤ ϕ ( x ) with a Gaussian prior of w :
[Equation]
where α and β are two positive parameters.
Figure 2 summarizes the model assumptions in (1) with a directed graphical model.
We estimate ˆ α and ˆ β by maximizing the model evidence
[Equation]
according to Algorithm 3 in You et al.
[90] and compute the likelihood of y ′ as follows:
[Equation]
For measuring covariate shift, we approximate the distribution of ϕ ( x ) with a Gaussian distribution
N (ˆ µ ϕ , ˆΣ ϕ ) , where ˆ µ ϕ and ˆΣ ϕ are estimated from the training data Φ .
Then we compute the density p (Φ ′ | Φ)= p (Φ ′ | ˆ µ ϕ , ˆΣ ϕ ) to quantify the covariate shift.
Finally, we compute the density at the logarithmic scale and this defines the proposed metric log p ( y ′ | Φ
[Equation]
, y , Φ)+log p (Φ ′ | Φ) .
Please refer to Appendix B.3 and B.4 for more details.
One distinctive aspect of our selection process is the cross-domain validation, embodied in the first term of (2).
Across different domains, there are domain-invariant and domain-specific features, where overfitting to the latter can severely harm the OoD generalization.
By evaluating on held-out domains, we are able to filter out models that fixate on domain-specific features.
To provide theoretical justification, an explicit analysis in the linear regression setting is conducted, where we show that the model with the optimal metric is the one that selects all domain-invariant features.
Despite the over-simplification, it does reflect the essence of our approach.
Due to page limit, the technical details are presented in Appendix B.5.
3.2 Model Ensemble with Feature Selection
The top-ranked PTMs in Section 3.1 are preferred for solving the OoD generalization task.
To further aggregate different PTMs, we consider assembling the top-ranked feature extractors and rewrite Φ= (cid:2) Φ (1) ,..., Φ ( k ) (cid:3) , where Φ ( i ) is the feature matrix from the i -th ranked feature extractor.
As we show in experiments, in most cases, aggregating features from multiple models can outperform any single model.
However, simply concatenating features inevitably introduces more noise.
As found in [87], non-informative but invariant features from training domains may only bring some noise that is irrelevant to the classification problem, and the accumulation of noise hurts the learnability of the
OoD generalization task while increasing the memory and computation cost.
Therefore, we propose a feature selection method under the Bayesian linear model framework in Section 3.1.
First, we impose a binary mask
z
=
( z 1 , z 2 ,..., z d ) ⊤ for the weight vector w = ( w 1 , w 2 ,..., w d ) ⊤ , where z i = 1 indicates that w i is an active weight in the top linear model, i.e., w i ̸ = 0 , meaning the corresponding feature is informative, while w i
≈
if z i =0 , indicating a noisy feature that should be screened.
Therefore the Bayesian feature selection is formulated by estimating the probability π i of z i with π i :=
p ( z i =1) and π = { π 1 ,π 2 ,...,π d } .
To facilitate the utility of the mask, we assume that the weights { w i } are independent of each other and each weight w i is drawn from either a slab prior or a spike prior [37] with the mean of zero:
[Equation]
We make the Bayesian treatment to the linear model in Section 3.1 by introducing gamma priors for all inverse variance terms:
[Equation]
and denote all hyper-parameters as ν = { ν i,j } .
In addition, we denote all latent variables as ξ = (cid:8) β, { w i , z i ,α i, 1 ,α i, 2 } di =1 (cid:9) .
Under certain conditions, maximizing marginal likelihood provably leads to consistent selection and obeys Occam’s razor phenomenon [27, 85], and thus screens non-informative features.
To estimate π i , the maximum marginal likelihood estimator of ( π , ν ) is given by
[Equation]
However, the direct maximization of (3) is intractable due to the integration over ξ .
One possible solution is to use EM algorithm [66].
In the E-step, we compute the conditional expectation:
[Equation]
Notice that evaluating the expectation involves the posterior distribution of ξ .
However in our case, it is not straightforward to obtain an analytical form of the true posterior distribution.
We instead
[Equation] [Equation]
[Equation] [Equation]
approximate it using variational inference [12] by introducing a tractable distribution
Q.
Considering the following objective function:
[Equation]
which is a lower bound of log p ( y | Φ; π , ν ) .
It has been shown the maximizer of L ( Q ) is the optimal approximator of p ( ξ | y , Φ; π , ν ) under the KL divergence.
To obtain an explicit solution, we consider the classical mean-field family [12], where variational distribution Q can be factorized into:
[Equation]
After all variational parameters in (4) are updated by running one-step coordinate gradient descent [12], in the M-step, we update π new and ν new by maximizing:
[Equation]
By repeating the E- and M-step, the estimator ( π new , ν new ) converges to an optimal solution.
We then screen those variables with converged prior π i smaller than the predefined threshold τ .
The pseudocode is provided in Algorithm 1 to illustrate the main idea of the proposed method, where ξ k denotes the k -th variable in the set ξ and ξ − k is the subset of all other variables except ξ k .
In the E-step, the optimal approximator Q ( ξ ) under the mean-field family takes the tractable form of the expectation of the joint distribution and the optimization of (5) in M-step is equivalent to substituting with corresponding variational parameters of Q ( ξ ) from E-step.
Our derivations for variational approximations and prior hyper-parameters optimization are listed in Appendix C.3.
However, the proposed algorithm still suffers from heavy computational cost: each iteration costs
O ( nd 2 ) .
To address this problem, we propose an efficient version based on stochastic variational inference [35].
A local estimator Q s ( ξ ) is established under stochastic approximation that enjoys less computational complexity and guarantees convergence to global optimum [65].
We successfully reduce the computation cost to O ( n s d 2 ) with n s ≪ n .
Readers can refer to Appendix C.4 for more detailed discussions and the complete algorithm for feature selection.
Experiments
In this section, we demonstrate the effectiveness of ZooD.
First, we evaluate the ability of our ranking metric to estimate OoD performance and compare it with the ground-truth performance and several existing IID ranking methods.
Second, we show that our aggregation method achieves significant improvements and SOTA results on several OoD datasets.
Finally, we demonstrate that ZooD requires significantly less computation, and, therefore, is practically scalable compared with naive fine-tuning.
Setup Details.
We use 35 PTMs with diverse architectures, pre-training methods, and pre-training datasets.
We divide the PTMs into three groups.
Group 1 consists of models with different architectures,
Group 2 consists of models pre-trained with different training methods, and Group 3 consists of models pre-trained on large-scale datasets.
We conduct experiments on six OoD datasets: PACS [43],
VLCS [24], Office-Home [77], TerraIncognita [10], DomainNet [63], and NICO (NICO-Animals & NICO-Vehicles)
Each of the datasets has multiple domains.
The standard way to conduct the experiment is to choose one domain as test (unseen) domain and use the remaining domains as training domains, which is named leave-one-domain-out protocol.
The top linear classifier is trained on the training domains only and tested on the test domain.
Each domain rotates as the test domain and the average accuracy is reported for each dataset.
To get ground-truth performance, we follow
DomainBed [29] to fine-tune top linear classifiers for the PTMs on these OoD datasets.
We adopt the leave-one-domain-out cross-validation setup in DomainBed with 10 experiments for hyper-parameter selection and run 3 trials.
We triple the number of iterations for DomainNet (5000 to 15000) as it is a large-scale dataset requiring more iterations [17] and decrease the number of experiments for hyper-parameter selection from 10 to 5.
More details on the experimental setup are in Appendix A.1.
4.1 Comparison with IID Ranking Metrics
IID ranking methods.
We divide existing ranking methods into two groups.
One group consists of methods that employ PTM’s classification layer for ranking.
These methods include NCE [76] and
LEEP [58].
The other group consists of approaches that only use PTM’s extracted features.
These methods include H-Score [8] and LogME [91].
Additionally, we also use kNN with k=200 [81] as a baseline.
Evaluation metrics.
To evaluate PTMs on OoD datasets with ranking methods, we follow leave-onedomain-out validation protocol [43].
For ZooD and kNN, we further adopt leave-one-domain-out validation for training domains and take average results as the performance prediction for the held-out test domain.
To compute the correlation between ranking scores and ground-truth performance, we use two metrics.
First, to compare the ranking of a transferability metric with accuracy, we employ
Kendall’s coefficient τ [38].
Unlike Pearson’s correlation, τ measures correlation based on the order of two measures.
Consequently, it is a better criterion for ranking.
Second, to measure the performance of transferability metric for top-model selection, we utilize weighted Kendall’s coefficient τ w [78].
The τ w gives more weight to the ranking of top-performing models compared with the rest of the models.
Therefore, it is a better comparative criterion for top model selection.
Results.
First, we compare our method with feature-based scoring methods: kNN, H-Score, and
LogME.
These methods, similar to our method, rank models based on the penultimate layer.
We compare ZooD with these methods for the full set of 35 PTMs.
We plot ranking scores and ground-truth accuracies in Figure 3.
For quantitative comparison, we also provide τ values.
It can be seen that ZooD is better correlated with fine-tuning accuracy than other ranking methods on most of the datasets.
For example, our method has a τ of 0.85 compared with LogME’s τ of 0.77 on Office-Home and a τ of 0.40 compared with LogME’s τ of 0.04 on TerraIncognita. [Table] [Table]
Furthermore, our metric is more stable and consistent.
Precisely, τ of ZooD varies between 0.40 ∼ 0.85 compared with 0.04 ∼ 0.80 for LogME, -0.08 ∼ 0.67 for H-Score, and 0.16 ∼ 0.86 for kNN.
The consistency of transferability metric across different datasets is critical since the purpose of a transferability metric is to estimate performance on a new dataset without having access to ground-truth accuracy.
Whenever an estimation metric is inherently unstable, it is hard to determine its reliability for a new dataset.
Note that our method uses a linear model with Gaussian error to approximate the top classifier.
This helps us achieve efficient model assessment, especially on small and medium-sized datasets in which the bias caused by model approximation is negligible compared with the estimation error due to insufficient data.
However, on DomainNet, things may be different.
The bias caused by model approximation dominants the evaluation performance on large datasets.
Therefore, our method does not outperform kNN on DomainNet.
Second, we compare our method with classification-layer based methods: NCE and LEEP.
For this comparison, we select a subset of our PTMs that have classification layers.
The results are illustrated in
Figure 4.
It can be seen that ZooD is also more stable and consistent than NCE and LEEP.
Moreover,
Our method achieves superior performance on the difficult real-world TerraIncognita dataset.
This dataset consists of obscure and blurry images captured by WildCams installed in different territories.
NCE has a negative correlation for this dataset.
On the other hand, our method, although not perfect, captures the relation in a better way.
For this challenging dataset, our method has a τ of 0.45 compared with 0.12 and -0.32 for LEEP and NCE, respectively.
Third, we compare the weighted Kendall’s coefficient of our method with other ranking methods.
The weighted Kendall’s coefficient is a better metric to gauge the performance of a metric for top model selection.
We also divide these results into two groups: comparison with feature-based scoring methods in Table 1a and comparison with classification-based scoring methods in Table 1b.
Our method outperforms feature-based scoring methods on 6 out of 7 datasets.
Similarly, it also outperforms both
[Table]
LEEP and NCE on 5 out of 7 datasets.
Moreover, our ranking method is more stable as it performs better on challenging datasets.
For example, it has τ w of 0.46 ∼ 0.92 compared with LogME’s τ w of 0.02 ∼ 0.90 and H-Score’s τ w of -0.20 ∼ 0.75.
In summary, transferability estimation of ZooD correlates better with ground-truth accuracy on most of the OoD datasets compared with previous ranking methods.
It also outperforms most feature-based metrics for model selection in terms of τ w .
Additionally, it is more stable and consistent across datasets, making it a better choice for pre-trained model selection.
4.2 SOTA Results with Our Selection Method
We also compare ZooD (model ranking and feature selection) with several recent SOTA OoD methods and demonstrate that it achieves substantial performance improvements.
We compare previous OoD methods with three versions of our method: 1) Single : fine-tune the top-1 model by transferability metric; 2) Ensemble : fine-tune an ensemble of the top-K models; 3) F. Selection : fine-tune an ensemble of the top-K models with feature selection, which is the expected result using ZooD. By fine-tuning, we mean using ERM with DomainBed settings to fine-tune a top linear classifier for the PTMs.
Their predictive performance and F. Ratio (the percentage of features used in F. Selection ) are listed in the last four lines of Table 2.
In all experiment results, except TerraIncognita (discussed in the next paragraph), our method achieves remarkable improvement against ERM and recent SOTA.
For Single , we list the improvements over the previous SOTA as follows:
+14% on Office-Home, +7.9% on PACS, +1.7% on DomainNet, and +0.4% on VLCS.
This result also shows that even without aggregation, using proper pre-trained model can improve OoD generalization by a large margin.
Notice that our method can be regarded as a complement to other OoD algorithms.
After selecting the top-ranked models, different OoD algorithms can be adapted to fine-tune the models.
The performance of Single does not outperform the previous SOTA on TerraIncognita.
This is because previous methods fine-tune the whole network.
In contrast, we only train a classifier on top of a fixed feature extractor.
TerraIncognita is a much more challenging dataset compared with other OoD datasets, as the majority of its images are obscured by the background.
Therefore it requires fully fine-tuning.
To show the effectiveness of ZooD with fully fine-tuning, we select the top-1 ranked model and fine-tune the whole model.
Our resulted model achieves a +2.6% improvement compared with the previous
SOTA.
One limitation of ZooD when aggregating multiple models is that fine-tuning the whole models is difficult due to the limitation of GPU memory.
However, for OoD tasks, fine-tuning the whole model may not perform better than fine-tuning the top classifier.
For example, the results of fine-tuning the full top-ranked models on PACS, VLCS and Office-Home are 90.6, 79.1 and 83.4, respectively.
Empirically, we find if a PTM is suitable for a given OoD task, fine-tuning the top classifier has better
OoD generalization than fine-tuning the full model.
As shown in Table 2, a simple model ensemble ( Ensemble ) provides fairly minimal improvement because it may introduce invariant but noisy features.
To efficiently utilize multiple models, we propose to select informative features in Section 3.2.
Here, we compare the performance improvement by F.
Selection with Single and Ensemble .
ZooD significantly outperforms both candidates while only using a small portion of aggregated features from top-K models.
Even on the most sophisticated DomainNet,
ZooD can improve predictive performance by +2.4% compared with Single and +0.1% compared with
Ensemble .
To find the appropriate number K for the model ensemble, we performed an ablation study.
We varied the number of K, e.g. K ∈ { 3 , 5 , 7 } .
The performance changes are plotted in Figure 5.
We found the performance by aggregating top-3 models strikes the right balance between performance and computational complexity.
Hence, K =3 is set to the default value.
In summary, our ranking metric in ZooD is good enough to select a model that can outperform the previous SOTA methods without adding any bells and whistles.
Furthermore, feature selection in
ZooD can efficiently utilize informative features from top-K models to further improve OoD generalization.
In this work, we do not control for the impact of better PTMs.
Given a zoo of PTMs, our method aims to exploit the power of the zoo for OoD generalization.
We can further increase the power of the model zoo by adding more PTMs.
Based on extensive experimental results on various OoD datasets, we conclude ZooD makes it easy and efficient to exploit a large set of PTMs for OoD generalization.
4.3 Computational Efficiency of ZooD
We illustrate the precision and computational efficiency of ZooD by comparing it with brute-force fine-tuning in terms of GPU hours.
The results are shown in Table 1c.
ZooD provides a minimum of 1516 × speed-up for DomainNet and a maximum of 9859 × speed-up for PACS.
Cumulatively, our method took a total of 13 GPU hours to evaluate all the PTMs on all the datasets compared with
GPU hours (equivalent to 4 GPU years) for brute-force fine-tuning.
Therefore, ZooD is a scalable and practical method for OoD generalization.
Conclusion
Machine learning models rely on IID assumption, which is often violated due to constant distribution shifts in real-world applications.
In this work, we argue for leveraging a large set of PTMs to improve
OoD generalization and propose ZooD, a paradigm for efficient PTMs ranking and aggregation.
Our paradigm avoids the computationally-prohibitive fine-tuning by ranking PTMs based on quantifying their inter-class discriminability and inter-domain stability, and selecting the most informative features from top-ranked PTMs ensemble.
Extensive experiments show ZooD is superior in ranking correlation with the ground-truth performance and achieves SOTA results on various OoD benchmarks.
Acknowledgments
Awais’s work, in part, was also supported by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant (No. 2021-0-02068, Artificial Intelligence Innovation
Hub and No.RS-2022-00155911, Artificial Intelligence Convergence Innovation Human Resources
Development (Kyung Hee University)).
We gratefully acknowledge the support of MindSpore for this research.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Section 4.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] Mainly in the Appendix. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] Will be released upon publication.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Setup Details in Section 4 and Appendix A.1.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
No]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Section 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Appendix A Experiments A.1 Complete Details of Experiment Setup
In this section, we provide a detailed experiment setup we have used.
For completeness purposes, this section also includes details already mentioned in the main paper.
Pre-trained models.
We use 35 PTMs having diverse architectures, pre-training methods and pretraining datasets.
Group 1 consists of models with different architectures.
This group consists ofdifferent architectures (CNNs and ViTs) trained on ImageNet-1k.
The architectures are as follows:
ResNet-50, ResNet-152 [30], ResNeXt-50 [82], DenseNet-169, DenseNet-201 [36], Inception v1 [72],
Inception v3 [73], MobileNet v2 [69], EfficientNet-B2, EfficientNet-B4 [74], Swin-T, Swin-B [50].
Group 2 consists of models pre-trained with different training methods.
We use 10 ResNet-50s trained via following pre-training methods: Adversarial Training [53], BYOL [28], MoCo-v2 [18], InsDis [81],
PIRL [54], DeepCluster-v2 [14], PCL-v2 [46], SeLa-v2 [4, 15], SwAV [15].
Group 3 consists of models pre-trained on large-scale datasets.
We used 13 different models trained on ImageNet-22k [67],
YFCC-100M [75], IG-1B-Targeted [84], WebImageText [64].
A summary of the PTMs can be found in Table 3.
Datasets.
We use six OoD datasets for our experiments.
The details of these datasets are listed here.
PACS [43] consists of 9,991 images from four domains (art, cartoons, photos, sketches) and seven classes.
VLCS [24] consists of 10,729 images from four domains (Caltech101, LabelMe,
SUN09, VOC2007) and five classes.
Office-Home [77] has four domains (art, clipart, product, real) of common objects in office and home settings.
The dataset has a total of 15,588 images belonging to 65 classes.
TerraIncognita [10] contains photos of wild animals taken by camera traps installed at four different locations.
It has a total of 24,788 images from 10 classes.
DomainNet [63] is one of the most challenging OoD datasets.
It has 586,575 images from six diverse domains (clipart, infographics, painting, quickdraw, real, sketch) belonging to 345 classes.
NICO [31] consists of nearly 25, images from two superclasses: NICO-Animals (10 classes) and NICO-Vehicles (9 classes).
We split the images of NICO-Animals and NICO-Vehicles into multiple domains according to [5] and combine validation and test sets as one domain to form four domains, separately.
Ground-truth performance.
To get ground-truth performance, we train linear classifiers on top of
PTMs following DomainBed [29].
The authors of DomainBed [29] argue for the hyper-parameter selection to be a part of the method selection criteria.
Based on this argument, they propose a rigorous test bench.
We follow their training and evaluation protocol, including dataset splits, hyper-parameter settings, optimizer, etc.
We adopt the leave-one-domain-out cross-validation setup in DomainBed with 10 experiments for hyper-parameter selection and run 3 trials.
We triple the number of iterations for
DomainNet (5000 to 15000) as it is a larger dataset and requires more training [17] and decrease the number of experiments for hyper-parameter selection from 10 to 5.
IID ranking methods.
We divide existing ranking methods into two groups.
The first group consists of methods that employ PTM’s classification layer for ranking.
These methods include NCE [76] and
LEEP [58].
The second group consists of approaches that only use PTM’s extracted features.
These methods include H-Score [8] and LogME [91].
Additionally, we also use kNN with k=200 [81] as a baseline.
Evaluation metrics.
To evaluate PTMs on OoD datasets with ranking methods, we follow leave-onedomain-out validation protocol [43].
For ZooD and kNN, we further adopt leave-one-domain-out validation for training domains and take average results as the performance prediction for the held-out test domain.
To compute the correlation between ranking scores and ground-truth performance, we use two metrics.
First, to compare the ranking of a transferability metric with accuracy, we employ
Kendall’s coefficient τ [38].
Unlike Pearson’s correlation, τ measures correlation based on the order of two measures.
Consequently, it is a better criterion for ranking.
Second, to measure the performance of transferability metric for top-model selection, we utilize weighted Kendall’s coefficient τ w [78].
The τ w gives more weight to the ranking of top-performing models compared with the rest of the models.
Therefore, it is a better comparative criterion for top model selection.
A.2 Extended Ranking Results
In this section
[Table]
we provide detailed and raw results for all 35 models on all six OoD datasets.
Specifically, we provide raw scores assigned by all the ranking methods to all PTMs.
We also provide accuracy of each model after fine-tuning.
A more interpretable and visual analysis of these scores are provided in section 4.1 of the main paper.
We provide these raw scores here to help aid reproducability and to help other researchers for easier benchmarking.
It is important to note that getting these results, especially accuracy results, is computationally expensive, which may hinder future progress.
For instance, on large DomainNet dataset, it takes 711 GPU days of training to get all ground-truth performance.
Therefore, providing these raw scores can significantly help future researchers.
The results are provided in the following tables.
Table 4 shows results on PACS and VLCS, Tableshows results on Office-Home and TerraIncognita, Table 6 contains results on NICO-Animals and
NICO-Vehicles
[Table]
and Table 7 contains results on DomainNet.
[Table] [Table]
B Model Ranking in ZooD
In this section, we present more details about the proposed ranking metric and algorithm.
B.1 Preliminaries: setup, problem and strategy
Suppose that:
• Model zoo.
We have a collection of PTMs as learned feature extractors:
[Equation]
, where ϕ k ( x ) is a d -dimensional feature extractor that maps X to R d . • Dataset.
A multi-domain dateset is collected for solving a domain generalization problem: D = {D 1 , D 2 ,..., D m } , with D i = (cid:8) ( x ij ,y ij ) , 1 ≤ j ≤ n i (cid:9) , where m is the number of observed domains and D i is the set of data points under the i -th domain.
The total sample size is n = (cid:80) i n i .
• Problem.
The objective is to select a PTM ϕ from M such that the optimal top classifier f based on the selected feature extractor ϕ , i.e. the whole predictor is f ◦ ϕ
( x ) , has good prediction performance on the domain generalization task.
To proceed further, we need more notations as folllows:
• For any domain i , we rewrite D i
= { y i , x i } where
[Equation]
[Table]
• Given a feature extractor ϕ , the learned feature matrix is denoted by
[Equation]
• For any i ∈ [ m ] , we denote Φ − i and y − i as
[Equation]
We can break the model selection problem down into two questions. 1).
When generalizing to unknown domains, are the learned features stable enough to avoid extrapolating predictions? 2).
Are the learned features informative enough to ensure that the correlation between features and labels is stable across different domains?
To answer these two questions
[Equation]
we compute the following two quantities:
• p (Φ i | Φ − i ) , which measures covariate shift between Φ i and Φ − i , indicating whether the validation input is a rare sample compared with the training input; • p ( y i | Φ i , y − i , Φ − i ) , which measures the discriminability and correlation shift between Φ i and y i given the training data Φ − i
and y − i .
We thus propose a metric by assembling the above quantities for PTMs ranking:
[Equation]
where λ is a tuning parameter that unifies the scale of the correlation shift and the covariate shift.
In our implementation, the tuning parameter is taken to be the ratio of two standard deviations:
[Equation]
which is also used in Ye et al.
[Table]
B.2 Model Assumption
Since the correlation between ϕ ( x ) and response variables y may be non-linear, we need to make further assumptions and approximations.
Let each y be independently generated from a unknown distribution: p ( y | Φ ,f ) .
Assume this distribution is unimodal and the mode is denoted by µ , we can take Taylor expansion of log-likelihood at the mode
[Equation]
where Λ= −∇ y ∇ y log p ( y | ϕ ( x ) ,f ) (cid:12)(cid:12) y
= µ .
The above transformation is the Laplace approximation [51] and the quadratic term implies the rationality of the Gaussian approximation.
Similar to You et al.
[91], the top model over a learned feature extractor ϕ is approximated with a linear model:
[Equation]
where ϵ is Gaussian noise with variance β − 1 .
We assume the prior distribution of the weights w is a zero-mean isotropic Gaussian distribution governed by a hyperparameter α :
[Equation]
and the conditional distribution of the target variable y given ϕ ( x ) is a Gaussian distribution:
y
[Equation]
) , w
[Equation]
x
[Equation] [Equation]
or
[Equation]
Recall the notations y i , Φ i , y − i and Φ − i in Appendix B.1.
Then we have
[Equation]
I
[Equation]
and y
[Equation]
i |
[Equation]
i
[Equation]
w ∼N (Φ
[Equation]
i w ,β − 1 I n − n i ) .
In the next section, we present the details of estimating the two hyperparameters α and β.
Appendix B. shows how to compute the conditional density p ( y i | Φ i , y − i , Φ − i ) and p (Φ i | Φ − i ) in the proposed metric (6).
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
B.3 Parameter Estimation
If we introduce a uniform prior distribution over α and β , the posterior distribution for α and β is
[Equation]
where the prior distribution p ( α,β ) is assumed to be a uniform distribution over α and β .
Then the values of ˆ α and ˆ β are obtained by maximizing the density function p ( y − i , Φ − i | α,β ) , which is also the model evidence over { y − i , Φ − i } .
The density function p ( y − i , Φ − i | α,β ) is obtained by integrating over w :
[Equation]
According to the model assumptions in Appendix B.2:
[Equation]
then the likelihood function of α and β is
[Equation]
where E ( w ) is the energy function of w ,
i
[Equation]
e.
[Equation]
Given y − i and Φ − i
[Equation]
then the posterior distribution of w is
[Equation]
where
[Equation]
Notice that
E
[Equation]
Then we have E ( m − i )
=
− 12 m ⊤− i A − i m − i + β 2 y ⊤− i y − i .
We rewrite w = w − m − i + m − i and obtain that
[Equation]
Therefore,
[Equation]
Then we have log
[Equation]
and obtain ˆ α and ˆ β by maximizing log L ( α,β ) , i.e.,
[Equation]
We can find that the objective function here is the same as Eq.(2) in You et al.
Then we use the fix-point iteration algorithm [91, 90].
The detailed inference procedure is presented as follows.
Let λ i and v i be the i -th eigenvalue and eigenvector of the matrix β Φ ⊤− i Φ − i .
That is ( β Φ ⊤− i Φ − i ) v i
=
λ
i v i .
Then we have
[Equation]
The stationary points of log L ( α,β ) with respect to α satisfy
[Equation]
with γ = d (cid:88) i
= λ
i α
+ λ i .
Notice that the eigenvalues
λ
i are proportional to β .
Hence d λ i / d β = λ i /β .
Then the stationary points of log L ( α,β ) with respect to β satisfy
[Equation]
B.4 Computing Metric
In this section, we present the details of computing the covariate shift p (Φ i | Φ − i ) and the correlation shift p ( y i | Φ i , y − i , Φ − i ) .
Then we can plug these two quantities into (6) to compute the proposed metric.
Covariate shift.
Leaving the i -th domain out, we compute the density p (cid:0) Φ i (cid:12)(cid:12)
Φ − i (cid:1) to check whether the learned feature ϕ ( x ) is stable such that the distribution shift between Φ i and Φ − i is not significant.
We approximate the distribution of ϕ ( x ) with a Gaussian distribution N ( µ ϕ , Σ ϕ ) and empirically estimate the parameters µ ϕ and Σ ϕ from the training inputs Φ − i ∈ R ( n − n i ) ×
d .
That is,
[Equation]
where 1 n − n i is a ( n − n i ) -length one vector.
Then we compute the density of Φ i according to
N (ˆ µ ϕ , ˆΣ ϕ ) :
[Equation]
Correlation shift.
Given ˆ α and ˆ β , we have
[Equation]
We write ˆ m − i
= ˆ β ˆ A − 1 − i Φ ⊤− i y − i and ˆ A − i
=
ˆ α I d
+ ˆ β Φ ⊤− i Φ − i .
According to (7),
[Equation]
To proceed further
[Equation]
we denote
[Equation]
Similar to (7)
[Equation]
we have
[Equation]
Plugging (9) and (10) into (6), we obtain the value of the proposed metric.
Remark.
Given y − i , Φ − i , ˆ α and ˆ β , the posterior distribution of w is
[Equation]
Further,
[Equation]
By calculating the integral, we can deduce
[Equation]
Therefore we can also use this distribution to calculate p ( y i | Φ i , y − i , Φ − i ) directly.
Throughout this paper, we use the formula (8) to calculate the correlation shift.
B.5 Cross-Domain Validation Selects Invariant Features
To justify our proposed selection method, and provide more intuition, we conduct explicit analysis in a linear regression setting.
Despite the over-simplification, it does reflect the essence of our approach.
From this base case, adaptions to more complicated and realistic assumptions can be made.
Data Assumption Suppose we have data in different domains with domain invariant and domainspecific features, with respect to the response variable y .
Denote the set of invariant features to be iv , which are assumed to be unit-norm and orthogonal to each other.
Without loss of generality, let data in domain D be x =( x iv , x D ) where x iv ∈ R d ∗ denotes the domain invariant features and x D ∈ R d − d ∗ denotes domain specific ones.
Let x iv be fixed.
The domain-specific features can have non-zero correlation with x iv such that
[Equation]
where A D ∈ R d ∗ × ( d − d ∗ ) , and e D ∼ N (0 ,s 2 I d − d ∗ ) .
For different domains, assume the correlation to be independently random, i.e., A D ’s are i.i.d.
matrices with independent entries with mean 0 and variance 1.
Given the features x , assume the response y only depends on x iv such that
[Equation]
where β =( β iv ,β D ) with β D =
and ϵ follows N (0 ,σ 2 ) .
Model Assumption Let the model candidates be linear models fitted to different subsets of the features and there are in total 2 d different combinations.
Denote the fitted parameters to be ˆ β ∈ R d with only the selected dimensions being non-zero.
Let the selection be ϕ , which is a subset of { 1 ,...,d } .
We want to show that our proposed statistics, in the cross-validated fashion, will prefer the optimal one with ϕ = iv .
The optimality is in the sense that it achieves the best goodness-of-fit, measured by the square loss.
More Notations Let ( X , y ) , ( ˜ X , ˜ y ) be independent datasets in two domains to be cross validated.
For any vector (matrix), we use subscript to denote part of it with selected rows (columns).
For instance, a model candidates with feature dimensions ϕ will only fit y ∼ X ϕ and the resulting ˆ β will only be nonzero on ˆ β ϕ .
For a set ϕ , denote | ϕ | be to its cardinality and ¯ ϕ to be its complement.
In our proposed test statistics, there are two terms to be assessed.
The first term is essentially the goodness-of-fit of ˜ y and ˜ X ϕ · ˆ β ϕ , which is of critical importance for selecting the invariance and consistent features across different domains.
The second term can be seen as some regularization.
In this section, we will focus on the first term, and to make things really simple, we consider expected lloss as the measure for goodness-of-fit.
The estimated ˆ β can be explicitly written as
[Equation]
Given y =
X β
+ ϵ
[Equation]
we can write
[Equation]
Thus
[Equation]
The expected l 2 loss can be expressed as
[Equation]
is wrongly selected features.
Easy to verify that E ( I 1 ) =
E ( I 3 ) = 0 and E ( I 1 I 3 ) = 0 , since ˆ β is independent with ˜ A, ˜ e , which are both mean zero.
[Equation]
For I 3 , we can further write
[Equation]
Therefore,
[Equation]
If ϕ = iv , the above quantity is minimized with E ϵ, ˜ ϵ,e,A, ˜ e, ˜
A
[Equation]
∥ ˜
[Equation]
−
[Equation]
β
[Equation]
=
[Equation]
C Feature Selection in ZooD
In this section, we present more details about the PTMs ensemble and feature selection in Section 3.2.
The top-ranked PTMs in Section 3.1 are preferred for solving the OoD generalization task.
To further aggregate different PTMs, we consider assembling the features by using PTMs as feature extractors
[Equation]
where Φ ( i ) is the i -th ranked feature extractor and [ · ] denotes the row concatenation operation.
As we show in experiments, in most cases, using aggregated models can significantly outperform any single model.
However, the rough ensemble will inevitably introduce more noise.
According to the definition of OoD learnability proposed by Ye et al.
[87], non-informative but invariant features from training domains may only bring some noise, and the accumulation of noise hurts learnability of the
OoD generalization task.
Therefore, we propose a Bayesian feature selection method based on the
Gaussian linear framework in Section 3.1.
C.1 Bayesian Variable Selection
In the Bayesian literature, the variable selection problem can be efficiently solved by introducing, for each variable w i , a binary mask z i ∈{ 0 , 1 } [48, 16, 83, 86], which are given by Bernoulli distributions governed by probability coefficient π .
Let z
=
{ z i } di =1 and p ( z ; π )=
d (cid:89)
[Equation]
From a generative perspective, these masks determine whether the weight w i is generated from a slab or a spike prior [37].
If z i = 1 , then w i will follow a slab prior with diffusing probability density; if z i = 0 , w i will have a spike prior with probability mass concentrated around 0, and thus should be discarded.
Specifically, we assume
[Equation]
Denote w = ( w 1 ,..., w d ) ⊤ and α i, 1 and α i, 2 control the shape of the w i distribution and should be reasonably large for α i, 2 .
Conditioned on w i , each data point y n is assumed to be independently drawn from a linear model with mean w ⊤ ϕ ( x ) and additional Gaussian noise with inverse variance β :
[Equation]
The model specification is completed by introducing conjugate Gamma priors over the inverse variance β
and { α i, 1 ,α i, 2 } di =1 :
[Equation]
Denote the set of Gamma prior parameters as ν =
{ ν i,j } and all latent variables as
[Equation]
Then the variable selection problem can be solved by estimating π = { π 1 ,π 2 ,...,π d } with π i = p ( z i =1) .
We can find the maximum likelihood estimator of the probability coefficient π of Bernoulli masks and then screen the variables if π i is smaller than the pre-defined threshold τ .
C.2 Variational EM Algorithm
Given the dataset { y , Φ } , the maximum marginal likelihood estimator of ( π , ν ) is given by
[Equation]
However, direct maximization of (12) is intractable due to the integration over ξ .
EM algorithm [66] might be a solution here.
In the E-step, we compute the conditional expectation
[Equation]
which involves inferring posterior p ( ξ | y , Φ; π , ν )
However, this is not straightforward to obtain due to the complexity of our model setup.
MCMC [57] is a common tool for this problem, but suffers from intensive computation, thus hard to extend to large-scale data.
We instead use approximate Bayesian inference in Section C.3.
In the M-step, we update π and ν by maximizing the expectation
[Equation]
By repeating the E and M steps, the estimator ( π new , ν new ) converges to an optimal solution.
We show this method has satisfying performance for the underlying variable selection problems in synthetic data and the prevailing OoD dataset.
C.3 Variational Inference
In the E-Step, computation of E ξ (cid:2) log p ( y , ξ | Φ; π , ν ) (cid:12)(cid:12) y , Φ; π old , ν old (cid:3) involves inferring posterior p ( ξ | y , Φ; π , ν ) .
However, due to the complexity of our model setup, no analytical form of the posterior distribution can be found.
We instead approximate true posterior distribution by variational inference [12].
The main idea involves the introduction of a set of distributions Q , which should ideally be easy to compute and provide a good approximation to the true posterior distribution.
We consider the following transformation of the marginal likelihood
[Equation]
where L ( Q ) denotes the variational lower bound.
The key point is that, through proper choice of Q distribution, L ( Q ) can be readily evaluated, and thus by maximizing the lower bound, we generally find the Q distribution, which is the best approximation within the considered family.
Here we factorize
Q over each latent variable, such that
[Equation]
which holds for classic mean-field family [11]
By denoting { m , λ , ˜ π } = { m i ,λ i ,π i } di =1 and ˜ ν
= { ˜ ν i,j } , an optimization-free form over all possible Q has been established, which can lead to minimization of
KL divergence between variational distribution Q ( ξ ) and true posterior p ( ξ | y , Φ; π , ν )
[Equation]
where denote ξ k as the k -th variable in the set ξ and ξ − k is the subset of all other variables except ξ k .
For models in conjugate families, the optimal Q ∗ ( ξ k ) has the same form as its prior distribution.
We then establish the optimization step for arbitrary variational parameters set { m , λ , ˜ ν , ˜ π } to approach the true posterior:
[Equation] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where the variational expectations are given by
[Equation]
Since the optimization steps for each variational parameter are mutually dependent, we can use coordinate gradient descent [12] starting by current Q ( ξ ) t − 1 from the last iteration.
After one-step optimization, variational parameters of Q ( ξ ) t are used in computation of E ξ ∼ Q ( ξ ; π old ,
[Equation]
old ) t (cid:2) log p ( y
[Equation]
π
[Equation]
) (cid:3) , thus finishing E-step.
During this procedure, the lower bound L ( Q ) will continuously increase until reaching its maximum value.
Therefore, the value of L ( Q ) can be used as a useful indicator for convergence of algorithm [19].
C.4 Algorithm Details
The proposed model contains a set of prior hyper-parameters π , ν , which is exactly what we want to estimate for feature screening.
In Bayesian literature, hyper-parameter selection can be automated from data through a procedure named “ARD” [52].
The original “ARD” procedure proposes a selection based on the value of model evidence.
However, in many cases including ours, this evidence is intractable.
Fortunately, it’s also feasible to use variational lower bound L ( Q ) as a substitute.
Learning prior hyper-parameters π , ν leads to the minimization of KL divergence.
This can be rationalized by the decomposition of L ( Q ) :
[Equation]
Thus by setting derivatives of each hyper-parameters with respect to L ( Q ) to 0, it’s easy to see L ( Q ) is maximized when all hyper-parameters are set to posterior parameters:
[Equation]
However, the proposed algorithm still suffers from heavy computational cost: Each iteration costs
O ( nd 2 ) .
Thus to relieve computation burden and memory usage, we leverage our method with stochastic approximation leading to the EM algorithm with stochastic variational inference [35].
In each iteration, we sample a random subset of entire data with size n s .
Fitting our algorithm over this subset for the current iteration, we obtain a local optimal estimator denoted by Q s ( ξ ) .
In M-step these intermediate variational distributions by factorizing Q s ( ξ ) will be used to learn hyper-parameters π and ν and simultaneously as the starting point for subsequent estimator in the next iteration.
In the end, we successfully reduce the computation cost to O ( n s d 2 ) with n s ≪ n , while maintaining the guarantee of convergence to the global optimum [65].
In our experiments, we collect variational probabilities of { ˜ π i } di =1 from the last three runs and early-stop the algorithm if its difference with the current probability is smaller than the pre-defined threshold ϵ or reaches the maximum iteration times.
Variational EM algorithm for Bayesian feature selection is summarized in Algorithm 2.
Note that we initialize m by linear regression and the initialization of ˜ ν is set to ν .
In our experiments, we often deal with the multivariate case.
If the underlying task involves multivariate regression or classification, i.e., Y ∈ R n × K , we can run the proposed EM algorithm on each dimension and take the union of all selected features.
Therefore, our feature selection procedure can be used in almost all prevailing models and tasks.
t ; 11: if t ≥ 3 then 12:
[Equation]
Early Stop if | π t − π mean | <
ϵ ; 14: end if 15: end for
C.5 Theoretical Result
It has been shown that our method, as well as others in Bayesian variable selection, has potentially strong selection consistency [48, 16, 83, 86].
Consider the following model with inverse Gamma prior:
[Equation]
where i runs from 1 to d , q N ,τ 0 ,N ,τ 1 ,N are constants that depend on sample size N , and IG ( α 1 ,α 2 ) is the Inverse Gamma distribution with shape parameter α 1 and scale parameter α 2 .
Under regular conditions (See conditions 4.1–4.5 in [56]), selection consistency is established:
Theorem 1.
Assume regular conditions hold, under the model with inverse Gamma prior, we have p (cid:0)
z
= t | Y ,σ 2 (cid:1) p −→ 1 as n → ∞ , that is, the posterior probability of the true model goes to 1 as the sample size increases to ∞ .
More related works on Bayesian feature selection can be found in [26, 55].
C.6 Simulation Study
In this section, we will conduct a series of simulations to verify selection performance on an i.i.d.
dataset with varying sizes and dimensions.
Here, we consider cases in the standard multivariate regression.
We first generate each input predictor from a standard normal distribution: x ni ∼ N (0 , 1) for i =1 ,...,d , and thus we generate response variables by subsequently sampling β j ∼ Uniform(1 , 3) for j =1 ,...,k<d and y n ∼ N ( (cid:80) ki
=1 β i x ni , 1) .
We then vary the values of d and k to find the potential influence in terms of True Positive Rate (TPR) and False Positive Rate (FPR).
The results are shown in
Table 9.
We repeat each case 50 times and present the mean and variance of TPR and FPR.
The hyper-parameter setting is listed in Table 8.
We vary n s to study the influence of batch size.
Overall, our method illustrates the experimental selection consistency.
When n>d , our method almost always selects the correct k variables with TPR close to 100% and successfully screens all unnecessary variables with
FPR equal to 0% .
Even under the less informative circumstance when n has an equal or less amount than d , our method can still achieve great selection results with TPR above 90% .
As n goes up, there is a uniform improvement in all cases in terms of TPR and FPR
[Table]
[Table] [Table]