Leveraging free energy in pretraining model selection for improved fine-tuning

Thank you for your comments. We are happy that you found that the method clear and appreciate our measure of downstream adaptability which "relies solely on the pretraining datasets". We have addressed your comments and questions to the best of our ability below.

Please consider raising your score or confidence if your concerns have been resolved. Thank you.

Weakness 1:

"...it lacks comparative analysis with prior research..."

Please see our top-level response on Comparisons at the top of this page.

Weakness 2:

"There are concerns regarding the validity of the theoretical assumptions, Assumptions 5.1 and 5.2..."

Assumption 5.1 is grounded in the original literature on Local Learning Coefficients (LLC), specifically regarding LLC estimation at local minima. Additionally, we noted on Line 330 that Assumption 5.2 follows from Yamazaki et al. (2007), which analyzed distribution shift scenarios. We believe these are rather mild assumptions actually.

We have also included an additional statement in the manuscript addressing the feasibility of these Assumptions (see the subsection in blue titled "Interpretation and Feasibility of Assumption 5.2"). To provide more context, please also see our response to Reviewer SycF’s Weakness 4, where we further clarify Assumption 5.2 by explaining it in terms of distributional support—emphasizing the need for the pretraining distribution to cover a sufficiently large support relative to the downstream distribution.

Weakness 3 / Question 4:

The experimental setup is limited

Please see our top-level response on Additional Experiments. We have now included additional experiments to the manuscript that demonstrate the same relationship.

Question 1:

...the proposed pretraining free energy seems to serve as a general model selection criterion, raising questions about its rationale and necessity, since model selection is often focused on specific downstream tasks.

Thank you for raising this question. We understand the concern about balancing a general model selection criterion with performance on specific downstream tasks. Our approach prioritizes generality, aiming to provide a selection criterion that applies across a wide range of potential downstream tasks. While tailoring checkpoints for specific tasks might improve performance in some cases, the pretraining free energy criterion is designed to work without such task-specific adjustments, making it practical and versatile in scenarios where downstream tasks may not be known during pretraining.

Question 2:

Assumption 5.2 appears overly generalized. It may not hold when two distributions overlap insufficiently or when higher-order statistical moments differ significantly.

In regards to the feasibility of Assumption 5.2 for real-world settings, we have included a subsection addressing the interpretation and feasiblity of this assumption (in blue in the updated manuscript).

In short, we specifically focused on experimental settings where the pretraining dataset is much larger and more complex than the downstream dataset; cf. Kornblith et al., 2019. In our experiments, (and as reflected in practice) we achieved this by using pretraining datasets with a substantially larger set of image classes than the downstream dataset. We agree that if this were reversed; i.e., the pretraining dataset had substantially fewer classes than the downstream dataset, the relationship we establish in Prop 5.3 is uninformative. This would be similar to the example you provide where $r^0(x,y) \sim N(0, 0.1)$ and $r^1(x,y) \sim N(0, 1).$ Taken to this extreme, we also agree it would not make sense to apply our pretraining free energy selection criterion.

Question 3:

The distinction between the downstream free energy strategy proposed by the authors (line 245) and existing free energy criteria is unclear.

Note that we did not claim the free energy criterion for model selection is a novel contribution, and we highlight its wide usage in statistics in the "Relationship to Prior Work" where we discuss the distinction with the classic model selection criterion. Instead, the novelty and value of our work lies in applying this classic criterion to the area of pretraining and fine-tuning, where it has not been previously examined as a model selection tool for evaluating checkpoint adaptability. Our goal in Section 4.1 was to provide the necessary theoretical background for readers less familiar with free energy in the context of model selection.

Thank you for responding to my questions, they have resolved some of my concerns. However, the reply to Q2 cannot fully address my concerns.

Theoretically, Assumption 5.2 remains challenging to satisfy in cases where the pretraining dataset contains more classes than the downstream dataset (Simple case: $r^0(x, y)\sim p_0(y=0)N(0,0.1)+p_0(y=1)N(1,0.1)+p_0(y=2)N(2,0.1), r^1(x,y)\sim p_1(y=0)N(0,1), p_0(y=i)=1/3, p_1(y=0)=1$). I recommend introducing additional constraints on this assumption, such as specifying certain statistical properties of the distributions, to enhance its general applicability. Otherwise, this assumption might not be broadly perceived as valid.

Furthermore, I notice that Reviewer SycF raised W1, which also touches on the distributional differences between $r^0$ and $r^1$. I recommend including experimental results to validate this assumption. For example, you could use GMMs to estimate the feature distributions of the pretraining and downstream datasets as proxies for $r^i$, and verify whether the required distributional differences and proportions hold.

We would like to thank the reviewer for recognizing the relevance of our work, as well as for appreciating our theoretical approach. We are glad that the paper was generally clear and accessible, even with the mathematical depth required by our analysis.

We are working diligently on point-by-point responses to all reviewers, but we wanted to immediately address your concerns about the general applicability of our approach. Below, we provide clarifications, which we hope will highlight the strength of our contributions and eventually encourage a reconsideration of the scores.

Theoretical concerns (Weakness 1 and 4)

We believe there may be a misunderstanding regarding the role of the constant $D$, which we will first clarify. In Proposition 5.3, the constant $D$ is a very minor character. Recall our goal is to relate $mK^1(w^{\ast 1}) + \lambda^1(w^{\ast}) \log m$ to a quantity that only uses $K^0$ and $\lambda^0.$ We begin by writing $K^1(w) = f(w) + D,$ where $f(w)$ is the first expression in Line 344. We then establish $M K^0(w)$ as an upper bound on $f(w)$ which naturally extends to an overall upper bound on $K^1(w)$ that is $MK^0(w)+D.$

Thus we establish that $K^1(w)=f(w)+D < MK^0(w) + D$. Our statement about disregarding the constant $D$ during optimization is in regards to the inequality $f(w) + D < MK^0(w) + D$. Note that the toy example you suggest where minimizing $g(x)=0$ but having $f(x) = 10e5$ simply cannot occur in our setting.

Thank you for pointing out this confusion; we will clarify this point in the text.

Next, we would like to address the real-world applicability of the relationship between pretraining and downstream free energy established in Proposition 5.3, as we believe this is your primary concern regarding the meaningfulness of our contribution. Your example of horse versus car images is a valuable thought experiment. You are correct that when $r^0(x,y)$ and $r^1(x,y)$ have disjoint label support, this would violate Assumption 5.2, which is specifically designed to prevent this situation by requiring controlled overlap between pretraining and downstream distributions. Specifically, if the support of the pretraining distribution $r^0$ is too small relative to the support of the downstream distribution $r^1$, the constant $M$ would become infinite, violating Assumption 5.2. Finally, you are correct that Yamazaki et al., particularly in their synthetic examples, actively avoid situations where Assumption 5.2 is violated.

However, note that, as you correctly observe in your Weakness #4, in order to more reasonably satisfy Assumption 5.2, we specifically focused on experimental settings where the pretraining dataset is much larger and more complex than the downstream dataset. This setting has also been studied in prior work, as you noted (e.g., Kornblith et al., 2019). In our experiments, we achieved this by using pretraining datasets with a substantially larger set of image classes than the downstream dataset. We agree that if this is reversed; i.e., the pretraining dataset has substantially fewer classes than the downstream dataset, the relationship we establish in Prop 5.3 is uninformative. Taken to the extreme, we also agree it would be quite silly to apply our pretraining free energy selection criterion if the pretraining dataset contains only horse images and the downstream dataset contains only car images.

In summary, we appreciate your insights, which have highlighted the importance of interpreting Assumption 5.2 with practical considerations in mind. We will incorporate this perspective into the paper to clarify its implications for real-world applications in the pretrain-then-adapt paradigm.

Thank you for your comments. We are very happy that you found the “rigorous theoretical analysis" insightful and find that our asymptotic pretraining free energy strategy could be "helpful for pretraining". As you point out, the pretraining free energy does "not rely on downstream task data" which we believe makes it a valuable and widely applicable proxy for assessing the quality of a pretraining checkpoint. We have addressed your comments to the best of our ability below.

Please consider raising your score or confidence if your concerns have been resolved. Thank you.

Weakness 1:

The experimental section is overly simplistic...

Please see our top-level response on Additional Experiments. We have included additional experiments to the manuscript that demonstrate the same relationship.

Weakness 2:

The lack of comparison with existing methods weakens the persuasiveness of the results.

Please see our top-level response on Comparisons.

Weakness 3:

"Full meta-test fine-tuning" and "Few-shot meta-test fine-tuning" should be presented as parallel points.

Thank you. We have incorporated your suggestions on presenting “Full meta-test fine-tuning” and “Few-shot meta-test fine-tuning” in parallel and improving consistency in cross-referencing formats and symbol representations throughout the paper.

Question 1:

Does the theory proposed offer additional insights or applications wrt learning rates, batch sizes, and momentum?

While it’s true that prior work has suggested that larger learning rates, smaller batch sizes, and increased momentum can improve transfer performance, our contribution lies in verifying these findings within a rigorous theoretical framework. By using pretraining free energy as a selection criterion, we confirm that these strategies are indeed effective for improving downstream transferability.

Question 2:

...the paper assumes that the classification head $v$ of the pretraining task and $u$ of the downstream task share the same dimensionality. How is this assumption specifically applied in the theoretical analysis?

Thank you for giving us the chance to clarify this. By assuming $u$ and $v$ are of the same dimensionality, we can unambiguously write $p(y|x,w)$ to refer to both the pretraining and fine-tuning model. This shows up in the rest of the theoretical development which only refers to $p(y|x,w)$.

Question 3:

Can the proposed method be applied to unsupervised pretraining processes and situations where downstream tasks differ from the pretraining tasks?

Yes absolutely, the proposed method can be applied to unsupervised pretraining. The pretraining free energy can be then defined for the model $p(x|w)$ rather than $p(y|x,w)$.

As a followup to the note on Comparisons above, we have also conducted further analysis quantitatively comparing the pretraining free energy with the pretraining geometric complexity and neural collapse. Please see also a detailed section in the Appendix (see Appendix E, in blue) in the updated draft.

In short, to assess the relationship between pretraining metrics and downstream performance, we computed Pearson correlation coefficients. These correlations were calculated between three pretraining metrics (geometric complexity, neural collapse, and free energy) and two downstream fine-tuning metrics (full fine-tuning transfer accuracy and average 5-shot transfer accuracy). We utilized model checkpoints obtained from our CIFAR-FS experiments, with models trained on ResNet-18 to convergence.

As shown in the table below, pretraining Free Energy demonstrates a substantially stronger correlation with downstream performance compared to the other evaluated metrics.