Early Period of Training Impacts Adaptation for Out-of-Distribution Generalization: An Empirical Study

Thank you for your time and constructive feedback.

W1/Q1: Comparisons with Baselines In this paper, our goal was to focus on evaluating the effect of gradual unfreezing (as a tool) specifically within early training dynamics, rather than proposing a SOTA method. We selected gradual unfreezing as an intervention method due to prior studies on its effectiveness in achieving SOTA in various OOD settings.

W2: Gradual unfreezing shares similarities with the approach in [1] for transfer learning with a pre-trained backbone, as both involve selectively tuning parameters over time to enhance OOD performance. Gradual unfreezing can be seen as a variant of the approach in [1], offering more fine-grained algorithmic control by tuning layer-by-layer instead of focusing solely on the classification head and backbone. In Section 5.2, we demonstrate that gradual unfreezing aligns gradient similarity across batches during early training, particularly in the classification head, further supporting this connection.

W3: Causal Interpretation of Sharpness and OOD Performance Thank you for this interesting suggestion. While our experiments establish a correlation, causality would indeed require interventions to control sharpness independently of other factors. Optimization processes (back-propagation / parameter updates), inherently affect sharpness/Fisher Information regardless of freezing or unfreezing. Currently, we lack an effective method to isolate and control these changes, which we acknowledge as a limitation and included in our limitations section (L918). We have some additional experiments in Appendix F.2 and F.3, where sharpness is regularized midway through training, yet also yield better test OOD results, further suggest that early sharpness/Fisher Information (FI) measurements may be noisy.

Q2:How does the gradual unfreezing strategy compare to other parameter-freezing approaches during early training? Could the similarities and differences be further clarified? We used MNIST with ResNet18 training from scratch and we visualized the training dynamics for LPFT (head first, then unfreeze the layers in the backbone all together) [1] and reverse unfreezing (bottom-up). Using LPFT and with a reverse selection of layers to unfreeze during training after the unfreezing of classification head. The training dynamics are now included in the revised draft Figure 11 and Figure 12, in the appendix.

Freezing parameters during the early stages of training generally increases sharpness compared to standard training. When parameters are unfrozen, it introduces fluctuations to the training dynamics but rapidly reduces sharpness. These findings suggest that the absolute value of sharpness/Fisher Information (FI) in the early training phase is not a reliable indicator of generalization when accounting for trainable parameters. (Note: This does not contradict findings on sharpness at convergence, as our focus is solely on the early training phase.) Instead, changes in training dynamics (sharpness/FI) is more likely to be signals of a “phase change” [2] in the learning process rather than directly predicting OOD generalization when factoring in the trainable parameters (hence, leads to our experiments with the heuristic algorithm).

Q3: To the best of our knowledge, no prior work has specifically focused on this aspect of OOD behaviour. In our Appendix (F.3), we also found when keeping the trainable parameters constant, by applying a regularizer that reduces the FI [3] also has a similar effect. While we have not yet experimented with other conventional OOD methods, we aim to explore them in the future. As previously mentioned, we observed that the Fisher Penalty also demonstrates similar behaviour, where minimizing sharpness from the beginning does not improve OOD performance.

Reference:

[1] Kumar, Ananya et al. “Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution.” International Conference on Learning Representations (2022).

[2] Achille, Alessandro et al. “Critical Learning Periods in Deep Networks.” International Conference on Learning Representations (2018).

[3] Jastrzebski, Stanislaw et al. “Catastrophic Fisher Explosion: Early Phase Fisher Matrix Impacts Generalization.” International Conference on Machine Learning (2020).

We would like to thank reviewers for their time and feedback.

W1: We recognize the importance of providing clearer intuition on why gradual unfreezing supports OOD generalization. In our experiments, we explored three different settings: one involving training from scratch and two involving fine-tuning with pre-learned features. We agree that the two fine-tuning settings are comparable to the approach of Kumar et al. [1]. Restricting trainable parameters acts as a form of regularization, promoting better gradient alignment and increased robustness to spurious correlations. We think this is the underlying mechanism at play when training from scratch (section 5.2).

W2: We appreciate your insightful suggestion on the additional work on spurious correlations. Our empirical evidence seems to suggest this could be related, especially when existing covariate shifts evaluation dataset don’t control for spurious correlations. We have updated our related work section (L108) to include the suggested recent work on early training and spurious correlation.

W3: We acknowledge the limitation in generalizing this heuristic and intend to expand our study in future work to better validate robustness across real-world scenarios. Our goal is to use heuristics to validate our hypothesis, and we believe our experiments achieve this objective. We really appreciate the reviewer's suggestion and include related discussion in our limitations section.

W4: Ablation Study on Timing of Gradual Unfreezing Thank you for this insightful suggestion. We performed additional experiments where gradual unfreezing was delayed to mid- and late-training phases. We first allow full parameter training, then during the mid- or late- training, freeze all the parameters, and perform unfreezing (standard training → freezing all parameters → gradual unfreezing). Here we report the results using MNIST when training from scratch (using the same training hyperparameters specified in our current draft appendix, hyperparameter search 3 values of the unfreezing interval k, results averaged over 6 random seeds).

The best ID/OOD results are as of following:

• Freezing-Unfreezing starts at 500 steps, best results: 99.18 / 32.29
• Freezing-Unfreezing starts at 1000 steps, best results: 99.22 / 30.71
• Freezing-Unfreezing starts at 2000 steps, best results: 99.26 / 39.12

For comparison, the early period unfreezing gives 98.98 / 63.99.

Q1/Q2 - Clarifications on Figure 5 The values of k adapted in Figure 5 are randomly selected to help with visualisation. Regarding the “rapid change followed by stabilization phase,” our observations indicate that, depending on the dataset and initialization, metrics like tr(F) and sharpness can stabilise at varying times. This is evident in the differences seen across Figure 5 subplots, emphasising that heuristic adjustments may be necessary based on training dynamics.

Q3: We appreciate the suggestion to replace "removal of interventions" with "gradual unfreezing" for clarity. We revised this to ensure the terminology is consistent and accessible for first-time readers.

We thank the reviewer for their time and suggestions.

We would like to clarify a misunderstanding. As stated in our original submission, "only ID data is used for model selection (L199)," which aligns with the reviewer’s suggestion. All training dynamic metrics were calculated using training data, as described in L952.

We appreciate the reviewer’s comments on the necessity of unfreezing layers sequentially versus randomly. To clarify, top-down unfreezing is not proposed as a novel method but rather as an analytical tool for our empirical study (we add discussion to this in Sec 3.3 in the revision). Prior work [1-5] has demonstrated that top-down unfreezing performs best in practical OOD generalization scenarios. Specifically, [2] and [3] compared top-down unfreezing with reverse schedules (bottom-up) and alternative selection strategies (e.g., random or Fisher-based), finding that these approaches performed worse or comparably. Based on this evidence, we chose to focus on the top-down schedule for this study.

Additionally, we present results for MNIST training from scratch using the same setup, where the best-performing bottom-up averages for ID/OOD across six runs are 99.17 / 44.73, respectively. For reference, the best-performing top-down performances are 98.98 / 63.99.

We have made a significant effort to study a diverse range of models, architectures, and domains, including ResNet18, VGG11, SimpleCNN, pretrained Vision Transformers, and XLM-RoBERTa, across both training from scratch and fine-tuning settings aligning with the models used in prior studies. Despite limited academic compute resources, we conducted a large number of experiments to produce our empirical results, over 300 experiments for Figure 2 alone. We believe our current results provide robust and sufficient evidence to support our conclusions.

We emphasize that the observed correlation between OOD performance changes and the timing of intervention removal is both novel and important. Additionally, our training dynamic analysis uncovers a compelling pattern: early metrics such as sharpness or Fisher Information (FI) do not reliably indicate OOD generalization. This is particularly noteworthy, as prior studies often assume trainable parameters are constant or that all parameters are trainable, which ignores the efficient tuning approaches in current transfer learning settings.

Finally, we would like to emphasize that the primary goal of this work is to demonstrate the existence of a critical period in OOD generalization and analyze its impact. The focus is not on proposing a novel method to improve OOD generalization but rather on gaining an empirical understanding of this critical period. Therefore, the magnitude of performance improvement is not the central concern of this study.

We hope these clarifications address your concerns and further highlight the contributions of our study. Thank you for your constructive feedback.

References:

[1] Kumar et al. “Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution.” International Conference on Learning Representations (2022).

[2] Lee et al. “Surgical Fine-Tuning Improves Adaptation to Distribution Shifts.” International Conference on Learning Representations (2023).

[3] Liu et al.. “FUN with Fisher: Improving Generalization of Adapter-Based Cross-lingual Transfer with Scheduled Unfreezing.” North American Chapter of the Association for Computational Linguistics (2023).

[4] Howard and Ruder. “Universal Language Model Fine-tuning for Text Classification.” Annual Meeting of the Association for Computational Linguistics (2018).

[5] Reinhardt et al. “Improving Vision-Language Cross-Lingual Transfer with Scheduled Unfreezing.” Proceedings of the 3rd Workshop on Advances in Language and Vision Research (ALVR) (2024).

Thank you very much for the discussion.

• We will include an error bar to all our figures in a future revision.

• Checkpoint selection & Figure 2: To clarify, for a given K value, a checkpoint is selected based on ID results, and the corresponding OOD results are plotted for analysis. This approach demonstrates the early training effect shown in Figure 2, as evaluating on OOD data is essential to observe this behaviour without an alternative. We are happy to revise the "maximum OOD accuracy improvement" to “maximum possible OOD accuracy improvement” to emphasize on this in a revision.

• Figure 2 vs Table 1: As the reviewer suggested, the best possible improvements in Figure 2 are achieved by looking at OOD data. In contrast, in Table 1 the K is determined by the heuristic algorithm purely on training data learning dynamics only. This shows the potential of achieving such maximum possible OOD improvements in a realistic setting (i.e. no access to OOD data during hyperparameter selection).

• Thank you very much for the additional suggestions on writing and pulling limitation sections into the main paper. We will incorporate the suggestions in a revision.

• Fig. 5: Only GU is used in Figure 5, with different K values. Std is the vanilla standard training (no interventions).

To answer your additional questions:

• E.g. Figure 5 (a): around the first 100 steps. For this figure, the learning dynamics are recorded every 5 training steps.
• Is kL < N: To achieve better OOD results, kL needs to be smaller than N. Prior work typically uses kL = N (e.g., in [1-3]).
• Appendix B mentioned three models, but Table 3 mentions 4: we will fix this inconsistency.
• DomainNet requires longer time to converge.
• the ResNet18 is not really unfrozen with a large parameter k? The largest value k in Figure 2 or Figure 3 is always small or equals to N (L239 in our paper) to ensure all blocks were unfrozen. To clarify, block is determined based on the implementation namespace (L941 in our paper for ResNet). For example, the last unfreezing k in Figure 2 for MNIST, is at 800.
• In Figure 2, you actually show the unfreezing schedule parameter K: yes, as indicated in our figure axis, captions and paragraph 1 in S5.1.

Thank you for your additional suggestions on writing. We hope these clarifications address your concerns and further highlight the contributions of our study.

The best kL is always substantially smaller than N, i.e., in the early period of training. All training hyperparameters (epochs or steps) are specified in Table 3. To translate the epochs into steps, here are the numbers:

MNIST - optimal k around 200, training steps 4.6K

CIFAR10 - optimal k around 3000, training steps 78K

CIFAR100 - optimal k around 1500, training steps 78K

OfficeHome - optimal k around 50, training steps 5K

DomainNet - optimal k around 200, training steps 15K

XNLI - optimal k around 1000, training steps 184K

SQUAD - optimal k around 1500, training steps 49K

1] Kumar et al. “Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution.” International Conference on Learning Representations (2022).

[2] Howard and Ruder. “Universal Language Model Fine-tuning for Text Classification.” Annual Meeting of the Association for Computational Linguistics (2018).

[3] Raffel, Colin et al. “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.” J. Mach. Learn. Res. 21 (2019): 140:1-140:67.

We would like to thank reviewers for their time and feedback.

W1: We appreciate the question regarding the rationale for starting gradual unfreezing (GU) at the end of a sharp change in sharpness. The decision to explore sharpness as a timing indicator is motivated by empirical observations of when GU begins to yield positive OOD results. While we do not claim that sharpness alone predicts OOD generalization, our study highlights how sharpness changes can help us identify a phase in early training where sensitivity to mini-batch gradients diminishes, thus potentially stabilizing training. This empirical finding suggests that the moment of sharp change in sharpness could serve as a practical indicator, even though sharpness itself is not a direct measure of OOD performance during training.

W2: We’d like to clarify that our heuristic intervention algorithm is also for validating the hypothesis. To prevent confusion, we now revised the section of our paper related to the heuristic algorithm intervention. We updated the original “case study” section into “hypothesis validation” (L456) to be explicit.

W3: We will revise these tables to include the absolute gain.

Q1: Extra computation is required to track the training dynamics (the worst-case sharpness is also more costly) at least in the early period of training. We use the heuristics to arrive at the insights that change in sharpness and training dynamics as a signal to OOD generalisation, rather than prescribing it as a method. We have updated our “limitations” section (L918, appendix A) to discuss this.

We hope these clarifications address your concerns and further highlight the contributions of our study.

Thank you for pointing this out. We apologize for the confusion. It was a typo on our part, we intended to refer to S6.1 (instead of S5.1). In this section, we discuss the training dynamics, including sharpness.

We would like to thank reviewers for their time and feedback.

W1: We demonstrate empirically that when training with partial or dynamically varying parameters, sharpness in the early period of training does not predict OOD generalization. We believe this finding is interesting and important because it challenges the commonly held belief that lower sharpness throughout the entire training process is crucial for generalization. As methods involving partial or dynamic parameter training gain popularity, our findings highlight an important aspect that warrants further attention. Additional supporting cases are also detailed in our Appendix.

To the best of our knowledge, the new insights provided in this work are:

• The early stages of training significantly influence OOD performance under covariate shift.
• The number of trainable parameters is a critical factor overlooked in previous studies.
• Empirical evidence shows that low sharpness or Fisher Information during the early training phase does not predict OOD generalization.
• Relative changes in training dynamics during the early period are more informative for OOD generalization than absolute values.

W2: We would like to clarify that our study does not propose a new algorithm, we aim to provide new insights rather than to establish SOTA OOD performance. We use gradual unfreezing as a method for investigation (L20 in the abstract of our original submission). It is selected due to prior work showing its effectiveness in achieving SOTA results in various transfer learning settings [1-5]. We reflect this in our updated revision (Section 3.3).

We also revised the section of our paper related to the heuristic algorithm intervention. We updated the original “case study” section into “hypothesis validation” (L456) to prevent misunderstanding. Additionally, we updated the corresponding “summary of findings” (L508) to emphasize that our intervention algorithm is for confirming insights and showing potential applications rather than rather than a prescriptive method.

W3: We are happy to report preliminary results on Domain Net using the same hyperparameters specified in the paper. We included these results in the revised version of the paper (L506 and Table 7).

We hope these clarifications address your concerns and emphasize our study's key contribution: demonstrating the significance of the early training phase for OOD generalization, particularly when considering trainable parameters and methods like freeze training.

References:

[1] Kumar et al. “Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution.” International Conference on Learning Representations (2022).

[2] Lee et al. “Surgical Fine-Tuning Improves Adaptation to Distribution Shifts.” International Conference on Learning Representations (2023).

[3] Liu et al.. “FUN with Fisher: Improving Generalization of Adapter-Based Cross-lingual Transfer with Scheduled Unfreezing.” North American Chapter of the Association for Computational Linguistics (2023).

[4] Howard and Ruder. “Universal Language Model Fine-tuning for Text Classification.” Annual Meeting of the Association for Computational Linguistics (2018).

[5] Reinhardt et al. “Improving Vision-Language Cross-Lingual Transfer with Scheduled Unfreezing.” Proceedings of the 3rd Workshop on Advances in Language and Vision Research (ALVR) (2024).