A Bayesian Model Selection Criterion for Selecting Pretraining Checkpoints

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.

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.

As 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.

In 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).

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.

Regarding 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.

Regarding 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.

Supplementary: ...In line 235 column 2, it says "under mild assumptions"...The assumption is missing from both the main paper and the supplementary material.

In 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.

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?

Our 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.

Essential References Not Discussed

We 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?

Other Strengths and Weaknesses: The idea is interesting ..but the paper is not well written or self-contained.

We hope that our planned clarifications around Equation (4) and any added references will address your concern about self-containment.

In 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?

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.

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.

Indeed 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.

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.

(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.

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.

In 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.

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.

Experimental Designs or Analyses: I think some more variety would be good, rather than just focusing on image classification

We 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.

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.

The 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.

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.

In 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.

Regarding 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?

Questions For Authors:

1. How computationally costly is it to compute WBIC - is this something that drastically increased the overhead of the experiments in the paper?
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?
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.
4. How confident are you that this finding would apply in fields beyond image classification, as presented in the CIFAR and ImageNet-mini experiments?

1. 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.
2. 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.
3. see above,
4. 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.

Thank you again for your insights, which will help make our paper better!

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.

Weaknesses

1. There is a lack of discussion about the relationship between model generalization performance and energy. Please provide more explanations about this question.
2. The model used in this paper is some small models. Are there the same results guaranteed on larger models?

In 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.

In 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.