Pre-train, fine-tune, interpolate: a three-stage strategy for domain generalization
Abstract
The goal of domain generalization is to train models that generalize well to unseen domains.
To this end, the typical strategy is two-stage: first pre-training the network on a large corpus, then fine-tuning on the task’s training domains.
If the pre-training dataset is large enough, this pre-training is efficient because it will contain samples related to the unseen domains.
Yet, large pre-training is costly and possible only for a few large companies.
Rather than trying to cover all kinds of test distributions during pre-training, we propose to add a third stage: editing the featurizer after fine-tuning.
To this end, we interpolate the featurizer with auxiliary featurizers trained on auxiliary datasets.
This merging via weight averaging edits the main featurizer by including the features mechanisms learned on the auxiliary datasets.
Empirically, we show that this editing strategy improves the performance of existing state-of-the-art models on the DomainBed benchmark by adapting the featurizer to the test domain.
We hope to encourage updatable approaches beyond the direct transfer learning strategy.
Introduction: the art of fine-tuning
Distribution shifts between train and test data can reduce the performance of deep neural networks [1, 2].
Thus, learning models that generalize well to unseen domains is critical as those models are now vastly deployed for real-world applications in many different contexts [3, 4].
The current paradigm for domain generalization is transfer learning [5, 6]: we pre-train a model on a large dataset and then fine-tune it on the available domains for the task.
The goal of the pre-training is to obtain the best possible featurizer.
Though different learning objectives generalize differently [7, 8, 9, 10, 11], recent works suggest that the main differences lie in the pre-training data distribution [12, 13, 14, 15] and the learnt features diversity [16, 17, 18].
To mitigate distribution shifts, this dataset should ideally contain diverse images with diverse concepts; it would transfer diverse features that can match domains [19], facilitate the optimization [20, 21], and overall generalize well on all kinds of distributions.
Yet, under a limited pre-training budget, the ideal pre-training is task-dependent [22, 13, 23].
Amortizing the computational cost is at the cornerstone of the “foundation models” paradigm [24].
Yet, only a few companies are capable of collecting the necessary huge-and-diverse data and deploying the hardwares.
This is a major limitation to the community with risks of centralization, monetization, privacy issues and lack of transparency [25].
To reduce the reliance on proprietary models, there is a need of alternative strategies involving individual stakeholders leveraging only open-source datasets such as ImageNet [26].
A promising lead is inter-training : recent works [27, 28, 29, 30, 31, 32] leveraged auxiliary datasets to enrich the featurizer.
They fine-tune the model on a task, “intermediate” between the pre-training and the target tasks.
We show the benefits but also the main limitation of inter-training: the auxiliary task should be adapted to the test domain, yet unavailable at initialization in domain generalization.
[Figure]
Weights averaging
In this paper, we propose a novel editing strategy to enrich the features; we edit the featurizer after fine-tuning, as illustrated in Figure 1.
Our featurizer is interpolated towards featurizers trained on auxiliary datasets.
Operating after training enables us to choose the adequate interpolating coefficients at test time, and thus to cheaply adapt the featurizer in function of the test domain.
This is a radical departure from existing transfer learning strategies, that tried to anticipate all kind of distributions at initialization.
Applied on the reference DomainBed benchmark [33] in Section 4, editing after training improves the performance of state-of-the-art domain generalization strategies.
Our main contributions are the following:
We highlight a key limitation of the inter-training strategy in Section 2. 2.
We show that we can enrich the featurizer after fine-tuning and propose in Section 3 a threestage strategy to improve domain generalization; pre-training on a large corpus, fine-tuning on the available training domains and then editing with auxiliary datasets.
Related work.
Our strategy follows the “updatable machine learning” [34] literature, further discussed in Section 5.
Notably, we are inspired by the recent success of weight averaging (WA) to merge models trained independently [35, 36], even when trained on different tasks [37, 38, 39].
Our work is also related to the editing literature [40, 41, 42] which tries to correct mistakes as they appear: none tackled domain generalization nor leveraged weight averaging.
Inter-training before fine-tuning
θ ImageNet are ResNet50 ImageNet weights those specialized on DomainNet.
From each initialization, we train 20 networks on the “Clipart”, “Product” and “Photo” domains, leaving apart the
OOD domain “Art”.
Finally, we either average all weights (DiWA [36]) or select the best one based on the train-domain validation accuracy (ERM).
A promising strategy to enrich features.
As illustrated in Figure 1, inter-training adds an intermediate learning task on an auxiliary dataset between pre-training and fine-tuning on the target task.
The idea is to leverage similar datasets to enrich and specialize the features.
Indeed, in the vicinity of the target task there are similar auxiliary tasks that would facilitate the learning of the adequate features.
With ImageNet [26] as the pre-training dataset, we experiment inter-training on DomainNet [28] for domain generalization on OfficeHome [43] in Figure 2.
The initialization is thus (1 − λ )
θ ImageNet + λθ DN where θ ImageNet are the ImageNet weights and θ DN are those weights specialized on DomainNet.
We observe that full inter-training ( λ = 1 ) improves the accuracy on the out-of-distribution (OOD) domain “Art” compared to the direct transfer learning from ImageNet ( λ = 0 ).
Inter-training requires a priori knowledge from the test domain.
Yet, inter-training suffers from catastrophic forgetting [44] of the features previously learnt on ImageNet: learning new features does not fully preserve those from initialization.
This explains why the inter-training on DomainNet is detrimental on the domain “Product”: λ ≈ 0 performs best for DiWA.
Thus, a more “robust initialization” can usually be found for 0 < λ < 1 in a similar fashion to “robust fine-tuning” [45]; by interpolating between the two extremes, we can find a better trade-off between pre-trained features and auxiliary features.
Critically, the best value for λ is domain-dependant, and cannot be detected using the other domains; this finding contrasts with recent works [46] observing “a clear linear relationship between IID accuracy and OOD accuracy”.
In conclusion, inter-training helps only when the auxiliary task is selected with some prior knowledge of the future test domains; DomainNet helps for OfficeHome and PACS [47] but will be detrimental for TerraIncognita [48] in Section 4.
Editing after fine-tuning
We then train 20 networks on the “Clipart”, “Product” and “Photo” domains, leaving apart the OOD domain “Art”.
To obtain θ HOME , we either average all weights (DiWA [36]) or select the best one based on the train-domain validation accuracy (ERM).
At test time, we use (1 − λ ) θ HOME + λθ DN where θ DN are specialized on DomainNet.
We argued in the previous section that we need some knowledge of the future test domains to best enrich the features with auxiliary tasks.
This motivates a more flexible strategy that would be applied after fine-tuning, when some test samples could be available.
To this end, we propose our editing strategy, illustrated in Figure 1; we leverage features from the auxiliary tasks not at initialization, but at inference time, without requiring a retraining step.
In brief, we merge the featurizers fine-tuned on the target task with those fine-tuned on auxiliary tasks.
To do so, we leverage weights averaging [35, 49], which enriches features (see Appendix B) and empirically succeeds to merge networks even when trained on different tasks [37, 38, 39] despite the non-linearities [50, 51].
Thus we predict using (1 − λ ) θ target + λθ aux as the featurizer; the classifier (on top of the featurizer) remains the one fine-tuned on the target task as the auxiliary tasks may not have the same classes.
We validate the benefits of editing for OfficeHome in Figure 3 and Appendix D, and on PACS in
Appendix E. We note that editing consistently helps.
We also remark similarities with the curve obtained via inter-training in Figure 2; the main advantage is that the editing operation is done after fine-tuning, and thus can be adapted cheaply to the test domain.
Experimental results on the DomainBed benchmark
Setup.
We evaluate our strategy on the DomainBed [33] benchmark, which imposes the code, the training procedures and the ResNet50 [52] architecture.
Each domain is successively considered as the test while others are for training: domains are split into 80% (used as training and evaluation) and 20% (used as validation).
The featurizer is pre-trained on ImageNet, and the classifier is initialized with linear probing [53] to prevent features distortion: from this shared initialization, we fine-tune the whole network with 20 hyperparameter configurations sampled from Table 5.
For ERM, θ target is the one with highest validation accuracy.
DiWA [36] uniformly averages the weights from all training runs, removing the need for model selection.
We apply DiWA on DomainNet [28] to obtain θ aux by averaging the 20 networks trained on DomainNet.
At inference time, we use (1 − λ ) θ target + λθ aux . λ is either arbitrarily set to 0 .
5 , or optimized on the validation dataset of the test domain: in the latter case, we select the λ with the highest validation accuracy among 20 sampled from the uniform distribution U (0 , 1) .
Thus, for fair comparison, scores are reported with DomainBed’s test-domain model selection.
This experimental setup is further described in Appendix C.
Results.
In OfficeHome [43] in Table 1 and PACS [47] in Table 2, editing and inter-training consistently improve ERM and DiWA.
The gains are larger for domains where DomainNet brings useful information, e.g., on the “Clipart” domain from OfficeHome, and smaller on others, e.g., on the “Product” domain.
λ
= 0 .
5 works well because DomainNet was selected with some human’s prior as the auxiliary task for OfficeHome and PACS; yet, λ = 0 .
5 fails for TerraIncognita [48] in Tablebecause this dataset is not related to DomainNet’s domains.
In contrast, editing on TerraIncognita with oracle-domain validation for λ ∼ U (0 , 1) does not harm accuracies as λ ≈ 0 are selected (see
Appendix F).
This test-time adaptation after fine-tuning is not possible with inter-training, which therefore performs badly.
This highlights the main advantage of editing w.r.t.
inter-training; when the auxiliary dataset is not relevant, the editing strategy can ignore the auxiliary weights at no cost.
[Table]
[Table] [Table]
Limitations and future work.
We will investigate the impact of other auxiliary tasks, e
[Table]
g., iNaturalist [54] to improve performance on TerraIncognita.
We will also report scores with standard deviation for different pre-training strategies [55] and architectures [56]
[Table]
and open-source our code.
Finally [Table] we will analyze the arithmetic properties of fine-tuning and weights averaging [57].
Updatable machine learning for domain generalization
The DomainBed benchmark [33] showed that no invariance-based approaches [58, 59, 60] outperformed ERM for domain generalization .
Recently, weight averaging [49, 61] (WA) strategies did by averaging weights obtained along a training run [62, 63] or from multiple independent runs [35, 36] to increase diversity in predictions: similarly, editing with auxiliary tasks increases diversity in features.
WA relies on the linear mode connectivity across networks [64, 65] [Table] which notably holds when networks are trained from a shared pre-trained initialization on a shared task [66].
Yet, recent findings suggest that WA may succeed even with different initializations [67, 68] — modulo a simple weights permutation — and on different tasks [37, 38, 39].
The latter works are the most similar to us.
[37] motivates a Fisher-averaging in NLP; in contrast, we show that basic interpolation is efficient and discuss the selection of the interpolating coefficient.
[38] proposed an “embarrassingly simple parralelization”.
[39] is specialized for zero-shot CLIP [55] models.
Interestingly, our approach is also similar to “Robust fine tuning” [45]; while they rely on the initialization to generalize well, we learn appropriate auxiliary weights to edit the fine-tuned weights.
These works complement each other and all highlight the potential of the updatable machine learning [34] paradigm.
Conclusion.
We proposed a strategy for domain generalization: we edit the featurizer after finetuning by interpolating towards weights trained on auxiliary tasks.
We validate this approach on the DomainBed benchmark.
In contrast with centralized foundation models, we could leverage a repository of averageable specialized networks trained in a decentralized way, akin to a github repository for softwares [34].
We hope to pave the way towards deep models training in collaboration.
A Learning strategies as pseudo-equations
In Figure 1, we illustrated the different learning strategies.
In Equation (1), we now provide an analytical formulation of these strategies with pseudo-equations, where θ represents the featurizer’s weights and D the datasets.
[Equation]
B Similarity between weight averaging and representation averaging
Our editing strategy relies on the intuition that averaging weights enriches features: in this section, we present a preliminary experiment strengthening this intuition.
We consider two weights θ 1 and θ 2 obtained from two independent fine-tunings on OfficeHome from
ImageNet weights.
We are interested in their feature representation in the penultimate layer, just before the classifier, noted Φ θ 1 and Φ θ 2 when extracted over the validation dataset.
The representation (Φ θ 1 + Φ θ 2 ) / 2 averages the features extracted from both featurizers.
We also consider the features extracted by their WA, i.e., Φ ( θ 1 + θ 2 ) / 2 .
Table 4 shows the L 2 norm of the representation differences.
The two closest representations are the averaged representation (Φ θ 1 + Φ θ 2 ) / 2 ad the one from the weight averaged network Φ ( θ 1 + θ 2 ) / 2 .
It suggests that WA successfully encodes the two representations to construct a richer one.
In conclusion, while [35, 36, 49] argue that WA approximates predictions ensembling [69], this experiment suggests that WA also approximates features averaging.
[Table]
C DomainBed
We further detail our experiments on the DomainBed benchmark [33].
Training protocol.
OfficeHome [43], PACS [47], and TerraIncognita [48] are divided into four domains.
Each domain is successively considered as the test domain while other domains are used in training.
We follow the training protocol from https://github.com/alexrame/diwa , which is itself adapted from https://github.com/facebookresearch/DomainBed .
For each dataset and domain, we perform a random search of 20 trials on the mild hyperparameter distributions described in Table 5.
We use a ResNet50 [52] pre-trained on ImageNet, with a dropout layer before the newly added dense layer and fine-tuned with frozen batch normalization layers.
The optimizer is Adam [70].
Our classifiers are initialized with linear probing [53]; we first learn only the classifier (with the featurizer frozen) with the default hyperparameters defined in Table 5; the classifier’s weights are then used to initialize all subsequent runs.
Following [36, 62, 63], for OfficeHome, PACS, and
TerraIncognita, all runs are trained for 5 k steps with validation accuracy computed every 100 steps; for DomainNet, all runs are trained for 15 k steps with validation accuracy computed every 500 steps.
[Table]
Baselines.
ERM is the standard Empirical Risk Minimization.
Mixup is the classical data augmentation [71] applied across domains (Interdomain Mixup [72]).
SWAD [62] average weights along the trajectory of an ERM run.
DiWA [36] uniformly averages the weights obtained from different ERM runs and removes the need of model selection.
Editing.
A inference time, our featurizer is (1 − λ ) θ target + λθ DN , where θ DN is a featurizer specialized on DomainNet [28] and λ the interpolating hyperparameter.
Specifically, θ DN is obtained by applying
DiWA’s protocol on DomainBed for DomainNet, but with all domains available for training: thus θ DN is the weights average of 20 networks fine-tuned on all six domains from DomainNet.
Setting arbitrarily λ
= 0 .
5 performs well on OfficeHome and PACS only because DomainNet was selected with some human prior information for these datasets; yet, it fails for TerraIncognita that contains images very different from DomainNet’s.
Test-domain model selection.
For more adaptive editing, needed notably for TerraIncognita, we can use a more principled approach based on test-domain model selection, i.e., when validation samples belong to the test domain.
To do so, we sampled 20 values for λ from the uniform distribution
U (0 , 1) and select the best one based on the test-domain validation accuracy.
That’s why, for fair comparison, the scores reported for baselines are those with test-domain model selection.
In contrast with train-domain model selection, the test-domain model selection is arguably more realistic for real applications: indeed, after deployment of their models, the practitioners would validate that the performances are correct, and thus would label some samples.
We will ablate the importance of the validation dataset’s size in future revision.
D Editing on OfficeHome D.1 Results on OfficeHome from DomainBed
[Table]
[Table]
[Figure]
E Editing on PACS E.1 Results on PACS from DomainBed
[Table]
E.2 Editing curves w.r.t. λ
F Editing on TerraIncognita
[Table]
[Equation]