Model Selection with Model Zoo via Graph Learning

TL;DR

The paper addresses the challenge of selecting suitable pre-trained models from large model zoos for fine-tuning. It introduces TransferGraph, a graph-learning framework that frames model selection as a link prediction problem over a graph built from models, datasets, and rich metadata. By leveraging graph learners (eg, Node2Vec, GraphSAGE, GAT) and supervised predictors, it predicts fine-tuning performance and ranks candidates with higher accuracy than baselines. Across 16 real datasets spanning image and text domains, TransferGraph achieves up to a 32% improvement in correlation between predicted and actual fine-tuning results, demonstrating the practical value of modeling intrinsic model-dataset relationships for efficient model selection.

Abstract

Pre-trained deep learning (DL) models are increasingly accessible in public repositories, i.e., model zoos. Given a new prediction task, finding the best model to fine-tune can be computationally intensive and costly, especially when the number of pre-trained models is large. Selecting the right pre-trained models is crucial, yet complicated by the diversity of models from various model families (like ResNet, Vit, Swin) and the hidden relationships between models and datasets. Existing methods, which utilize basic information from models and datasets to compute scores indicating model performance on target datasets, overlook the intrinsic relationships, limiting their effectiveness in model selection. In this study, we introduce TransferGraph, a novel framework that reformulates model selection as a graph learning problem. TransferGraph constructs a graph using extensive metadata extracted from models and datasets, while capturing their inherent relationships. Through comprehensive experiments across 16 real datasets, both images and texts, we demonstrate TransferGraph's effectiveness in capturing essential model-dataset relationships, yielding up to a 32% improvement in correlation between predicted performance and the actual fine-tuning results compared to the state-of-the-art methods.

Figures (17)

• Figure 1: Illustration of the model selection problem setting.
• Figure 2: Average fine-tuned accuracy of the top 5 selected models compared between random selection strategy and our proposed solution learning from a graph along with metadata (example dataset: stanfordcars krause_3d_2013).
• Figure 3: Link prediction in the context of model selection
• Figure 4: Graph properties
• Figure 5: An overview of TransferGraph on model selection for fine-tuning, including model zoo construction (stage 1), training (stage 2-3) and model selection (stage 4).

Theorems & Definitions (3)

• Example 3.1
• Definition 3.1: Graph
• Example 3.2