A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
Authors: Michael Munn
Susan Wei
Venue: ICML 2025 poster (2025)
Keywords: transfer learning
Bayesian model selection
efficient fine-tuning
model adaptation
Primary area: probabilistic_methods->bayesian_models_and_methods
TL;DR: We propose a new Bayesian model selection criterion, the "downstream free energy", which measures a model's adaptability to new tasks. This criterion does not require downstream data and shows promising results in predicting fine-tuning performance.

Abstract: Recent advances in artificial intelligence have been fueled by the development of foundation models such as BERT, GPT, T5, and Vision
Transformers. These models are first pretrained on vast and diverse datasets and then adapted to specific downstream tasks, often with significantly less data. However, the mechanisms behind the success of this ubiquitous pretrain-then-adapt paradigm remain underexplored, particularly the characteristics of pretraining checkpoints that enhance downstream adaptation. We introduce a Bayesian model selection criterion, called the downstream free energy, which quantifies a checkpoint’s adaptability by measuring the concentration of nearby favorable parameters for a downstream task. We demonstrate that this Bayesian model selection criterion can be effectively implemented without access to the downstream data or prior knowledge of the downstream task. Furthermore, we provide empirical evidence that the criterion reliably correlates with improved fine-tuning performance, offering a principled approach to predicting model adaptability.

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
Michael Munn Susan Wei 2
Abstract
progress across natural language processing and computer vision with applications including text classification (Qiu et al., 2020), text generation (Li et al., 2024), image classifi- cation (Liu al., 2023b), object detection (Sanchez et al., 2020), medical imaging (Mormont al., 2018; Chen et al.,
Recent advances in artificial intelligence have been fueled by the development of foundation models such as BERT, GPT, T5, and Vision Transformers. These models are first pretrained on vast and diverse datasets and then adapted to specific downstream tasks, often with signifi- cantly less data. However, the mechanisms behind the success of this ubiquitous pretrain pvetrain-then-a thtntthennapppe paradigm remain underexplored, particularly the characteristics of pretraining checkpoints that en- hance downstream adaptation. We introduce a Bayesian model selection criterion, called the downstream free energy, which quantifies check- point's adaptability by measuring the concentra- tion of nearby favorable parameters for the down- stream task. We demonstrate that this Bayesian model selection criterion can be effectively imple- mented without access to the downstream data or prior knowledge of the downstream task. Further- more, we provide empirical evidence that the crite- rion reliably correlates with improved fine-tuning performance, offering a principled approach predicting model adaptability.
2019; Ke et al., 2021), autonomous driving (Kim & Park, 2017) and robotics (Jaquier et al., 2023).
As a result, there is a growing body of research that aims to better understand the theoretical reasons behind the success of this terain-ttten-daapp paradigm (Galanti al., 2022: Munn et al., 2024). One of the key open questions is to understand how to select pretraining checkpoints which are optimal for adaptation. A number of practical heuristics have emerged through experimental intuition and empirical analysis (Liu et al., 2023a), but a principled theoretical framework for effective checkpoint selection still lacking.
To address this, we repurpose well-established concepts from Bayesian statisties and propose downstream free en- ergy as pretraining model selection criterion. Downstream free energy measures the negative log of the concentration of well-performing network weights near a pretraining check- point when evaluated on downstream data. In statistical lingo, this nothing more than the (negative log) marginal likelihood where the integral is restricted to local neighbor- hood around the pretraining checkpoint. Intuitively, lower downstream free energy indicates higher concentration of parameters in parameter space for which the model is more adaptable and capable of generalizing well on downstream tasks. In short, checkpoints with lower downstream free energy are better suited for adaptation and thus should be preferred during pretraining.
1. Introduction
The advent of foundation models has significantly reshaped the landscape of modern machine learning (Bommasani et al., 2021). Trained on expansive, diverse datasets using su- pervised or self-supervised learning methods, these models learn generalized representations that can then be success- fully adapted (or finetuned) to a wide array of downstream tasks, often where there is significantly less data or limited computational resources (Bengio, 2012; Brown et al., 2020). This pretrain-then-adape paradigm has emerged as dom- inant and highly successful technique driving significant
Although the use of downstream free energy pretraining model selection criterion has strong theoretical grounding in Bayesian statistics, it comes with an unfortunate caveat: to compute it requires access to the downstream dataset which may not be available to the practitioner during pretrain- ing. However, under certain distributional shift conditions between the pretraining and downstream data, it is possi- ble to overcome this limitation. Namely, we introduce the pretraining free energy. which is computed solely on the pretraining data, and show that minimizing it serves as a reliable proxy for minimizing the downstream free energy (see Proposition 5.1). Together, these insights provide a
'qqual contribution 'Google Research, New York, USA DDept. of Econometrics and Business Statistics, Monash Univer- sity, Melbourne, Australia. Correspondence to: Michael Munn <munn @ @googlc.com>>
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
solid justification for using the pretraining free energy as model selection criterion during pretraining. This strategy is particularly advantageous when pretraining is intended to be general purpose, as is the case with most foundation models.
We experimentally confirm (Section 6), using varied datasets and architectures, that lower pretraining free energy not only enhances downstream adaptability (Figure 2 and Figure 3) but also exhibits a stronger correlation with adaptability compared to other pre- training metrics (Table I).
0.52
learning rate betch size momentum
0.50
2. Relationship to Prior Work
Frietnn Tranderr Accurcy
0.48
Implicit bias in transfer learning. The term implicit bias refers to the tendency of optimization processes, such as stochastic gradient descent (SGD), to inherently guide the model's learning dynamics towards solutions with proper- ties which are not explicitly prescribed by the loss function (Neyshabur et al., 2017; Soudry et al., 2018: Gunasekar et al., 2018). For example, the selection of training hyperpa- rameters, such the learning rate and batch size, can have significant effect on the optimization efficiency as well as on the quality of the learned model (Keskar et al., 2017; Masters & Luschi, 2018: Goyal, 2017; He et al., 2019; An- driushchenko et al., 2023). As a result, there has been con- siderable effort to understand the mechanisms which govern these implicit biases during model training. However, the effect of implicit bias in transfer it impacts successful downstream domain adaptation growing but less explored area of research (Lippl & Lindsey, 2024; Kumar et al., 2022).
0.46
0.44
0.42
0.40
0.38
10000
20000
30000
40000
50000
Pretreining Free Energy
learning rate batch size momentum
Frsetuns Aug S-chot Trarsder Arcuraayy
C.00
0.58
0.58
how
0.54
0.50
10000
20000
Pretraining free Energy
30000
40000
50000
In transfer learning, the ability to identify and leverage pretraining biases to predict and improve downstream test error is highly valuable. Recent work of (Liu et al., 2023a; Galanti et al., 2022; Munn et al., 2024) can be viewed as establishing relationships of the form
Figure 1. We plot pretraining free energy versus types of trans- fer accuracy (top and bottom) for checkpoints at the end of pretrain- ing. As expected, checkpoints with lower pretraining free energy, across various pretraining hyperparameters such as learning rate, batch size, and momentum, show higher transfer accuracy. The size of the icons represent magnitude of the hyperparameter value; e.g., a larger triangle means higher momentum. The reported val- ues averaged over five random sceds. See Section 6 for details.
downstream test error X pretraining characteristic. (I)
Ideally, these pretraining characteristics are sensitive to fac- tors which can be manipulated by practitioners, thus allow- ing for deliberate influence and intentional design during pretraining. Furthermore, any such pretraining characteris- should be accessible using only pretraining data, since knowledge the downstream task data is typically not available. It is worthwhile to note that (Liu et al., 2023a; Galanti et al., 2022; Munn et al., 2024) mainly consider the linear probe as their fine-tuning method while we consider full fine-tuning.
To justify our theoretical results, we exploit certain pretrain- ing mechanisms that are known to reduce the pretraining free energy, such as larger learning rates, smaller batch sizes and higher momentum (Lau et al., 2025). We then verify that these mechanisms, which lead to reduced pretraining free energy, in turn correlate with improved downstream adaptation performance. A preview of these results is pre- sented in Figure 1. In summary, our contributions are:
(Liu et al., 2023a) explore the role of implicit bias in lan- guage modeling and establish an empirical relationship be- tween the pretraining flatness (measured by the trace of the Hessian of the pretraining loss) and the downstream test accuracy. Their experiments verify that lower pretraining flatness, which they show is effectively regularized by SGD, strongly correlates with better downstream performance. Al- though this work does not provide a formal bound as in (1), it offers valuable empirical evidence on how the implicit
We introduce the downstream free energy as novel model selection criterion for quantifying downstream adaptability (Section 4.1).
We prove the downstream free energy can be controlled by the pretraining free energy (Proposition 5.1) and provide insight into how this free energy perspective informs practical pretraining heuristies (Section 5.1).
2

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
flatness regularization of SGD acts to benefit transfer learn- ing. This is particularly beneficial since techniques exist for explicitly minimizing loss landscape sharpness; e.g., (Foret et al., 2021; Wen et al., 2023).
collection of models, Mi,..., M,, the task of choosing optimal model for some given data is known as model selection. There are different (and sometimes irreconcilable) model selection criteria; but, in general, all model selection criteria attempt to balance fit and complexity. A particu- larly appealing Bayesian model selection criterion is the free energy criterion which is widely used and accepted in the both the statistical and machine learning literature (Hinton & van Camp, 1993; Kass & Raftery, 1995; MacKay, 2002; Robert et al., 2007). The free energy model selection criterion says we should pick the model with the lowest free energy. Since the free energy is the negative log of the marginal likelihood, also known as Bayesian model evi- dence, free energy minimization is equivalent to marginal likelihood maximization. To our knowledge, this work rep- resents the first application of the free energy criterion in the domain of transfer learning.
(Galanti et al., 2022) examine the efficacy of transfer learn- ing through the lens of neural collapse, a recently observed phenomenon which characterizes the geometry of last-layer features and weights for overparameterrzed classification networks (Papyan et al., 2020). They show through theory and experiments that the neural collapse exhibited during pretraining generalizes to new classes of the downstream task as well, thus enabling successful model adaptation. Drawing on the formalism described in (1), (Galanti et al., 2022) can be seen as deriving theoretical bounds of the form
downstream test error
downstream neural collapse neural collapse.
pretraining
However, despite supporting neural collapse as a beneficial pretraining characteristic, practical methods to explicitly regularize it are lacking.
3. Problem Setup
Here, we shall mainly treat the supervised setting though the theory developed below applies equally to the unsupervised setting. During pretraining, for input x and target y, we employ a probabilistic model p" (y/ar, w) parameterized by w E W C IP. Throughout, we assume the pretraining model p (,)l,, w) depends on through a neural network f (x) = aua(v" moalv""Slll (o())) where w = (v, Here do de- notes the feature extractor parameterized by 0 and v the weights of the linear head. The final activation is denoted Cour: e.g., softmax or sigmoid for classification tasks.
(Munn et al., 2024) make progress in this direction by means of the geometric complexity, model complexity measure introduced and analyzed in (Dherin et al., 2022). They prove that the geometric complexity of the model's ' learned feature representations upper bounds the model neural col- lapse. Furthermore, their experiments verify that techniques which implicitly reduce this geometric complexity during pretraining (such as large learning rates, small batch sizes and increased 2 regularization) in turn put regularizing pressure on the pretraining neural collapse leading to im- proved transfer test accuracy.
For fine-tuning, we attach new linear head u to the backbone de resulting in a neural network f (xx) oour (u"To (z))) where w' = (u. 0) with u potentially hav- ing different dimension to v. The fine-tuning probabilistic model is denoted p² (yla,, w²) where the dependence on x is through /ef
Our key contribution is the identification of free energy as a novel and significant pretraining characteristic which exhibits direct theoretical and empirical connections govern- ing successful downstream model adaptability. We prove in Section 5 that, similar to neural collapse, the pretraining free energy bounds from above the downstream free energy. In addition, we establish (see Appendix A) a theoretical link between downstream free energy and the downstream Bayesian prediction, providing thcoretical guarantees on the downstream Bayes test error. Together, these theoretical results, viewed in the context of (1), imply
Given a pretraining checkpoint w" = (v",0"), we initialize f at (uo, 6*) where uy is randomly initialized. All param- eters of w' are then fine-tuned via stochastic optimization. In this work, we employ limited fine-tuning where the lin- ear head undergoes standard training, while the backbone remains mostly frozen, with updates governed by a separate, smaller learning rate. This approach particularly useful in scenarios with limited downstream data, where the dif- ferential learning rates help to prevent overfitting or loss of purpose representations; cf. (Lee et al., 2022).
downstream
downstream free energy
pretraining free energy.
Bayesian test error
Furthermore, using mechanisms established in (Lau et al., 2025) which are known implicitly regularize the pretrain- ing free energy-such as large learning rates, small batch sizes, and increased momentum we experimentally ver- ify (see Section 6) that lower pretraining free energy does indeed lead to improved fine-tuning performance.
general-purpose
For theoretical convenience, we will assume that u and v share the same dimensionality'
This way, we can use p(y|æ, w) to denote both the pretrain-
'This cases exposition by avoiding having to distinguish be- tween p (v) w)) and p' (yke,, w'). Note, we do not adhere to this restriction our experiments.
Bayesian model selection criterion. The idea of using free energy has its roots in Bayesian model selection. Given a
3

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
ing and fine-tuning models. Let the true (and unknown) pretraining (i 0) and fine-tuning (i 1) joint distribu- tions be denoted
selection (Balasubramanian, 1997) using a local version of the free energy criterion to select among the local optima. This amounts to comparing the downstream free energies between elements of Uo. We now define the downstream free energy associated to an element of Up-
y) (((())))) i = 0, 1;
and define the pretraining (i = 0) and fine-tuning (i =) test loss to be
4.1. Downstream free energy
With datasets D and D as above, let = |D0| and m [D'| Informally, we might say that a pretraining checkpoint w- = (v", 0*) E Uo is good candidate for adaptation there are many weights in the vicinity of 0 with low fine-tuning test loss; i.e., low values of K (w). One way to make this mathematically precise is via the downstream free energy
K' (w) = Eri()) ExI DKr, (r'((()()))))))))))))) ( w)).
Let 2 and D be datasets drawn from the pretraining and downstream distributions (resp.) and
the corresponding pretraining (i = 0) and fine-tuning (i
1) sample losses be
R (w))== = Di
Σ (I.
(logr' - log p(yla, w))
F' (", == log Z (B,BB()))
(1)
which is the negative log of a local marginal likelihood
Note that minimization of ('w) and (www) with respect to w can recover the standard cross-entropy loss and squared loss frequently employed in deep learning. Indeed, if we drop the entropy term in K' and K', which does not depend on w, we obtain the negative log likelihoods, for i = 0,
Z' (B, = / expl-mK (p()) dw. (2) (,,()))
Here (pw)) is prior over the model parameters w, and (w := (w = (o*,0) ||0 00 S 1/73 is the 7- neighborhood around w" with v" frozen. Note that large values of y force us to stay near 0 and thus, ultimately, stay near the pretraining checkpoint w' = (v",0") as well.
L' Lww)= == -E,((,.)) (a.w) log plyl(,, w) (w))
1
logr(ylar,w). (x.w)ED)
Σ
|D'
Taken together, equations (1) and (2) imply that a large con- centration of weights @ near 6 with low downstream test loss K! (w) results in large Z (B, and, equivalently, small F (B,("")) )). Thus, we propose the following down- stream free energy strategy for improved fine-tuning:
We double load test loss to mean either K' or loss to mean either Ks or L'"
L' and train
4. Pretraining and downstream free energy
In this section, we begin by introducing the downstream free energy mcasure of how suitable checkpoint is for downstream adaptation. We then introduce the pretraining free energy as proxy that can be measured solely using the pretraining data.
Pretraining checkpoints with lower downstream free energy are more likely to adapt successfully to downstream tasks.
Formally, we seek to find parameters w² E Uo which mini- mize the downstream free energy; i.e.,
Let Ug = [w= = (2, ta denote the set of local min- ima of the pretraining test loss K (ww). In our theoretical development, we will frequently refer to the elements of Uo as pretraining checkpoints. Note that the elements of Uo. being local minima of the test loss, generally differ from the actual checkpoints obtained during pretraining, which are governed by the training loss (w) (or equivalently, (w)). To bridge this gap between theory and practice, checkpoints should correspond to local minima of the train- ing loss. This ensures that the theoretical objects we analyze minimizers of the test loss are meaningfully related to their empirical counterparts.
arg min (w),(w))) w EUo
(3)
Before addressing the implementation of this free energy strategy, let's first understand the competing forces behind this model selection criterion. Given w² E Un. following the techniques set out in (Watanabe, 2009), the asymptotic expansion of F (B,((w")) in the sample size m is
F1 (B,("1))))
= mk (w=1) + X (ww") log m + Ollg log m), (4)
Given a single model a parametric family M = {p(y|æ, w) : w E W}- - with multiple optima (as neural net- works are prone to exhibit), we can perform internal model
where -= == arg min B, (w) K (w). Further discus- sion, including the derivation of equation (4), can be found in Section and Appendix B of (Lau al., 2025).
4

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
Remark 4.1. From (4), note that that downstream free energy of checkpoint w" is weighted sum of two things: the fit, as measured by K- (w"", ). and the complexity, measured by X' (www). This complexity measure X (w") was recently introduced as the local learning coefficient; see Lau et al. (2025). Lower local learning coefficient means lower model complexity. Note that a checkpoint with higher loss under the downstream distribution may still be preferred as long as its complexity is low enough to compensate. Furthermore, note that for pretraining checkpoints that are in the same level set of K2, the checkpoint with the lowest model com- plexity, as measured by X will have the lowest downstream free energy.
and /3 >> is an inverse temperature. Unlike 7 (B( and (I (w,("))) here the quantities 7 (" (w=); ); (8) and F (B, ); 3) are stochastic. We indicate this by dropping the overhead bar.
Analogous to (4), the asymptotic expansion of F (B,(w"); in n for w E Uo
F (B,
- (w + >² log + Op log log
where wnt arg B. K" (w). Note that the asymptotic expansion of F' (', )) in (4) involves the downstream test loss K' whereas the asymptotic expansion of (B, () in (7) involves the pretraining train loss Ko, To compare the two, we take the expectation over the dataset in (7), arriving at the following expansion involving only deterministic quantities:
The free energy strategy in (3) which uses F (B, (w"))) to select among candidate checkpoints in Un is conceptually sound but presents two notable implementation challenges. First, F1 (B, (w))), besides involving some unknown terms such as K2 is the negative log of an intractable integral. many techniques such as MCMC or variational inference are available to deal with intractable integrals.
This is not insurmountable
Do F (B, (w()
(8)
= =nokk (w*)) + X (wwwwww))))) log" + O(log log n).
The second, and more significant, issue is that applying F' (B, (w")) to select among checkpoints w* E Uo requires access to downstream data. This poses a problem because, in many practical scenarios, the downstream task may not be known or fully available during pretraining. To address this limitation, we introduce the pretraining free energy, analog of the downstream free energy but which can be computed using only the pretraining data. In Section 5 we show how these two quantities are related.
In the next section, we will use these asymptotic expan- sions to bound the discrepancy between the downstream and pretraining free energy.
5. Relationship between pretraining and downstream free energy
In this section, we show there is satisfying relationship between pretraining free energy and downstream free energy, asymptotically speaking. Relying on the leading order terms of the asymptotic expansion of the downstream free energy in (4), we can express the downstream free energy strategy in (3) as
Remark 4.2. Note that the free energy as defined in equa- tions (I) and (2) is not scale invariant with respect to pa- rameters W. Thus, for certain neural network architectures exhibiting strict scale invariance, such as those composed purely of ReLU activations, it's possible for global pa- rameter rescaling to leave model outputs and downstream accuracy unaffected, while potentially altering the free en- ergy in some non-trivial way. However, our investigation here centers on commonly deployed neural networks, which typically incorporate elements like normalization layers or weight decay that break strict parameter scaling invariance.
arg min un EUD [mK'(w"') + X (w"" log m|
(9)
where w² = arg min (ww) K- (www). To avoid requir- ing the downstream test loss K' we introduce the pretrain- ing asymptotic free energy strategy which relies only on the pretraining distribution and (under mild assumptions, be- low) serves as a viable proxy for (9). Formally, this strategy seeks a solution of the following optimization
4.2. Pretraining free energy
Similar to the downstream free energy defined in (1), we define the pretraining free energy for pretraining check-
arg min [nock (w") + X' log w0 EUio
(10)
point w- = (0*+ E Up as
Fdd (B,) - (B:,("))::
(5)
n log This strategy is supported by the following result whose proof can be found in Appendix C.
where 30 M logn
where
Proposition 5.1. Let w be a local minimum of K (w);; Lee. w² E Uo and yy be such that w is a local minimum of (w)); i.e. w E Up. Further suppose X >(ww) S (w
(B,
expl-n3k (p) dw
B,(w')
(6)
5

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
((a))) Define M == max(x.))) (aa))
< co. Then,
and
F3 = n3 (wz)) X (1
log n.
(13)
K'(w" (w= + A ) log m
Suppose K (weg) < K (w))) i.e., wa and wa are in differ- ent level sets and checkpoint w has lower pretraining test loss; but >" > ((w)), implying checkpoint w3 is less complex than checkpoint w Then it is entirely possible for Fa > F3 so that checkpoint wa will be preferred by (10) despite having higher pretraining test loss. In fact, this happens precisely when
m
(11)
< MK MK"(w")+ (w") D+ + logm
m
where D == floo
r ( r (z) y) dr dy. ()())
Proposition 5.1 justifies model selection using the asymp- totic expansion of the pretraining free energy as in (10). This follows from (11) by first multiplying both sides by m and then noting that minimizing
1 (((w)) - ((w3) M (w33) ("w)))
m
log m
mMK (w") + mD (ww" log m
Recall, m represents the number of examples in the down- stream dataset. Note that, when M is large. there's a smaller range of m under which the suboptimal pretraining check- point will be preferred. In other words, if the downstream distribution is very different to the pretraining distribution, the free energy strategy will look to the lower level sets of pretraining test loss.
is equivalent, up to constants, to minimizing
log n [mMKK(w") + (w("") (lw"")ogm) logm], log m
which leads us precisely to (10). To further illustrate Propo- sition 5.1, we include explanatory examples in Appendix D which interprets this result applied to Gaussian distributions.
Observation 2: When n3n >> log n, a checkpoint with lower pretraining test loss will always be preferred by the pretraining asymptotic free energy strategy in (10). Again, suppose we have two local minima wa w3 E Un but which are in different level sets of the test loss; i.c. K (wx) * (w3). Without a handle on So. we cannot decide which checkpoint has lower free energy since, as described above in Observation 1, the complexity term A also plays a role in comparing Fa and Fg.
There are some real-world scenarios for which Proposition 5.1 would be uninformative. For example, if the pretrain- ing data includes only images of horses while the down- stream data contains only cars, their label supports would be disjoint, leading to an infinite M. To address this, our experiments in Section focus on settings where the pre- training dataset is significantly larger and more diverse than the downstream dataset. This also reflects common practice in the field and an established heuristic in transfer learning; see also (Kornblith et al., 2019). Specifically, we achieve this by using pretraining datasets with substantially larger set of image classes. If this were reversed; i.c., the pretrain- ing dataset has substantially fewer classes than the down- stream dataset, the relationship we establish in Proposition 5.1 would be uninformative.
However, when n,g is significantly larger than log n, the first term in (10) dominates the second. In this case, the pretraining asymptotic free energy strategy prioritizes check- points with lower pretraining test loss K".
Using the definition of Ag in (10), the setting described here is equivalent to Mm >> log m, where m is the size of the fine-tuning dataset and M measures distributional shift. Since m already grows faster than log m, this may offer an intriguing insight which justifies the pretraining test loss as a heuristic for checkpoint adaptability.
5.1. Observations of the pretraining asymptotic free energy strategy
In this section, we present practical observations that fol- low from selecting pretraining checkpoints according to the pretraining asymptotic free energy strategy defined by (10).
Observation 3: For checkpoints with the same pretrain- ing test loss, the one with the lowest complexity is pre- ferred by the pretraining asymptotic free energy strat- egy in (10). Suppose we have two models wx w/ Un in the same level set of K"; i.e., same pretraining test loss K (wt) (ww)). As before, our strategy (10) directs to compare F. and F3 as defined in equations (12) and (13), resp.
Observation 1: A suboptimal checkpoint in terms of pre- training test loss can still be preferred by the pretraining asymptotic free energy strategy in (10). Suppose we have two models w w3- E Un: i.e., both models are local min- ima of the pretraining test loss K", In order to determine which model is preferred for fine-tuning, our strategy (10) directs us to compare
However, since the first terms are equal, selecting the pre- ferred pretraining checkpoint depends only on the model complexity, as measured by (w(w)) and (((w)). Thus, all else being equal, the strategy in (10) naturally prefers
F. = (n8kK((W)) + (((w() log
(12)
6

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
simple pretraining checkpoints over more complex ones for improved fine-tuning.
It is not hard to see that (16) is a local- ized adaptation of Watanabe's classic Widely Applicable Bayesian Information Criterion (WBIC) (Watanabe, 2013). The classic WBIC itself was developed because the stan- dard Bayesian Information Criterion (BIC) (Schwarz, 1978) is unsuitable for singular statistical models. Recall that a model is said to be 'regular' if its parameter-to-distrbuiii mapping one-to-one and its Fisher information matrix is positive definite for all possible parameter values; other- wise, it is singular. The key distinction of the pretraining WBIC, as defined in (16), and the classic WBIC is its local- ization through a Gaussian prior centered on the pretraining checkpoint w".
where 8 log n
5.2. Estimating pretraining free energy
So far, we have established the pretraining asymptotic free energy strategy
theoretically principled approach to pretraining model selection for improved finetuning. In this section, we show how to estimate the pretraining asymp- totic free energy required in (10) using only the sample pretraining train loss L This estimation technique, which we employ in our experiments (Section 6), enables the ap- plication of our proposed strategy in (10) for real-world machine learning scenarios.
The pretraining WBIC at checkpoint w" is good estimate of the (expected) pretraining free energy around w" defined by equations (5) and (6). Furthermore, WBIC(w"; ;,3)) can be reliably computed via SGLD sampling methods; see Lau et al. (2025, Appendix G).
We begin by focusing first on model selection for pretraining checkpoints in the same level set of K In this case, we can set 3o to an arbitrary value; we set 8o-- = 1. Next, note that the optimization objective in (10) can be equivalently ex expressed in terms of I since it differs only from K by a constant with respect to w. In other words, we have
Therefore, to apply the pretraining asymptotic free energy strategy in (10) to checkpoints with the same K we simply select the one with the smallest pretraining WBIC given by WBIC(w"; 3. Next, we empirically verify this strategy using the CIFAR dataset trained on ResNet-18.
arg min [nK](w") ) + X (w(")) log
= =arg min [nL (w") (² logn] EnE
(14)
To estimate the RHS of (14), we refer to recent work of (Lau et al., 2025) which shows that the Widely Ap- plicable Bayesian Information Criterion (WBIC) around w" E Un is an asymptotically unbiased estimator of mL (w*) + (w)") log n. This localized version of the WBIC is computed from the sample pretraining train loss i measured in the neighborhood (,(w") of the checkpoint w² as described below.
6. Experiments
The goal of our experiments is to evaluate how well the pretraining WBIC, which estimates the pretraining free en- ergy as described in Section 5.2, correlates with downstream performance. In order to measure the impact of lower pre- training WBIC, we apply mechanisms during pretraining which are known to implicitly regularize this quantity, as shown in (Lau et al., 2025). These include including large learning rates, small batch sizes, and high momentum.
Consider a localizing Gaussian prior which acts as a sur- rogate for enforcing the domain of integration given by (www). Specifically, let
We use the CIFAR-FS dataset (Bertinetto et al., 2019), derived from CIFAR-100 where the 100 classes are di- vided into 64 classes for meta- veta--raining, 16 classes for meta- validation, and 20 classes for meta-testing. We pretrain on the meta-training set and then assess model adaptability on the unseen meta-test set via limited fine-tuning described in Section 3. The vela-validation classes are not used.
(p(w)) ix exp(-7"1[101)) 7
which is centered at the origin with scale vector 7- = (?.....%%)) Since we only want to measure the free en- engy with respect to parameters 0 of the model backbone (recall, the fine-tuning setup described in Section 3), we take 7j do in the coordinates of and ?/ = in the coordinates of 0, where y is the same as the radius defining the neighborhood (w"": recall, (2).
Pretraining. For pretraining, we use all 64 classes from the CIFAR-FS meta-training set train a ResNet-18 model using stochastic gradient descent (SGD). We explore ranges of hyperparameter values for the learning rate, batch size and momentum. Interaction effects between these are not considered. Full experiment details for each hyperparameter sweep are provided in Appendix B. During training we track the pretraining train loss (first column of Figure 2) and the pretraining WBIC (second column of Figure 2). The hyperparameter settings for pretraining WBIC computation are provided in Appendix B.I.
Define the pretraining posterior distribution
(w; w0,""", 10",7777 7) x expl exp(-n,80 ( (w))) - w²). (15)
Following Lau et al. (2025), we define the pretraining WBIC at w E UoD by
(BII(("::") = [ [mo](c)) (w; w", ,,,,,) dw,
(16)
7

---

A Bayesian Model Selection Criterion for Selecting Pretraining Cheekpoints
2:05
Tate
17006
2.45
2:30
2:055
300000
2..43
0.33
90000
2.33
0.90
0.00
2000
0.35
20023
18003
1U8C3
7238)
4230)
10293
ARRS:
JC003
RARA
22000
1933)
(210) Sum
2.45
0.00
2.5
300000
0.33
3.0
Seyture Saatter
3.5
512
1.0
3.2
10290 20033
40E39
SC093 ME23 /2002
2200) 2238) FRXA
S03K (O0A) 1938)
1000 AAXO AARE BAKU XANO MARK 7900)
11602 JUECE AREZ MAKE MADE MAKE WRES
390000
Peequee Tranter AcDuaNc
Arp S-bbat Theesto Aec.umer
2.000
3.5
90000
2.0
1:42
1.5
1.25
1:99
3.0
9.45
2000
335
Iaretion
Sereton
Benats
Beratoon
Figure 2. Model checkpoints with lower pretraining WBIC (second column) consistently result in better transfer accuracy. both when fine-tuning on the full downstream dataset (third column) and in the few-shot setting (fourth column). Lower pretraining WBIC correlates with better downstream performance for Top row: larger learning rates, Middle row: smaller batch sizes, and Bottom row: increased momentum. Additional experiments mini-ImageNNt and VGG model yield similar results; Figure 3 and Appendix E.
Full meta-test fine-tuning uses the full meta-test dataset, consisting of all 20 meta-test classes with 600 examples per class. We use an 80/20 split for training and testing, with stratification within each class. In this setting new (randomly initialized) linear head is attached for the 20- class classification task, and the model is fine-tuned for 100 steps using SGD. This setting corresponds to the "Fine- tune Transfer Accuracy" metric (third column) in Figure 2. Hyperparameter details for this setting are in Appendix B.2.
5.2) and better downstream performance; see Figure 2. In particular, we see that increasing learning rate, decreasing batch sizes, and increasing momentum all result in lower pretraining WBIC, which in turn leads to better downstream performance. Note the Avg 5-shot transfer accuracy (fourth column) is typically higher than the finetune transfer accu- racy (third column); this is likely because the former only needs to learn 5 classes at time while the latter needs to learn 20 classes. Interestingly, we can view pretraining train loss (the first column of Figure 2) as a baseline comparison. We see that pretraining train loss often collapses to sim- ilar value as training proceeds, rendering it ineffective for distinguishing different fine-tuning behaviors.
Few-shot meta-test fine-tuning examines data-limited, few-shot scenario. A single few-shot task is created by ran- domly sampling 5 classes and 5 examples per class from the meta-test dataset, creating a dataset with 25 total train- ing examples. A new (randomly initialized) linear head is attached for the 5-class classification task, and the model finetuned for 100 steps using full batch gradient descent. The transfer accuracy is evaluated on 100 randomly selected test examples for each of the 5 classes. The overall transfer accuracy is averaged over 100 few-shot tasks. This setting corresponds to the "Avg 5-shot Transfer Accuracy" metric (fourth column) in Figure 2. Hyperparameter details for this setting are in Appendix B.2.
In Figure 1, we take each checkpoint at the end of pretrain- ing and plot its pretraining WBIC (called pretraining free energy there since the terminology had not been introduced) versus transfer accuracy. The left (right) plot of Figure I corresponds the third (fourth) column of Figure 2.
Comparison of with other pretraining metrics. As de- scribed in Section 2, recent work of (Galanti et al., 2022) and (Munn et al., 2024) examines the role of neural col- lapse and geometric complexity as effective pretraining met- rics for assessing the suitability of a model checkpoint for transfer learning. To compare the effectiveness of our free energy strategy against these other pretraining metrics, we
Results. In each of these two fine-tuning scenarios, we observe strong correlation between lower pretraining free energy (as measured by the pretraining WBIC, see Section
8

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
conducted a correlation analysis computing the Pearson cor- relation coefficients (Pearson & Galton, 1895) using model checkpoints obtained from training ResNet-18 model on CIFAR-FS to convergence; see Table 1.
proach would be instead identify computationally effi- cient "levers" that influence pretraining free energy, thus allowing us to improve downstream adaptation performance without relying on direct computation of the pretraining WBIC.
These experiments involved a comprehensive exploration of the hyperparameter space (see Appendix B.3). For each checkpoint, we compared the Geometric Complexity, Neu- ral Collapse, and Free Energy of the pretrained model to its downstream performance, measured via both full meta-test fine-tuning and few-shot meta-test fine-tuning. As indicated in Table 1, the pretraining Free Energy cxhibits a substan- tially stronger correlation with downstream performance than other metrics considered.
Impact Statement
This work proposes a novel theoretical framework for under- standing the mechanisms behind successful fine-tuning in machine learning. Our findings have the potential to guide development of more efficient fine-tuning strategies, reduc- ing computational costs and resource consumption, with implications for diverse applications like NLP and computer vision. As the primary focus of this work is theoretical, there are no direct societal consequences of our work that feel must be specifically highlighted.
Pretraining
Finetune Accuracy
Avg 5-shot Accuracy
Metric
Geometric
-0.767
-0.433
Complexity
Neural Collapse
0.632
1875
Free Energy
0.820
0.8901
Acknowledgments
Correlation comparision between pretraining metrics (ge- ometric complexity, neural collapse, and free energy) and down- stream performance (finetune and few-shot transfer accuracy).
Table
We would like to thank Javier Gonzalvo for helpful discus- sions, suggestions, and feedback during the development of this work.
7. Conclusion and Future Work
References
Andriushchenko, M., Varre, A. V., Pillaud-Vivien, L., and Flammarion, N. Sgd with large step sizes learns sparse features. In International Conference on Machine Learn- ing. pp. 903-925. PMLR, 2023.
In this work, we introduced the downstream free energy as Bayesian model selection criterion for quantifying the adaptability of pretraining checkpoints, offering a principled way to predict their performance on unseen downstream tasks. Our key insight is that checkpoints with lower down- stream free energy are more adaptable, making them ideal candidates for fine-tuning. Our empirical results across var- ied datasets (CIFAR-FS, mini-Imagenet) and architectures (ResNet, VGG) validate the utility of the pretraining free en- ergy as a practical checkpoint selection criterion, especially when downstream data is scarce or inaccessible.
Balasubramanian, V. Statistical inference, occam's 's razor, and statistical mechanics on the space of probability dis- tributions. Neural Computation, 9(2):349-368, 1997.
Bengio, Y. Deep learning of representations for unsuper- vised and transfer learning. In Proceedings of ICMI. workshop on unsupervised and transfer learning, pp. 17-
36. JMLR Workshop and Conference Proceedings, 2012.
Despite the promising results, some limitations remain. First, our analysis currently lacks a direct link between downstream free energy and downstream predictive per- formance. At the moment, provide rigorous connec- tion only when downstream adaptation is performed in Bayesian manner (see Appendix A). While Bayesian deep learning is not yet widely adopted due to its computational overhead, this link may become valuable as computational barriers are reduced, particularly fine-tuning scenarios.
Bertinetto, L., Henriques, J. F., Torr, P., and Vedaldi, A. Meta-leariig with differentiable closed-form solvers. In International Conference on Learning Representations, 2019. URL https: / http///pppwwweeepp//tpppppoe.eep net/ forum? id=HyxnZh0ct7.
Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., von Arx, S., Bernstein, M. S., Bohg, J., Bosse- lut, A., Brunskill, E., Brynjolfsson, E., Buch, S., Card. D., Castellon, R., Chatterji, N. S., Chen, A. S., Creel, K. A.. Davis, J., Demszky, D., Donahue, C., Doumbouya, M., Durmus, E., Ermon, S., Htchemendy, J., Ethayarajh, K., Fei-Fei, L., Finn, C., Gale, T., Gillespie. L. E., Goel, K., Goodman, N. D., Grossman, S., Guha, N., Hashimoto, T., Henderson, P., Hewitt, J., Ho, D. E., Hong, J., Hsu, K., Huang, J., Icard, T. F,, Jain, S., Jurafsky, D., Kalluri,
In addition, while our theoretical framework supports the use of free energy as selection criterion, the practical computation of the pretraining WBIC as in (16), remains challenging for large models which may possess tens or hun- dreds of billions of parameters. Developing tractable meth- ods for this computation remains a challenge and presents a significant direction for future work. An alternative ap-
9

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
P., Karamcheti, S., Keeling, G., Khani, F., Khattab, O., Koh, P. W., Krass, M. S., Krishna, R., Kuditipudi, R.,
He, F., Liu, T., and Tao, D. Control batch size and lear- ing rate generalize well: Theoretical and empirical evidence. Advances in neural information processing systems, 32, 2019.
Kumar, A., Ladhak, F,, Lee, M., Lee, T., Leskovec, J., Levent, I., Li, X. L., Li, X., Ma, T., Malik, A., Man- ning, C. D., Mirchandani, S. P., Mitchell, E.,, Munyikwa, Z., Nair, S., Narayan, A., Narayanan, D., Newman, B., Nie, A., Niebles, C., Nilforoshan, H., Nyarko, J. F., Ogut, G., Orr, L., Papadimitriou. I., Park, S., Piech, C.,, Portelance, E., Potts, C.,, Raghunathan, A., Reich, R., Ren, H., Rong, F., Roohani, Y. H., Ruiz, C., Ryan, J., R'e, C., Sadigh, D., Sagawa, S., Santhanam, K., Shih, A., Srinivasan, K. P., Tamkin, A., Taori, R., Thomas, A. W., Tramèr, F., Wang, R. E., Wang, W., Wu, B., Wu, J., Wu, Y., Xie, S. M., Yasunaga, M., You, J., Zaharia, M. A., Zhang, M.. Zhang, T., Zhang, X., Zhang, Y.. Zheng, L., Zhou, K., and Liang, P. On the opportunities and risks of foundation models. ArXiv, 2021. URL https: //crfm. stanford edu wawpttsssssssssss report...... .dff.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Hinton, G. and van Camp. D. Keeping the neural networks simple by minimizing the description length of the weights. In Proceedings of the Sixth Annual Conference on Computational Learning Theory, COLT '93, pp. 5-13, New York, NY, USA, 1993. Association for Computing Machinery. ISBN 0897916115 15. doi: 10. N0.114//183004.666000 URL https: /tse:////oooo/////// 10. 18.1115/111 htt.///6634..116066688 168306.
Jaquier, N., Welle, M. C., Gams, A., Yao, K., Fichera, B., Billard, A., Ude, A., Asfour, T., and Kragic, D. Trans- fer learning in robotics: An upcoming breakthrough? review of promises and challenges. The International Journal of Robolics Research, pp. 02783649241273566. 2023.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Chen, S., Ma, K., and Zheng, Y. Med3d: Transfer learn- ing for 3d medical image analysis. arXiv preprint arXiv: 1904.00044.055 2019.
Kass, R. E. and Raftery, A. E. Bayes factors. Journal of the American Statistical Association, 90(430):733-755, 1995. doi: 10. 1011080001624459.19951110000
URL https: wttp://aaaaaoonieeccooo///o/// 10.7955..147777...
Dherin, B., Munn, M., Rosca, M., and Barrett, D. Why neural networks find simple solutions: The many regu- larizers of geometric complexity. Advances in Neural Information Processing Systems, 35:2333 2022.
10. 11.118000112210
Ke, A., Ellsworth, W., Banerjee, O., Ng, A. Y., and Ra- jpurkar, P. Chextransfer: performance and parameter effi- ciency of imagenet models for chest x-ray interpretation. In Proceedings of the conference on health, inference, and learning, pp. 124, 2021.
Dhillon, G. S., Chaudhari, P., Ravichandran, A., and Soatto, S. A baseline for few-shot image classification. In Inter- national Conference on Learning Representations, 2019.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. P. On hurge-batch training for deep learning: Generalization gap and sharp minima. In Inter- national Conference on Learning Representations, 2017.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.
Kim, J. and Park, C. End-to-end ego lane estimation based on sequential transfer learning for self- self-driving cars. In Proceedings of the IEEE conference computer vision and pattern recognition workshops, pp. 30-38, 2017.
Galanti, T., György, A., and Hutter, M. On the Role of Neural Collapse in Transfer Learning, Jan- uary 2022. URL http: / /arxiv.vvvo org/ haa//ab/2200 /2122. 15121. 11Xiv:2112.15121 [es].
Kornblith, S., Shlens, J., and Le, Q. V. Do better imagenet models transfer better? In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 2661-2671, 2019.
Goyal, P. Accurate, large minibatch sg d: training imagenet in hour. arXiv preprint arXiv: (0..v::77060267777 2017.
Gunasckar, S., Lee, J. D., Soudry, D., and Srebro, N. Im- plicit bias of gradient descent on linear convolutional networks. Advances in neural information processing systems, 31, 2018.
Kumar, A., Raghunathan, A., Jones, R. M., Ma, I., and Liang, P. Fine-tuning can distort pretrained features and underperform out-of-distribution. In International Con- ference on Learning Representations, 2022.
10

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
Lau, E., Furman, Z.,, Wang, G., Murfet, D., and Wei, S. The local learning coefficient: A singularity- complex- ity measure. In The 28th International Conference Artificial Intelligence and Statisties, 2025. URL https: // //openrevice ew. net/ / forum? ? i d= av slsuI.
Papyan, V., Han, X., and Donoho, D. L. Prevalence of neural collapse during the terminal phase of deep learn- ing training. Proceedings of the National Academy of Sciences, (17(40))24652-446) 2020.
Pearson, K. and Galton, R. Vii. note regression and inheritance in the case of two parents. Proceed- ings of the Royal Society of London, 58(347-352): 240-242, 1895. doi: 10.1098/rsp..1895.000 18940000.. URL
Lee, Y., Chen, A. S., Tajwar, F., Kumar, A., Yao, H., Liang, P., and Finn, C. Surgical fine-tuning improves adapta- tion to distribution shifts. In The Eleventh International Conference Learning Representations, 2022.
hctres:/rooyalloocetyypooldiidsd doi/abs /10 1008 /ss///rsp.... 1895.0044
Li, Tang, T., Zhao, W. X., Nie, J.-Y., and Wen, J.-R. Pre- trained language models for text generation: A survey. ACM Computing Surveys, 56(9):1-39, 2024.
Qiu, X., Sun, T., Xu, Y., Shao, Y., Dai, N., and Huang, X. Pre-trained models for natural language processing: A survey. Science China technological sciences, 63(10); 1872-1897, 2020.
Lippl, S. and Lindsey, J. Inductive biases of multi-task learning and finetuning: multiple regimes of feature reuse. Advances in Neural Information Processing Systems, 37: 118745-118776. 2024.
Robert, C. P. et al. The Bayesian choice: from decision- theoretic foundations to computational implementation, volume Springer, 2007.
Liu, H., Xie, M., Li, Z., and Ma, Same pre-training loss, better downstream: Implicit bias matters for lan- guage models. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Pro- ceedings of the 40th International Conference on Ma- chine Learning, volume 202 of Proceedings of Ma- chine Learning Research, pp. 22188-22214. PMLR, 23-
Sanchez, S., Romero, H., and Morales, A. A review: Com- parison of performance metrics of pretrained models for object detection using the tensorflow framework. In IOP conference series: materials science and engineering. volume 844, pp. 01 IOP Publishing, 2020.
Schwarz, G. Estimating the Dimension of a Model. The Annals of Statistics, (22):41-4644 March 1978. ISSN 0090-5364, 2168-8966. doi: 10. 176344133 https::/pgrojecteecclln
29 Jul 2023a. URL https::/t//..c. / http://ooooe/g////// mlr. press/v202 wwt.002/11i2aaa.... html.
URL foorralssammalsso---stttttt.s..
Liu, Y., Zhang, Y., Wang, Y., Hou, F., Yuan, J.. Tian, J., Zhang, Y., Shi, Z., Fan, J., and He, Z. A survey of visual transformers. IEEE Transactions on Neural Networks and Learning Systems, 2023b.
vol volume-6/issue-2)
Est 10 10.1214/20s/1175333 Institute of Mathematical Statistics.
Publisher:
MacKay, D. J. C. Information Theory, Inference Learning Algorithms. Cambridge University Press, USA, 2002. ISBN 0521642981.
Simonyan, K. Very deep convolutional networks for large- scale image recognition. In International Conference on Learning Representations, 2014.
Masters, D. and Luschi, C. Revisiting small batch training for deep neural networks. arXiv preprint arXiv:1804.076122 2018
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. The implicit bias of gradient descent on sep- arable data, Journal of Machine Learning Research, 19 (70):1-77, 2018.
Mormont, R., Geurts, P., and Marée, Comparison of deep transfer learning strategies for digital pathology. In Proceedings of the IEFE conference on computer vision and pattern recognition workshops, pp. 2262-2271, 2018.
Watanabe, S. Algebraic Geometry and Statistical Learning Theory. Cambridge University Press, USA, 2009.
Munn, M., Dherin, B., and Gonzalvo, J. The impact of geometric complexity on neural collapse in transfer learn- ing. In The Thirty- Annual Conference on Neural Information Processing Systems, 2024. URL https: //openreview tere/fformm http://owwmmmdpppiidddood.e.e
Watanabe, S. A Widely Applicable Bayesian Infor- mation Criterion. Journal of Machine Learning Re- search, 14(Mar):867-897, 2013. ISSN ISSN 1533- 7928. URL http: http://wwww /www. jmlr. wwa.lr.org/ppppree
v14/wat 1l4bbettaaabl1 3a. html.
Neyshabur, B., Tomioka, R., Salakhutdinov, R., and Srebro, N. Geometry of optimization and implicit regularization in deep learning. arXiv preprint 1rXiv:1705.0307 2017.
Wen, K., Ma, T., and Li, Z. How sharpness- aware aware minimiza- tion minimizes sharpness? In The Eleventh International Conference on Learning Representations, 2023.

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
A. Theoretical guarantees on fine-tuning predictive performance
Here we discuss theoretical guarantees on downstream predictive performance when employing the version of the down- stream free energy strategy in equation 9. We would like to give an analysis of downstream predictive performance without being tied to a specific training algorithm e.g., SGD with momentum, ADAM, etc. Towards this end, we consider measuring predictive performance through quantities related to the downstream posterior distribution neural network weights:
p² (w; w", 77 x exp(-mk' (x w²)
(17)
This does not mean we are advocating for Bayesian prediction, but rather we believe the posterior distribution above contains highly relevant information that all sensible downstream training algorithms are sensitive to.
Since fine-tuning entails finding small perturbation of said w' which performs well on the downstream training dataset D we might consider an indicator of the downstream fraining performance to be given by
T (w") == Ewwspp (weee* .y) Kk (w).
(18)
Let us call equation 18 the downstream Gibbs training error. Select y such that w* is local minimum of K (w): w" Ug. Then, average, over the draw D', the expected downstream Gibbs training error is given by
i...
(ww) (w=l)- X X'(")
v² (w")
(19)
m
m
where v (w"). like the local learning coefficient A (w"", ), is a positive number called the singular fluctuation that is an invariant of the underlying modee-truth-prior triplet. Since v² (ww") always positive, the strategy equation leads us to select checkpoint that minimizes an upper bound on Epi Tm(w")
We can also look at the population counterpart equation 18 given by
Gm (w") == lewn'' (wmee x) (w)
(20)
Let us call equation 20 the downstream Gibbs test error. The expected value of this, over the draw of DD is given by
Epu Gm( = K' (w")) X >'ww)) + v (w=)
(21)
m
does not appear the strategy in equation 9 gives control over the (expected) downstream Gibbs test error.
Finally consider the test error resulting from Bayesian model averaging:
GBMA
( == (w) (e) (r' (y|x)) (x) [l (" (woe" "o)eee"")PPP/ll ,))(U/ w))
(22)
m
where the expectation over the posterior has been moved inside the logarithm. Let us call equation 22 the downstream Bayes test error. We have that
X² (ww) m
Ep
m CBAA (w*) = K (w=)))
+ o m
(23)
It is evident that the strategy in equation leads us to select checkpoint that minimizes upper bound on Ep GBMA (w*).
B. Experiment details
This section provides details for the experiment results presented in Figure and Figure 2. For these experiments we use the CIFAR-FS dataset (Bertinetto al., 2019) which has been pre-pe partitioo into 64 meta- tiitt-training classes, 14 meta-validation classes and 20 meta-test classes. Each class contains 600 examples. We use the meta-training dataset for pretraining and the meta-test dataset during fine-tuning. We do not use the validdttin dataset.
Random seeds To account for stochasticity, we repeat all experiments below with 5 different random seeds. These seeds control the randomness in the pretraining optimization trajectory, the train-test split and the fine-tuning optimization trajectory in full meta-test finetuning (Section B.2 below), and the construction of few-shot tasks in few-shot meta-test finetuning (Section B.2 below). The variability across the random seeds is reflected in Figure 2, although the error bands may not always be visible due to the wide scale of the y-axis in some cases.
12

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
B.1. Pretraining details
We pretrain a ResNet-18 (He et al., 2016) on the CIFAR-FS meta-training dataset (Bertinetto et al., 2019) using SGD with cross-entropy loss. We vary SGD hyperparameters such the learning rate, batch size, and momentum. We use plain SGD optimizer without any regularization nor schedule to avoid masking effects. We used random crop and random flip for data augmentation. Throughout training we report the pretraining train loss on the augmented data (Figure 2 first column) and the pretraining WBIC computed on the augmented data (Figure 2 second column). Note, we use the same SGLD hyperparameters to compute the WBIC across all experiments. That is, we use step size e = 2 x 10 10-7, chain length of 3,000 iterations, batch size of 2,048, 7 = 1.0, and 3 log n where n is the size of the pretraining dataset.
Learning rate. For experiments that vary the learning rate in Figure 2 (top row), for each learning rate value in {0.01, 0.05, 0.1, 0.2) run SGD without momentum with fixed batch size of 512 for 50,000 iterations. The WBIC estimations were performed every 2,000 iterations with the SGLD hyperparameters above.
Batch size. For experiments that vary the batch size in Figure (middle row), for each batch size in {16, 32, 64, 128, 256, 512] we run SGD without momentum with fixed learning rate of 0.05 for 50,000 iterations. The WBIC estimations were performed every 4,000 iterations with the SGLD hyperparameters above.
Momentum. For experiments that vary the momentum in Figure 2 (bottom row), for each momentum in {0.0, 0.2, 0.4, 0.6, 0.8) we run SGD with fixed learning rate of 0.01 and batch size of 512 for 80,000 iterations. The WBIC estimations were performed every 2,000 iterations with the SGLD hyperparameters above.
B.2. Fine-tuning details
We perform fine-tuning in two scenarios: full CIFAR-FS meta-test finetunining which uses all 20 classes of the meta-test set, and few-shot meta-test finetuning which consists of multiple tasks constructed from the CIFAR-FS meta-test dataset. In both settings we fine-tune ResNet-18 model initializing the weights of the ResNet backbone with the pre-training weights. The weights of the model head are randomly initialized.
Full meta-test fine-tuning. When fine-tuning on the full CIFAR-FS meta-test dataset, we use all meta-test classes and all 600 examples in each class. We then create an 80/20 train/test split, We use SGD with L regularization rate of 0.01 and with fixed learning rate of 0.0001 for the model backbone and fixed learning rate of 0.01 for the model head. We fine-tune for 100 steps using batch size of 128.
Few-shot meta-test fine-tuning. For few-shot fine-tuning, we use only part of the CIFAR-FS meta-test dataset by sampling 5-class classification tasks randomly from the 20 classes available in the meta-test dataset. For each of these 5 classes we sample 5 training examples to create a 5-shot dataset for fine-tuning. During fine-tuning, with full meta-test fine-tuning. we use a fixed learning rate of 0.0001 for the model backbone and fixed learning rate of 0.01 for the model head. We perform 100 steps of full-batch gradient descent (GD) with L regularization rate of 0.001 and then measure the model performance on 100 random test samples from each class. This constitutes a single task. Finally, we report the resulting accuracy rates averaged over 100 randomly chosen tasks.
B.3. Correlation Analysis for Table 1
To assess the effectiveness of our free energy strategy in comparison to these other pretraining metrics, we computed the Pearson correlation coefficients (Pearson & Galton, 1895) for each of the pretraining metrics () against the downstream (full meta-test fine-tuning transfer accuracy, few-show meta-test fine-tuning transfer accuracy) using model checkpoints obtained from experiments with CIFAR-FS, trained on ResNet-18 convergence.
These experiments, detailed in Section 6, involved a comprehensive exploration of the hyperparameter space. We swept across three hyperparameters (learning rate, batch size, and momentum), with six values for learning rate, six for batch size, and five for momentum. Each configuration was trained with five different random seeds, resulting in a total of 85 model checkpoints. For each checkpoint, we compared the Geometric Complexity, Neural Collapse, and Free Energy of the pretrained model to its downstream performance, measured via both full meta-test fine-tuning and shot meta-test fine-tuning. Notably, as indicated by the Pearson correlation coefficients in Table 1, the pretraining Free Energy exhibits a
13

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
substantially stronger correlation with downstream performance than other metrics considered.
C. Proof of Proposition 5.1
Proof. By definition of the test loss and rearranging terms via change of measure, for all w,
r' ('(aa)
K (ww)
log
(zz, v)drdy
pluler, w:)
/
(w(rr)
r (ulx))
r (a, () (xr uddady
log
p(ule,, w)
(vlxr)
r (0 v)
/ log
(vlx))
r (I, v) ,0 (x, v) dx dy
plyla, w)
(x,
+ / log
r (x.y)) dx dy
(y)a))
S MK (w) D.
Also, by definition of w+1, we have K! (w*1) K! (w=). Combining these two facts, we get K (w*) S MK* (w*) + D and obtain the conclusion in (11).
D. Examples of Proposition 5.1
In this section we provide two detailed examples involving Gaussian distributions which help to illustrate Proposition 5.1 in action.
Example (Covariate shift between pretraining and downstream distributions). Suppose (y) r (y)x)) r(yl=)) Our pretraining and fine-tuning joint model is p' (z, y/w) ply|æ,, wx)" Then we have X (²(")) X >'ww)) and (*(w)= = (a) Kxx, w) where K(2, w) = ( w)). Writing
/ K(#, w)
(x) ,0)))0(()) (x) dr x0(x)
E, (x) K(r, w)
(m) (m)
we have that M = max._,"(a))
< oo then
Bxy(2) ((x,w) S ME K K(,, w)
Putting this together we have D and
K (ww") K' (w)) < (w")
Suppose the two covariate distributions are Gaussians
11
(xx) x exp(
20
then M is finite if 0o > 01, in which case M = exp (N0 (#0-M1)²
2(0 -afj
Example (Nuisance parameter mismatch between pretrain and downstream distributions). Suppose the pretrain (i = 0) and downstream (i = 1) distributions are given by
r( v) = rlyl,, wo, of)
where r Wo. of?) = N N(fe, of) with (x) representing neural network with weight w. The pretraining and fine-tuning model are given by
(( = r(y x. w. rx))
Then we have (w²) = 21 (wx)) and M 06/01.
14

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
Learting
0,022
33000
2.33
0.00
8.400
73000
Wreeune
1.43
9:45
60000
90000
1.000
20539
30000
40053
10003
23300
30002
42300
20000
30590
53000
20068
39000
40005
flaviion
Beic1 Sum
0.90
3.5
532
00000
0.45
Sceture
0.90
20000
0.33
10000
20833
30000
40003
10003
23300
30002
42300
12200
20000
30000
80000
33000
20068
70000
40008
Avg Theedte Jocuurecy
30000
2.0
8,000
233
5000
Seysun Sandter
0.50
05
144
Iaretton
Dereton
Barato
Baratoo
Figure 3. Model checkpoints with lower pretraining WBIC (second column) consistently result in better transfer accuracy, both when fine-tuning on the full downstream dataset (third column) and the few-shot setting (fourth column). Lower pretraining WBIC correlates with better downstream performance for Top row: larger learning rates, Middle row: smaller batch sizes, and Bottom row: increased momentum.
E. Additional Experiments for mini-Imagenet; see Figure
E.1. Pretraining details
We pretrain VGG-16 (Simonyan, 2014) the mini-Imagenet meta-training dataset (Dhillon et al., 2019) using SGD with cross-entropy loss. We vary SGD hyperparameters such as the learning rate, batch size, and momentum. We use plain SGD optimizer without any regularization nor schedule to avoid masking effects. We used random crop and random flip for data augmentation. Throughout training we report the pretraining train loss on the augmented data (Figure 2 first column) and the pretraining WBIC computed on the augmented data (Figure 2 second column). Note, we use the same SGLD hyperparameters to compute the WBIC across all experiments. That is, we use step size e = 10-7, chain length of 1,000 iterations, batch size of 1,024, 7 1.0, and 3 log where the size of the pretraining dataset. The results are plotted in Figure 3.
Learning rate. For experiments that vary the learning rate in Figure 2 (top row), for each learning rate value in (0.0025, 0.005, 0.01} we run SGD without momentum with a fixed batch size of 512 for 50,000 iterations. The WBIC estimations were performed every 2,000 iterations with the SGLD hyperparameters above.
Batch size. For experiments that vary the batch size in Figure 2 (middle row), for each batch size in {16, 32, 64, 128, 256, 512} we run SGD without momentum with fixed learning rate of 0.01 for 50,000 iterations. The WBIC estimations were performed every 2,000 iterations with the SGLD hyperparameters above.
Momentum. For experiments that vary the momentum in Figure 2 (bottom row), for each momentum in (0.0, 0.1, 0.3, 0.5} we run SGD with a fixed learning rate of 0.005 and batch size of 512 for 50,000 iterations. The WBIC estimations were performed every 2,000 iterations with the SGLD hyperparameters above.
15

---

A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints
E.2. Fine-tuning details
We perform fine-tuning in two scenarios: full mini- Imagenet meta-test finetunining which uses all 20 classes of the meta-test set, and few-shot meta-test finetuning which consists of In both settings we fine-tune VGG-16 16 model initializing the weights of the VGG backbone with the pre- pre-training weights. The weights of the model head are randomly initialized.
multiple tasks constructed from the mini-lmagenet meta-test dataset.
Full meta-test fine-tuning. When fine-tuning on the full mini Imagenet meta-test dataset, we use all 20 meta-test classes and all 600 examples in each class. We then create an 80/20 train/test split. We use SGD with L regularization rate of 0,01 and with a fixed learning rate of 0.0001 for the model backbone and a fixed learning rate of 0.01 for the model head. We fine-tune for 500 steps using batch size of 32.
Few-shot meta-test fine-tuning. For few-shot fine-tuning, we use only part the mini- minne-mm meta-test dataset by sampling 5-class classification tasks randomly from the 20 classes available in the meta-lest dataset. For each of these classes we sample 5 training examples create 5-shot dataset for fine-tuning. During fine-tuning, with full meta-test fine-tuning, we use a fixed learning rate of 0.0001 for the model backbone and a fixed learning rate of 0.01 for the model head. We perform 100 steps of full-batch gradient descent (GD) with L regularization rate of 0.01 and then measure the model performance on 100 random test samples from each class. This constitutes a single task. Finally, we report the resulting accuracy rates averaged over 100 randomly chosen tasks.
16


Reviews:
[{"review_id": "raMKf3qJS3", "reviewer": "Reviewer_ek2K", "summary": "This paper introduce a Bayesian model selection criterion, called the downstream free energy, to improve fine-tuning performances. There are both theoretical and empirical results provided.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our paper. We are glad you think our paper is \"well-written\" and provides a \"clear statement of results\" supported by both \"theoretical and empirical evidence\". Below, we address the potential weaknesses you mentioned, and we hope these clarifications will encourage you to consider raising your score.\n\n> Weaknesses\n> 1. There is a lack of discussion about the relationship between model generalization performance and energy. Please provide more explanations about this question.\n> 2. The model used in this paper is some small models. Are there the same results guaranteed on larger models?\n\nIn regards to Weakness 1: Thank you for emphasizing the importance of the relationship between model generalization performance and energy. We agree that explaining this connection is crucial. In fact, Section 5.1 of our paper already provides a detailed discussion of how the free energy (and its complexity term) interacts with test loss to shape model generalization, through a series of Observations tied to our main proposition. If these details were inadvertently overlooked, we kindly invite you to revisit that section and let us know if anything remains unclear or incomplete.\n\nIn regards to Weakness 2: The theory we develop here does not make any assumptions on the model size or complexity. So, yes the same results should hold for larger models as well. However, as we state in the Section 7 'Conclusion and Future Work', the bottleneck is computation which can be challenging for very large models. This is an intriguing area of future work and we also suggest some alternative approaches to address this limitation there.", "claims_and_evidence": "Yes. Section 5 is about theoretical results, and empirical results are in Section 6.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes. Proposition 5.1 is correct.", "experimental_designs_or_analyses": "Yes. Section 6 is about empirical results. And there are also some empirical details in appendix.", "relation_to_broader_scientific_literature": "This paper is mainly related to model generalization performance, which mainly focusing on the controlling the upper bound of test error, using the information on training error and function class.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths\n1. This paper is well-written, which provide a clear statement of the results.\n2. There are both theoretical and empirical evidences to support the idea.\n\nWeaknesses\n1. There is a lack of discussion about the relationship between model generalization performance and energy. Please provide more explanations about this question.\n2. The model used in this paper is some small models. Are there the same results guaranteed on larger models?", "other_comments_or_suggestions": "See strengths and weaknesses.", "questions_for_authors": "See strengths and weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qDCTBuhLbD", "reviewer": "Reviewer_UWMc", "summary": "This paper proposes a new metric, pretraining free energy, which can be used to find a pretraining model checkpoint which is most adaptable for downstream finetuning tasks. The paper is largely theoretical, justifying this metric, although there are two experiments (one in appendix) showing that WBIC, which is used to approximate the pretraining free energy, correlates with downstream finetuning performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your careful attention to our paper and thoughtful review. We are glad you think our paper is \"clear\" and that the \"structure of the narrative was good\". We will do our best to answer your concerns regarding potential weaknesses below. \n\n> Experimental Designs or Analyses: I think some more variety would be good, rather than just focusing on image classification\n\nWe fully agree that a more diverse suite of experiments beyond image classification would show that the correlation between WBIC and downstream performance is not coincidental. Thank you for the suggestion, and we look forward to addressing it in follow-up work.\n\n> Experimental Designs or Analyses: it also sees (by eye) that there is a close link between strong pretraining performance and strong downstream performance; ruling out n empirical link there, as is discussed in Observations 1 and 2, would help the experimental results I think.\n\nThe reviewer suspects a confound: maybe WBIC (and strong downstream performance) both correlate with strong pretraining performance, rather than with each other. We actually have some counterexamples to this, for instance the third row of Figure 2. Note that towards the end of training, all momentum values share the same pretraining loss yet the downstream performance is quite different; the pretraining WBIC can pick this up. \n\n\n> Other Strengths and Weaknesses: Overall, I found most of the paper clear [...] the structure of the narrative was good in building up a more complete picture of the method being presented. That said, I found the paragraph starting on line 196 (left hand side), about how the checkpoints considered are not actually checkpoints, a bit confusing. I also found proposition 5.1 hard to follow.\n\nIn regards to the question about \"checkpoints considered are not checkpoints\": We apologize for the confusion. We will clarify that “pretraining checkpoints” in our theoretical discussion refers to local minima of the test loss, which may differ from the actual checkpoints saved during training. We will further clarify that in order for the theory to match the empirical analysis, we stipulate the actual checkpoints saved during training are local minima of the training loss.\n\nRegarding Prop 5.1 being hard to follow, do you mean that the statement of the proposition itself is hard to follow, the proof, or the discussion of how prop 5.1 is used to justify Eq (10), or something else? \n\n> Questions For Authors:\n>  1. How computationally costly is it to compute WBIC - is this something that drastically increased the overhead of the experiments in the paper?\n>  2. Do you think this method is both feasible, and will continue to scale, for much larger models - for example, 70B or 300B parameter LLMs?\n>  3. As a follow up ot the above - is this still useful if it doesn't, given this is the arguably the most significant field for pretraining and finetuning.\n>  4. How confident are you that this finding would apply in fields beyond image classification, as presented in the CIFAR and ImageNet-mini experiments?\n\n1. The original WBIC is very costly to compute. The local WBIC computed in this paper is much cheaper because the localizing prior forces the exploration to stay close to some parameter $w^*$. We compute local WBIC through SGLD sampling which is computationally efficient for deep learning models.\n2. Yes, there is no fundamental obstacle preventing the approach from scaling to much larger architectures—provided sufficient computational resources. In principle, this includes LLMs on the order of tens or hundreds of billions of parameters.\n3. see above, \n4. Our current theoretical results rely on conditions (e.g., mild distribution shifts) that are plausibly satisfied in the image classification tasks we considered. We are cautiously optimistic this approach could generalize to other tasks as well—potentially including text domains—but confirming that the same assumptions hold there would likely require additional theoretical and empirical investigation and is the focus of future work.\n\nThank you again for your insights, which will help make our paper better!", "claims_and_evidence": "The paper makes a number of claims:\n- That downstream free energy is a good proxy for the downstream performance of a model after finetuning. This seems to be generally justified via theory, and seems to hold based on the arguments presented in the paper.\n- That pretraining free energy is a more measurable alternative to downstream free energy which upholding similar performance prediction characteristics. Again, I think this hold, although I admit (discussed below), that I found this discussion relatively confusing - possibly a function of my background not aligning to that of the paper.\n- The WBIC can be used to approximate pretraining free energy without requiring expensive (and quite possibly intractable) integration calculation. This is demonstrated empirically for two experiments, and seems to hold, though as an empirically minded researcher I would have possibly liked to see this demonstrated in a couple of additional domains to ensure the finding holds generally (possibly in a task that was not image classification). That said, the key contribution of this paper is theoretical and so I do not believe that this limits the correctness or validity of the work.", "methods_and_evaluation_criteria": "As stated above, this paper is principally theoretical (despite, admittedly, dealing with a very empirical topic). As such, while I believe consideration of additional benchmarks would be good - possibly for larger models, such as LLMs, given this is where a large part of the pretrain-then-finetune regime gains have proven fruitful - I think the paper should be viewed with a more theory-focused lens. In this perspective, I believe the benchmarks used are valid and, while the empirical effectiveness of the work could be boosted, the experiments run in this paper provide enough support for it to stand.\n\nI appreciated the narrative, in which the method was built up - I felt this had a very logical flow, and took practicality into consideration (which can be rare for theoretical papers). As such,  the transition from downstream free energy -> pretraining free energy -> WBIC was very logical and I think makes a lot of sense for the problem at hand.", "theoretical_claims": "I attempted to follow the theoretical claims made throughout the paper, but admit that this work is beyond my background and thus I did not always follow 100%. One thing I was unsure about was why the terms in the pretraining free energy are stochastic, whereas they weren't for downstream free energy, and think a qualitative sentence explaining this would provide some needed clarity.\n\nI would suggest that other reviewers be more emphasised in this regard.", "experimental_designs_or_analyses": "The experiments seeem reasonably designed, and there is a description of the pretraining process and fine-tuning details in the appendix (including hyperparameters). I think some more variety would be good, rather than just focusing on image classifciation, to truly verify whether the correlation between WBIC and downstream performance is legitimate or coincidence (though attached to the theory, I think it should hold). I particularly emphasise this as it also sees (by eye) that there is a close link between strong pretraining performance and strong downstream performance; ruling out n empirical link there, as is discussed in Observations 1 and 2, would help the experimental results I think.", "relation_to_broader_scientific_literature": "The paper seems to contextualise themselves well against prior literature, although I am not an expert in this field. There are cmparisons against certain metrics which have been proposed in prior literature for similar problems.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Overall, I found most of the paper clear as a reader (particularly one who works in a different area to this work). I thought the structure of the narrative was good in building up a more complete picture of the method being presented.\n\nThat said, I found the paragraph starting on line 196 (left hand side), about how the checkpoints considered are not actually checkpoints, a bit confusing. I also found proposition 5.1 hard to follow.\n\nBesides that, I felt this was a good paper.", "other_comments_or_suggestions": "On line 267, in the right hand column, there is a missing full stop.", "questions_for_authors": "- How computationally costly is it to compute WBIC - is this something that drastically increased the overhead of the experiments in the paper?\n- Do you think this method is both feasible, and will continue to scale, for much larger models - for example, 70B or 300B parameter LLMs?\n- As a follow up ot the above - is this still useful if it doesn't, given this is the arguably the most significant field for pretraining and finetuning.\n- How confident are you that this finding would apply in fields beyond image classification, as presented in the CIFAR and ImageNet-mini experiments?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yn6VJA1k1H", "reviewer": "Reviewer_Jv2Y", "summary": "This paper introduces a Bayesian model selection criterion called the downstream free energy, which quantifies the adaptability of pretraining checkpoints for downstream tasks. By measuring the concentration of favorable parameters for the task, this criterion helps predict fine-tuning performance without requiring access to downstream data or prior task knowledge. Empirical evidence validates that the criterion reliably correlates with improved fine-tuning performance.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our paper. We note that the reviewer raised two concerns—(1) the absence of certain references ([1–5]) and (2) the limited dataset scope (CIFAR-100 and mini-ImageNet)—and offered no further questions or objections. You’ll find below our best efforts to address these concerns. Please consider raising your score if you are satisfied with them.\n\n> Essential References Not Discussed: There are several works [1-3] focusing on assessing the reusability or transferability of pre-trained models. However, this paper does not discuss these works.\n\nIndeed there are several studies (including those [1-3] mentioned here) which examine how to quantify transferability of pre-trained models. Below, in “Other Strengths and Weaknesses,” you also reference [4] and [5], but do not categorize them as “essential.” Can you please specify why you feel these particular works [1-3] are essential to the scope of our current paper? In particular, can you please clarify how these works directly inform or extend our results? Provided this clarification, we are happy to include these or any other references we would have accidentally missed. \n\n> Other Strengths and Weaknesses: Selecting pre-trained models for downstream tasks is a field with many existing works [1-5], but this paper does not discuss the differences from these works, nor does it compare them with these methods in the experiments.\n\n(Related to above) Our paper includes comparisons with established measures such as geometric complexity and neural collapse. These are equally heuristic or empirical in nature and comparable to [1–5]. Since our focus is on a Bayesian model selection approach, not on exhaustively benchmarking all transferability metrics, we believe our chosen references are sufficient to position this work in the broader literature. Can you please specify how [1–5] directly inform or critique the Bayesian framework we adopt in our paper? If so, we are happy to include these or any other references we would have accidentally missed. \n\n> Other Strengths and Weaknesses: The experiments only involve two datasets, CIFAR-100 and mini-ImageNet, which are relatively few in number and have a small dataset size for each.\n\nIn regards to our experiments, we used CIFAR-100 and mini-ImageNet because they are well-established benchmarks that allow rapid, reproducible testing of our approach. We view exploring larger datasets as an orthogonal direction that would not alter our main theoretical contributions. We appreciate your feedback and remain open to expanding our experiments to additional datasets in future work.", "claims_and_evidence": "The claims made in the submission are supported by evidence.", "methods_and_evaluation_criteria": "Yes, the proposed methods make sense for the problem at hand.", "theoretical_claims": "I have generally checked the proofs, but some details have not been thoroughly verified.", "experimental_designs_or_analyses": "I have generally reviewed the experimental design, which seems reasonable.", "relation_to_broader_scientific_literature": "The paper contributes to the study of Bayesian model selection.", "essential_references_not_discussed": "There are several works [1-3] focusing on assessing the reusability or transferability of pre-trained models. However, this paper does not discuss these works.\n\n[1] Tran et al. Transferability and Hardness of Supervised Classification Tasks. ICCV 2019.\n\n[2] Nguyen et al. LEEP: A New Measure to Evaluate Transferability of Learned Representations. ICML 2020.\n\n[3] You et al. LogME: Practical Assessment of Pre-trained Models for Transfer Learning. ICML 2021.", "other_strengths_and_weaknesses": "1. Selecting pre-trained models for downstream tasks is a field with many existing works [1-5], but this paper does not discuss the differences from these works, nor does it compare them with these methods in the experiments.\n\n2. The experiments only involve two datasets, CIFAR-100 and mini-ImageNet, which are relatively few in number and have a small dataset size for each.\n\n[1] Tran et al. Transferability and Hardness of Supervised Classification Tasks. ICCV 2019.\n\n[2] Nguyen et al. LEEP: A New Measure to Evaluate Transferability of Learned Representations. ICML 2020.\n\n[3] You et al. LogME: Practical Assessment of Pre-trained Models for Transfer Learning. ICML 2021.\n\n[4] Guo et al. Identifying Useful Learnwares for Heterogeneous Label Spaces. ICML 2023.\n\n[5] Zhang et al. Model Spider: Learning to Rank Pre-Trained Models Efficiently. NeurIPS 2023.", "other_comments_or_suggestions": "I have no other suggestions.", "questions_for_authors": "I have no other questions beyond those already mentioned.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "trCQ66g7L8", "reviewer": "Reviewer_dKNM", "summary": "This paper studies the problem of neural network model selection under the pretrain-then-adapt paradigm. Based on the pretraining data, multiple neural network checkpoints can be obtain roughly corresponding to different local minimums of the network parameters. To select a better choice that adapts well to downstream tasks, pretraining and downstream free energies are introduced as Bayesian model selection criteria. To deal with cases where downstream data are not available during model selection, relations between the two energies are explored so that an approximation of downstream free energy can be designed based on pretraining data, which the authors further refer to as pretraining WBIC. Several numerical experiments are conducted.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to read our work. We are concerned that the major criticisms do not fully justify a “reject” recommendation. Below, we show that the issues raised can be readily resolved with minor clarifications or references, rather than indicating any fundamental flaw.\n\n> Claims and Evidence: This paper is not self-contained..for Eq.(4) no derivation is given...also, the complexity term $\\lambda^1(\\omega^{*1}$) lacks its definition.\n\nAs we state in the paper, it is possible to derive Eq (4) using techniques set out in Watanabe’s book, the details of which can be found in [Lau, 2023]. We will add a precise reference to the relevant sections/pages in [Lau, 2023]. Because this expansion is well-established in singular learning theory, many works reference it rather than re-deriving it fully. We hope you will agree that adding explicit sections/pages to our reference of [Lau, 2023] is a minor fix that does not warrant rejection.\n\nIn regards to the definition of $\\lambda^1({\\omega^{\\ast 1}})$, we believe here you mean $\\lambda^1({\\omega^{\\ast}})$ since the former does not appear in the paper. With respect to $\\lambda^1({\\omega^{\\ast}})$, recall that this quantity (which represents the complexity measure of $\\omega^{\\ast}$) is defined implicitly as the $\\log m$ coefficient in Eq (4). We feel this is sufficient since we never have to reckon with this term on its own, only as it appears in the asymptotic expansion in (4).\n\n> Claims and Evidence: The intuitive argument and...the definiion of downstream free energy in Eq.(1) .. is not very compelling. It seems that the scale of parameters would strongly affect the energy...I also wonder how does the use of batch/layer normalization affect the effectiveness of this downstream free energy.\n\nRegarding scale, this is what we think the reviewer is expressing, please correct us if we’ve misinterpreted. The reviewer is hinting that in neural network architectures with some form of scale invariance such as ReLU networks, multiplying the entire parameter set by some constant might not change the function’s outputs—and thus wouldn’t degrade or improve downstream accuracy—while the free energy quantity we define could shift in a non-trival way. We think this is a valid point. However our experiments and our intention revolve around realistic neural networks that are deployed in practice which rarely exhibit strict parameter scaling invariance. We will add a brief discussion in the final version to clarify this point. Thank you for raising this.\n\nRegarding batch/layer norm, we note that for our experiments we trained models with (e.g. Resnet) and without batch norm (e.g. VGG) and see the same effect. We do not expect this factor would affect the effectiveness of our approach. \n\n> Supplementary: ...In line 235 column 2, it says \"under mild assumptions\"...The assumption is missing from both the main paper and the supplementary material.\n\nIn the main text at line 235 col 2, we wrote a parenthetical “under mild assumptions, below” to refer to the assumptions in Propn 5.1. We will ensure the final version states these assumptions explicitly, The broken reference in the supplementary will also be corrected.\n\n> Relation to Broader...: How does this differ from the commonly used WBIC criteria in Bayesian model selection. Is this free energy idea a new interpretation for WBIC, or is the proposed method actually different from WBIC?\n\nOur localized WBIC can be viewed as the classical WBIC with a Gaussian prior centred on a pretraining checkpoint. While we have taken care to rigorously define our localized WBIC, we welcome the suggestion to make the distinction with classical WBIC more explicit and will add a dedicated paragraph discussing how our “pretraining WBIC” compares with the classical WBIC. We see this as a straightforward clarification that in no way invalidates our approach.\n\n> Essential References Not Discussed\n\nWe cited Lau et al. (2023) for our local WBIC approach but we are happy to cite the original WBIC paper as you suggest. Thank you. However, we do not consider references on classic WAIC that you mention here to be relevant for our work. Can you please articulate in which sense the WAIC references are essential or give some more detail as to how you see WAIC directly intersecting with our methodology?\n\n> Other Strengths and Weaknesses: The idea is interesting ..but the paper is not well written or self-contained.\n\nWe hope that our planned clarifications around Equation (4) and any added references will address your concern about self-containment. \n\nIn regards to being \"not well-written\", can you please provide more precise feedback on any sections which remain confusing or unclear. We note that Reviewers ek2K and UWMc explicitly praised our writing, but we will certainly incorporate any further suggestions to improve readability. Could you please indicate which specific aspects of the writing needs improvement so we can address them directly?", "claims_and_evidence": "1. This paper is not self-contained. As one of the most important equations, only Watanabe's book is referenced for Eq.(4) but no derivation is given. If Eq.(4) is directly taken out of the book, maybe a detailed reference to sections&pages would help. Similarly, the complexity term \\lambda^1(\\omega^{*1}) also lacks its definition.\n\n2. The intuitive argument \"Intuitively, lower downstream free energy indicates a higher concentration of parameters in parameter space for which the model is more adaptable and capable of generalizing well on downstream tasks\" along with the definiion of downstream free energy in Eq.(1) using the parameter ball B_\\gamma(\\omega^*) is not very compelling. It seems that the scale of parameters would strongly affect the energy, which not necessarily affect the model performance. For non-neural network examples, this energy definition may not be ideal for a local-scale mixture model. Similarly, I also wonder how does the use of batch/layer normalization in the neural network affect the effectiveness of this downstream free energy.", "methods_and_evaluation_criteria": "See Claims And Evidence.", "theoretical_claims": "See Claims And Evidence.\n\nProposition 5.1 seems correct.", "experimental_designs_or_analyses": "Seems fine.", "relation_to_broader_scientific_literature": "in Section 5.2, the authors estimates the pretraining free energy using Eq.(14), which is called the pretraining WBIC. How does this differ from the commonly used WBIC criteria in Bayesian model selection. Is this free energy idea a new interpretation for WBIC, or is the proposed method actually different from WBIC? If it's the latter, perhaps more discussion of their distinctions could be discussed.", "essential_references_not_discussed": "Here are some key relevant references not cited/discussed:\nWAIC: Watanabe, Sumio. \"A widely applicable Bayesian information criterion.\" The Journal of Machine Learning Research 14.1 (2013): 867-897.\nWBIC: Watanabe, Sumio. \"A widely applicable Bayesian information criterion.\" The Journal of Machine Learning Research 14.1 (2013): 867-897.\nMore on WAIC: Gelman, Andrew, Jessica Hwang, and Aki Vehtari. \"Understanding predictive information criteria for Bayesian models.\" Statistics and computing 24 (2014): 997-1016.\nWAIC for latent variable models (potentially related to neural networks): Merkle, Edgar C., Daniel Furr, and Sophia Rabe-Hesketh. \"Bayesian comparison of latent variable models: Conditional versus marginal likelihoods.\" Psychometrika 84.3 (2019): 802-829.", "other_strengths_and_weaknesses": "The ideas is interesting and worth further exploration, but the paper is not well written or self-contained at this point.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See Claims And Evidence & Relation To Broader Scientific Literature.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel Bayesian model selection criterion based on the concept of downstream free energy, which quantifies the adaptability of a model checkpoint to downstream tasks by measuring the concentration of favorable parameters in its vicinity. A key strength of the method is that it can be computed without access to downstream data or prior knowledge about the downstream tasks. The proposed criterion is theoretically motivated, practically useful, and shown to be strongly correlated with downstream performance.

Overall, the reviewers appreciated the novelty and quality of the submission, and I share their positive assessment. The idea of downstream free energy is both original and insightful, with sound theoretical backing and practical implications. While some concerns were raised regarding the omission of related work, I believe these issues are minor and can be addressed in the final version. I urge the authors to ensure a thorough revision addressing all reviewer comments, such as incorporating the missing references.
Decision: Accept (poster)