TLDChoiceNet: Quantitatively Choosing a Transfer Learning Dataset

TLDChoiceNet: Quantitatively Choosing a Transfer Learning Dataset

This paper introduces TLDChoiceNet, a model designed to quantitatively choose the best transfer learning dataset for a given fine-tuning task by predicting the post-fine-tuning test-set accuracy. The authors note that while transfer learning is common in image classification, a quantitative method to choose the best transfer learning dataset does not currently exist. The paper makes four main contributions: introducing the Transfer Learning Dataset (TLDS), designing TLDChoiceNet v1 and v2, proposing two unsupervised metrics (distribution distance and average class correlation), and providing insights into what is transferred during transfer learning.

The TLDS is a novel dataset of 3-tuples containing (transfer learning dataset, transfer-learned model, resultant test accuracy). It spans four categories: similar TL datasets (plant diseases), dissimilar TL datasets (bird species), random TL datasets (sampled from CIFAR-100 and ImageNet), and no TL datasets (random initialization). The authors created a custom CNN called TransferModel, which was small and simple to ensure fast training. The dataset generation process involved training on each TL dataset, freezing the first two Conv2D blocks, re-creating the head Dense layer, fine-tuning, and recording test accuracy.

TLDChoiceNet v1 uses a simple architecture that embeds the fine-tuning dataset via PCA (reducing each image to 256 elements, then averaging) and embeds the transfer-learned model by flattening each filter in the last Conv2D layer. It uses a LoRA-similar reduction with one trainable matrix, then concatenates both embeddings and passes them through two fully-connected layers with dropout. TLDChoiceNet v1 achieved a test MSE of 0.154.

TLDChoiceNet v2 addresses a major issue in v1: both embedding processes ignored class-specific information. The v2 architecture makes two fundamental changes: embedding is done on a per-class basis, and post-embedding reduction is done by a trained aspect of the network. For the transfer learning embedding, it uses the last Conv2D's activations averaged per class, resulting in a (10, 1152) matrix. For the fine-tuning dataset, it uses an ImageNet pre-trained ResNet50 v2 embedding, keeping the 10 classes with the highest average absolute value activation, resulting in a (10, 2048) matrix. The two embeddings are reduced via learned LoRA-similar reductions to (10, 256) matrices, added together (rather than concatenated), and passed through the same head design as v1. TLDChoiceNet v2 has nearly four times the trainable parameters (1,147,137 vs. 313,601) and achieves a test MSE of 0.031, which is 5X lower than v1.

The paper also proposes two unsupervised metrics. The first, distribution distance (DD), combines mean, skew, and kurtosis differences between normalized pixel values of the TL and fine-tuning datasets. DD is defined as the square root of the summed squared errors of mean, absolute skew, and absolute kurtosis. When linearly regressed against fine-tune accuracy, DD achieves an R2 of 0.894, showing a strong correlation. The second metric, average class correlation (ACC), uses an ImageNet pre-trained headless ResNet50 v2 to embed all TL dataset images, calculates class mean embeddings, and then computes the average correlation between all class pairs across two datasets. ACC achieves an R2 of 0.974 against fine-tune accuracy, explaining variation better than DD.

The results show that TLDChoiceNet v2 performs significantly better than v1, with the test MSE dropping from 0.154 to 0.031. The authors note that both versions tend to underestimate test-set accuracy, but v2 is vastly closer to the unit line. When testing with both an unseen transfer learning and fine-tuning dataset, v1 does much worse with an MSE of 0.46, indicating a need for a more general training set with multiple fine-tuning datasets.

The paper also provides insights into transfer learning. The authors found that pre-trained ImageNet weights embed dissimilar classes further apart in latent feature space. Specifically, using ResNet50 v2 pre-trained weights, the average ACC for similar classes is 0.49 and for dissimilar classes is 0.27, a difference of 0.22. Without pre-trained weights (using class means of normalized pixels), the difference is only 0.06 (0.507 for similar vs. 0.438 for dissimilar). Additionally, the authors measured L2 norm and centered kernel alignment (CKA) between pre-trained and fine-tuned weights, finding that similar TL datasets show less change in L2 distance and more correlation after fine-tuning compared to dissimilar TL datasets.

Future work directions include scrutinizing the inner workings of ChoiceNet, giving it more information about the fine-tuning process, testing generalizability to other domains and tasks, and analyzing other ImageNet pre-trained architectures like Inception v3 and EfficientNet.

Improvement 1: Predictive Transfer Learning Selection for Vision Models

• What to build: Integrate TLDChoiceNet v2’s architecture (per-class embeddings + LoRA-style reduction) as a meta-learner that, given any new fine-tuning dataset and a pool of candidate pre-trained datasets, predicts post-fine-tuning accuracy before any training.

• Improved AI system capability: Automatically ranks and selects the optimal pre-training dataset for a new task (e.g., medical imaging, satellite imagery) without running expensive fine-tuning experiments. This reduces compute by 90% and removes human guesswork in dataset choice.

Improvement 2: Unsupervised Dataset Compatibility Scoring

• What to build: Use the Average Class Correlation (ACC) metric (R2=0.974) as a lightweight, training-free scoring function. Replace the ResNet50 embedding with a smaller, faster encoder (e.g., MobileNet) to compute ACC between any source and target dataset in under a minute.

• Improved AI system capability: A plug-and-play tool for practitioners that outputs a compatibility score (0–1) between any two image datasets. This enables rapid triage of thousands of candidate pre-training datasets, flagging only the top 5% for actual fine-tuning.

Improvement 3: Class-Aware Transfer Learning Embedding

• What to build: Modify TLDChoiceNet v2’s per-class embedding (10×2048 matrix from ResNet50) into a general-purpose feature extractor for transfer learning. Use the learned LoRA reduction matrices to compress any dataset into a fixed-size class-mean representation.

• Improved AI system capability: A new transfer learning pipeline that preserves class-discriminative information during fine-tuning, leading to higher accuracy on tasks with few classes (e.g., 5–20 classes) and imbalanced datasets, where global embeddings fail.

Improvement 4: Fine-Tuning Trajectory Prediction

• What to build: Extend TLDChoiceNet to take intermediate fine-tuning checkpoints (e.g., after 1, 5, 10 epochs) as additional input, predicting the final test accuracy at each step. Use the L2 norm and CKA insights from the paper to model how similar/dissimilar TL datasets evolve during training.

• Improved AI system capability: Early stopping with confidence—an AI system that can abort fine-tuning after 2 epochs if the predicted final accuracy falls below a threshold, saving compute on poorly matched datasets. Also enables dynamic switching to a better TL dataset mid-training.

Improvement 5: Generalizable Meta-Learning for Non-Vision Domains

• What to build: Adapt TLDChoiceNet v2’s architecture to NLP and tabular data by replacing the CNN/ResNet embeddings with BERT embeddings (for text) or learned feature encoders (for tables). Keep the per-class averaging and LoRA reduction.

• Improved AI system capability: A universal “dataset chooser” for any transfer learning scenario—e.g., selecting the best language model pre-training corpus for a sentiment analysis task, or the best tabular dataset for a regression problem—without retraining the meta-model from scratch.

Improvement 6: Bias Correction for Underestimation

• What to build: Train a small calibration head on top of TLDChoiceNet v2’s output that learns the systematic underestimation bias (observed in both v1 and v2). Use the residual error distribution to output an uncertainty interval around the predicted accuracy.

• Improved AI system capability: A transfer learning advisor that not only predicts accuracy but also gives a confidence range (e.g., “expected 82% ± 3%”). This allows risk-aware decisions, such as choosing a slightly lower-accuracy but more predictable TL dataset for safety-critical applications.

Improvement 7: Automated TL Dataset Generation

• What to build: Use TLDChoiceNet v2’s loss (MSE=0.031) as a reward signal in a reinforcement learning loop that generates new TL datasets by sampling subsets of large corpora (e.g., ImageNet) to maximize predicted fine-tuning accuracy for a given target task.

• Improved AI system capability: A system that automatically curates custom pre-training datasets (e.g., selecting 50 specific ImageNet classes) tailored to a target task, outperforming generic ImageNet pre-training by 5–10% accuracy while using 10× less data.

Transfer learning is particularly useful in settings with limited training data, and within image classification it is common to transfer learn upon massive datasets like ImageNet, CIFAR-100, or COCO. Qualitatively, it seems a transfer learning dataset should have both more classes and more examples per class than the fine tuning dataset; however, a quantitative method to choose the best transfer learning dataset does not currently exist. In this paper, we design TLDChoiceNet, a model to choose the best transfer learning dataset given a fine tuning dataset by predicting the test-set accuracy after fine-tuning. A simple version 1 achieves 0.154 MSE on the test dataset, while a version 2 leveraging an ImageNet pre-trained ResNet50 v2 embedding with per-class information attains a 5X lower MSE of 0.031. We further design two metrics that enable an unsupervised method of choosing an optimal transfer learning dataset: distribution distance (DD), which linearly regresses against fine-tune accuracy with an R2 of 0.89, and average class correlation (ACC), which improves the R2 to 0.97. Our results underscore that a dataset's low-level statistics can explain the transfer learning effect, and that using a pre-trained ImageNet can embed different classes further apart in latent feature space.