Learning the Regularization Strength for Deep Fine-Tuning via a Data-Emphasized Variational Objective

TL;DR

This work targets the inefficiency of grid-search hyperparameter tuning for regularization in deep transfer learning. It introduces a data-emphasized ELBo (DE ELBo) that scales the likelihood with a factor $\kappa$ (recommended as $\kappa = D/N$) to balance data fit and model complexity, enabling simultaneous learning of backbone and head parameters and hyperparameters $\lambda,\tau$ in one training run. The method leverages a variational approximation $q(w)V$ and closed-form updates for the hyperparameters, yielding substantial reductions in compute time while maintaining competitive held-out accuracy on CIFAR-10 and Oxford-IIIT Pet. The proposed approach reduces the need for exhaustive grid searches, supports effective hyperparameter learning in high-dimensional settings, and holds promise for wider applicability in related learning paradigms.

Abstract

A number of popular transfer learning methods rely on grid search to select regularization hyperparameters that control over-fitting. This grid search requirement has several key disadvantages: the search is computationally expensive, requires carving out a validation set that reduces the size of available data for model training, and requires practitioners to specify candidate values. In this paper, we propose an alternative to grid search: directly learning regularization hyperparameters on the full training set via model selection techniques based on the evidence lower bound ("ELBo") objective from variational methods. For deep neural networks with millions of parameters, we specifically recommend a modified ELBo that upweights the influence of the data likelihood relative to the prior while remaining a valid bound on the evidence for Bayesian model selection. Our proposed technique overcomes all three disadvantages of grid search. We demonstrate effectiveness on image classification tasks on several datasets, yielding heldout accuracy comparable to existing approaches with far less compute time.

Figures (2)

• Figure 1: Model selection comparison between the ELBo (left) and our data-emphasized ELBo (DE ELBo) (right) for two ResNet-50s trained on CIFAR-10 $N=1000$. For both models, we fix the estimated posterior $q$ and vary $\lambda, \tau$. Takeaway: Without enough training data or with too many model parameters, the ELBo has a preference for simpler models.
• Figure 2: Test-set accuracy on CIFAR-10 (top row) and Oxford-IIIT Pet (bottom row) over training time for L2-SP with MAP + grid search (GS) and our data-emphasized ELBo (DE ELBo). We run each method on 3 separate training sets of size $N$ (3 different marker styles). Takeaway: Our DE ELBo achieves as good or better performance at small dataset sizes and similar performance at large dataset sizes with far less compute time. To make the blue curves, we did the full grid search once (markers). Then, at each given shorter compute time, we subsampled a fraction of all hyperparameter configurations with that runtime and chose the best via validation NLL. Averaging this over 500 subsamples at each runtime created each blue line.