Model Ratatouille: Recycling Diverse Models for Out-of-Distribution Generalization

TL;DR

This work tackles out-of-distribution generalization by exploiting the diverse features learned by multiple fine-tunings of the same foundation model. It introduces Model Ratatouille, a recycling approach that performs parallel target fine-tunings from diverse auxiliary-task initializations and only at the end averages their weights, leveraging linear mode connectivity to improve generalization. Empirically, Ratatouille sets SoTA on DomainBed, demonstrates that diversity across auxiliary tasks drives gains, and provides insights into when LMC-based averaging holds, as well as the trade-offs for in-domain tasks. The paper also discusses embracing an updatable, open-community paradigm for sharing initializations and collaboratively updating models. Overall, Ratatouille offers a robust, scalable path to harness diverse auxiliary knowledge for better OOD robustness and a blueprint for collaborative model improvement.

Abstract

Foundation models are redefining how AI systems are built. Practitioners now follow a standard procedure to build their machine learning solutions: from a pre-trained foundation model, they fine-tune the weights on the target task of interest. So, the Internet is swarmed by a handful of foundation models fine-tuned on many diverse tasks: these individual fine-tunings exist in isolation without benefiting from each other. In our opinion, this is a missed opportunity, as these specialized models contain rich and diverse features. In this paper, we thus propose model ratatouille, a new strategy to recycle the multiple fine-tunings of the same foundation model on diverse auxiliary tasks. Specifically, we repurpose these auxiliary weights as initializations for multiple parallel fine-tunings on the target task; then, we average all fine-tuned weights to obtain the final model. This recycling strategy aims at maximizing the diversity in weights by leveraging the diversity in auxiliary tasks. Empirically, it improves the state of the art on the reference DomainBed benchmark for out-of-distribution generalization. Looking forward, this work contributes to the emerging paradigm of updatable machine learning where, akin to open-source software development, the community collaborates to reliably update machine learning models. Our code is released: https://github.com/facebookresearch/ModelRatatouille.

Figures (20)

• Figure 1: The different fine-tuning strategies discussed in this paper: vanilla fine-tuning oquab2014learning, moving average izmailov2018 and variants Wortsman2022robust, model soups Wortsman2022ModelSA and DiWA rame2022diwa, inter-training phang2018sentence, fusing choshen2022fusing and our proposed model ratatouille. They all start with a pre-trained foundation model. Some strategies fine-tune the pre-trained model on auxiliary tasks (thin solid arrows ): these auxiliary fine-tunings can be performed by different contributors of the community on their own data. Then, all strategies perform fine-tuning on the target task of interest (thick solid arrows ). Finally, the weights fine-tuned on the target task are used as is, or are averaged (dashed arrows ) into a final model. Ratatouille (i) enables compute parallelism throughout training, (ii) maximizes the amount of diversity in models' predictions, (iii) achieves state-of-the-art performance in DomainBed gulrajani2021in, the standard computer vision benchmark for OOD generalization and (iv) does not incur any inference or training overhead compared to a traditional hyperparameter search.
• Figure 2: Illustrations of (a) different linear mode connectivity (LMC) conditions, and (b) model ratatouille. In subplot (a), we illustrate \ref{['obs:1']}, about LMC between two checkpoints along the same target fine-tuning; \ref{['obs:2']}, about LMC between two target fine-tunings; \ref{['hyp:1']}, about LMC between two auxiliary fine-tunings; and \ref{['hyp:2']}, about LMC between two target fine-tunings initialized from auxiliary weights satisfying \ref{['hyp:1']}. In subplot (b), we offer a diagram of our proposed recycling strategy, where we (i) fine-tune a pre-trained model on auxiliary tasks, (ii) plug a linear probe on the pre-trained model and the auxiliary fine-tunings, (iii) fine-tune on the target task from each auxiliary weights, and (iv) return their weight average as the final model.
• Figure 3: Explorations on q-diversity kuncheva2003measures and its positive impact on accuracy for the OOD test domain "Art" from OfficeHome. In (a), we compute the diversity between pairs of models either directly fine-tuned from ImageNet, either inter-trained on DomainNet: having one model from each initialization increases diversity. In (b), we plot this diversity along the 5k training steps. In (c), we observe that the more diverse the models, the higher the accuracy gain of their weight average compared to the average of their individual accuracies. In (d), we average $M$ models: a proportion $(1 - \mu)$ start directly from ImageNet, the others $\mu$ are inter-trained on DomainNet. The accuracy of the weight average is maximized when $\mu\approx 0.5$.
• Figure 4: \ref{['fig:pacs0_lmc_hyp1_ood', 'fig:vlcs0_lmc_hyp1_ood', 'fig:home0_lmc_hyp1_ood', 'fig:terra0_lmc_hyp1_ood', 'fig:came0_lmc_hyp1_ood']} validate \ref{['hyp:1']} by plotting $\lambda \to \mathrm{acc}_\mathrm{te}\mathopen{}\mathclose{\left(\mathopen{}\mathclose{\left(w^\mathrm{lp}, (1 - \lambda) \cdot \phi_a^{\mathrm{aux}} + \lambda \cdot \phi_b^{\mathrm{aux}}\right)\right)$, where $w^\mathrm{lp}}}$ is the linear probe of $\phi_{\mathrm{IM}}^{\mathrm{pt}}$, and $\phi_a^{\mathrm{aux}}$ and $\phi_b^{\mathrm{aux}}$ are fine-tuned on the two auxiliary datasets in the legend "Dataset$_a$ to Dataset$_b$". \ref{['fig:pacs0_lmc_hyp2_ood', 'fig:vlcs0_lmc_hyp2_ood', 'fig:home0_lmc_hyp2_ood', 'fig:terra0_lmc_hyp2_ood', 'fig:came0_lmc_hyp2_ood']} support \ref{['hyp:2']} by plotting $\lambda \to \mathrm{acc}_\mathrm{te}\mathopen{}\mathclose{\left((1 - \lambda) \cdot \theta_a + \lambda \cdot \theta_b\right)$ where $\theta_a$ and $\theta_b$ are fine-tuned on the target task starting respectively from $(w^\mathrm{lp}}, \phi_a^{\mathrm{aux}})$ and $(w^\mathrm{lp}, \phi_b^{\mathrm{aux}})$. We encounter two exceptions to \ref{['hyp:2']} (\ref{['fig:terra0_lmc_hyp2_ood', 'fig:came0_lmc_hyp2_ood']}), due to the fact that neither the auxiliary nor the target task bear enough similarity with the pre-training task.
• Figure 6: OOD accuracy ($\uparrow$) for model ratatouille when increasing the number of auxiliary tasks and uniformly averaging all fine-tuned weights. For each target task, we consider the first domain as the test OOD; the other domains are used for training.

Theorems & Definitions (1)

• Remark 1