Optimal split of compute between pretraining and fine-tuning in ridge regression
Compute-Optimal Pretrain--Fine-tune in Ridge Gradient Descent
Machine Learning
Summary
When training a machine learning model in two steps—first pretraining it on general data, then fine-tuning it for a specific task—there is a fixed amount of computing power to use. Spending more compute on pretraining means less is left for fine-tuning, and vice versa. The authors analyze this trade-off in a simple, understandable mathematical setting to find the best way to divide computing resources. They show that the best split depends on how the initial training directions relate to the later task and the structure of the data used for fine-tuning.
What this means in practice
- •For machine learning engineers: Allocate training compute efficiently for better model adaptation in two-stage learning with limited resources.
- •For data scientists: Design experiments that balance pretraining and fine-tuning efforts by understanding data geometry and spectral properties.
A theory result. No direct application yet.
Authors
Alex Buna, Fanghui Liu, Patrick Rebeschini
Abstract
Pretraining followed by fine-tuning introduces a compute-allocation problem: under a fixed training budget, compute spent improving the upstream objective reduces the compute available for downstream adaptation. Despite its practical importance, this trade-off is not yet well understood theoretically, even in simple models. In this paper, we cast this allocation as a compute-split problem under a two-stage pretrain--fine-tune procedure with fixed total optimisation budget, using regularised least squares trained by gradient descent as a tractable setting. We characterise the optimal split under data-dependent evaluation geometries induced by the fine-tuning problem. Our results show that the allocation depends on how pretraining directions affect fine-tuning predictions and how fine-tuning shifts are seen through downstream data geometry. In particular, the relevant quantities are determined by prediction-relevant spectral components of the pretraining and fine-tuning empirical covariances. Technically, the analysis relies on a basis-invariant, eigenspace-level spectral decomposition, together with perturbative control of the non-commuting pretraining and fine-tuning dynamics.