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

## Abstract

Abstract Prior research has found that differences in the early period of neural network training significantly impact the performance of in-distribution (ID) tasks. However, neural networks are often sensitive to out-of-distribution (OOD) data, making them less reliable in downstream applications. Yet, the impact of the early training period on OOD generalization remains understudied due to its complexity and lack of effective analytical methodologies. In this work, we investigate the relationship between learning dynamics and OOD generalization during the early period of neural network training. We utilize the trace of Fisher Information and sharpness, with a focus on gradual unfreezing (i.e. progressively unfreezing parameters during training) as the methodology for investigation. Through a series of empirical experiments, we show that 1) selecting the number of trainable parameters at different times during training, i.e. realized by gradual unfreezing – has a minuscule impact on ID results, but greatly affects the generalization to OOD data; 2) the absolute values of sharpness and trace of Fisher Information at the initial period of training are not indicative for OOD generalization, but the relative values could be; 3) the trace of Fisher Information and sharpness may be used as indicators for the removal of interventions during early period of training for better OOD generalization.

## 1 Introduction

While deep neural networks have achieved impressive results in their training tasks, they are often sensitive to distribution shifts (i.e. out-of-distribution, OOD) during inference. As in many applications of deep neural networks, the training and testing data may come from different distributions, such as training with perfect images or texts but inference with noise-corrupted data [, ], data obtained from different time periods [, ], or across different languages and domains [, , , , ], to name a few. Failure to generalize to the OOD setting degrades the models’ robustness and reliability.

A plethora of research has found that differences in the early period of training have a significant impact on the in-distribution (ID) performance [, , , , inter alia] for a wide range of settings. These settings include classification, multimodal learning, training from scratch, parameter-efficient fine-tuning or even federated learning. The wide observation of such a period in machine learning applications suggests that the early period of learning is generally important for neural network training [], which is sometimes analogous to biological processes (such as the critical learning period in animals [, ]).

In particular, prior literature identifies that modifications (or interventions) to the optimization process are critical ways to shape the early period of training. Training techniques such as weight decay [], learning rates  [, ], data augmentations [, ] (such as with MixUp []) or adding noise to weights [] impact learning dynamics early on, and can significantly improve or hamper the final task results depending on the time of application or removal. Yet, there is very limited work exploring how the early period of training impacts generalization to OOD data during testing, as well as “when” the early period of training significantly affects the outcome of final (OOD) results.

[] found that using gradual unfreezing [] (i.e., progressively releasing trainable parameters during training at fixed time intervals) can impact the trace of Fisher Information in the early period of training; however, the work was focused on which parameters to select for better cross-lingual transfer (where cross-lingual transfer is a form of natural OOD generalization). Besides the trace of Fisher Information, other sharpness metrics have also been used to study the generalization of network training, especially after the success of methods such as Sharpness-Aware Minimization (SAM) [, , ], however, sharpness is less used in prior work to study the early period of training.

In this work, we delve deeper and advance our understanding of how the early period of training impacts OOD generalization through a series of empirical investigations. We first employ gradual unfreezing [] as a method to intervene in the dynamics of the early period of training from scratch for OOD generalization and investigate how the early period of training impacts OOD generalization. This is done in a simple and controlled setting. Next, we investigate which metrics are effective for studying the early period of training for OOD generalization. Here, we utilize Fisher Information and sharpness and show how these metrics change in the early period of training with gradual unfreezing, and examine their impact on the final solutions.
After this, we focus on whether there is an optimal time (i.e. “when”) to remove intervention to achieve a good trade-off between ID and OOD generalization. Along the way, we also try to answer the questions: how early does the “early period” begin and can our observations extend to other settings without gradual unfreezing or with different model architectures?

To summarize, in this paper, we show that 1) when considering the number of trainable parameters at a time, i.e. realized by gradual unfreezing – has a minuscule impact on ID results, but greatly affects the generalization on OOD data; 2) we also show that the value of sharpness and trace of Fisher Information during the initial period of training are not indicative for OOD generalization, but their relative values could be; 3) the trace of Fisher Information and sharpness may signify the removal of interventions during the early period of training.

## 2 Related Work

Early period of neural network training. Under the standard usage of the term generalization (in-distribution, where training and testing data are assumed to be from the same distribution), prior work [] shows that the early period of training of neural networks exhibits a “critical learning period” when trained from scratch. Regularization and interventions applied in this critical period affect final task results.

[] indicates that when learning with a lower learning rate, Fisher Information exhibits an “explosion” which impedes generalization. Applying regularization to the trace of Fisher Information alleviates the negative impact of the high Fisher Information. [] shows the termination of MixUp early in training and switching to standard empirical risk minimization helps with better ID generalization. [, ] shows that even winning “lottery tickets” emerge in the early period of training with large learning rates. The critical learning period is also found in many other settings, such as in multimodal models [], in linear models [], in transformers [] and Federated learning []. However, prior work focuses on in-distribution generalization, ignoring the setting of OOD generalization.

[] shows the early period of training relates to cross-lingual transfer performance in the parameter-efficient fine-tuning (PEFT) setting with transformer models, connecting the early period of training with progressive parameter unfreezing and OOD generalization (as cross-lingual data are naturally-occurring OOD data). We utilize findings in [] and base our work on gradual unfreezing []. In this paper, we focus on 1) a general setup (i.e. training from scratch), and 2) the characterization of the early period of training and its relationship to OOD generalization.

Fisher Information, sharpness and generalization.
Fisher Information has been studied in many prior works such as [, ] to investigate and improve optimization behaviour. Similarly, sharpness is another popular metric used to study optimization behaviour and its relationship to generalization.

[] found a correlation between sharpness and the ratio of learning rate to batch size. [, , ] give theoretical guarantees on the generalization error using sharpness-related measures and conduct large empirical experiments showing that sharpness-based measures correlate with generalization. However, there have been debates on whether sharp minima (such as a high largest eigenvalue of the training Hessian, $\lambda_{max}$) imply poor generalization [] and demonstrate the limits of $\lambda_{max}$ in explaining in-distribution generalization []. Most of the current research efforts are towards analyzing the loss landscape at convergence, for in-distribution generalization.

For OOD, [] shows adaptive sharpness is not a stable measurement for OOD generalization of the final solution. However, the relationship between metrics such as Fisher Information, sharpness and OOD generalization in the early period of training is unclear.

## 3 Preliminaries

### 3.1 Fisher Information Matrix (FIM)

To investigate the training process, we first look at the Fisher Information [].
Fisher Information reflects the local curvature and measures the amount of information with respect to network parameters, i.e. how sensitive the network predictions are to the changes in parameters. A larger Fisher Information indicates that a small change in network parameters can change the output significantly, which can be interpreted as a “sharper” loss landscape.

Let $x$ be the inputs and $y$ be the labels of a dataset $D$. Given a neural network that is parameterized by $w$. The Fisher Information is defined as:

$$ $\displaystyle F(w)=\operatorname{\mathbb{E}}_{P_{x,y}}\left[\nabla_{w}\log p_{ w}(\hat{y}|x)\nabla_{w}\log p_{w}(\hat{y}|x)^{T}\right].$ (1) $$

Estimating the full Fisher Information is usually expensive, prior work shows that the
trace of the Fisher Information ($\operatorname{\texttt{tr}(F)}$) correlates well with the full Fisher Information and can be used for real applications and capture signals during the learning process [, , , inter alia]. Using the empirical data distribution $\hat{Q}(x)$:

$$ $\displaystyle\operatorname{\texttt{tr}(F)}=\operatorname{\mathbb{E}}_{x\sim \hat{Q}(x)}\operatorname{\mathbb{E}}_{\hat{y}\sim p_{w}(\hat{y}|x)}||\nabla_{w }\log p_{w}(\hat{y}|x)||^{2}.$ (2) $$

### 3.2 Sharpness

Let $\mathcal{L_{D}}(w)=\frac{1}{|D|}\sum_{(x,y)\in D}\log p_{w}(y|x)$ be the loss over training datasets of a neural network parameterized by $w$, and $\delta$ be a small perturbation drawn from a noise distribution, such as a Gaussian distribution $\mathcal{N}(0,\rho^{2}diag(c^{2}))$. The definition of worst-case and average-case sharpness are [, , , ]:

$$ $\displaystyle S^{\rho}_{avg}=\operatorname{\mathbb{E}}_{\delta\sim\mathcal{N}( 0,\rho^{2}diag(c^{2}))}\mathcal{L_{D}}(w-\delta)-\mathcal{L_{D}}(w),$ (3) $$

$$ $\displaystyle S^{\rho}_{worst}=\max_{\|\delta\odot c^{-1}\|_{p}\;\leq\;\rho} \mathcal{L_{D}}(w-\delta)-\mathcal{L_{D}}(w),$ (4) $$

where $\odot c^{-1}$ is element-wise multiplication.

The sharpness here refers to how rapidly the loss changes with respect to the changes in the model parameters.(^1^11The sharpness can be negative.) While both the Fisher Information and sharpness are used for investigating loss landscapes and generalization, they offer different views of the training process.

Both $S^{\rho}_{avg}$ and $S^{\rho}_{worst}$ are studied in the prior literature for generalization. For instance, [, , ] show that the worst-case sharpness correlates better with generalization, at convergence. While prior work believe that flatter (less sharp) minima in the loss landscape lead to better generalization in neural networks [, , , ], these metrics’ attribution to the early period of training and how their early period trends are related to OOD generalization is understudied.

### 3.3 Gradual Unfreezing

Gradual unfreezing [] is a simple tuning method that progressively increases the number of trainable parameters (i.e. unfreeze, layer-by-layer) of the neural network from the top to the bottom of the network with a fixed interval of training steps, $k$ (i.e. the unfreezing time). In this paper, we used a modified formulation of gradual unfreezing [], where we progressively unfreeze the parameters during the early period in the training top-down and for “blocks” of parameters (a block of parameters can be a single layer or several consecutive layers, in our case, we use the namespace of the parameters used in standard implementations to determine blocks). See Appendix A for details and the algorithm.

## 4 Experimental Setup

To investigate our research questions, we perform two sets of experiments. First in a controlled setup, with datasets such as CIFAR10 [] and models such as ResNet [] to observe patterns and validate our hypothesis. Then, we experiment with the complex setup with transformer models. We use $\rho=0.01$ to calculate the sharpness (both average-case and worst-case) with 15 samples, and $L2$ norm for the worst-case sharpness. We normalize the $\operatorname{\texttt{tr}(F)}$ by the number of trainable parameters. We use the Auto-PGD algorithm [] as implemented in [] (we refer the readers to the original papers for details) for computing worst-case sharpness, as it is a hyperparameter-free estimation method.

#### Controlled Setup.

In this setup, following [], we use ResNet-18 and VGG-11 as the neural networks for this study (training from scratch). We perform experiments on classic image classification datasets and use MNIST [], CIFAR10 [] and CIFAR100 [] for training. For out-of-distribution evaluation, we use the corrupted corresponding evaluation datasets, named MNIST-C [], CIFAR10-C [] and CIFAR-100-C [] (results averaged across corruptions and severities). The ID evaluation sets are the original test sets respectively.

We use the SGD optimizer and the CIFAR datasets are applied with standard augmentations (i.e. random crops and flips). We report results over 6 random seeds for MNIST (due to the high variance in OOD results), and we use 4 random seeds for other datasets. Other hyperparameters such as learning rate or weight decay are specified in Appendix B.
In our experiments, the default learning rate specified in Appendix B is denoted as $lr_{d}$ and we also experimented with reduced learning rates which are 1/10th of the default, specified as 0.1*$lr_{d}$.

#### Complex Setup.

We also conduct experiments using transformers. As pre-train then fine-tuning becomes a popular way to adapt general foundational models to downstream tasks, we examine the cross-lingual transfer (train with English data, test with other languages) task using parameter-efficient fine-tuning (PEFT) with the LoRA [] adapters (HuggingFace PEFT [] implementation). This setting is parallel to our controlled setting because: 1) only English data (i.e. ID data) is used for training and validation (other language data are for testing), this is a natural setting of OOD generalization with parallel evaluation protocol to the image tasks; 2) using LoRA adapters allows us to inject randomly initialized parameters for learning, which is analogous to our controlled setting.

We train with SQuAD [] (English, question and answering task) and MNLI [] (English, natural language inference task), and evaluate on XQuAD [], MLQA [] and XNLI []. We use XLM-RoBERTa [] as the pre-trained multilingual transformer backbones and AdamW [] as the optimizer. We report results across 4 seeds for all experiments in this setting. Please see Appendix B for hyperparameters.

## 5 Gradual Unfreezing in the Early Period of Training Can Improve Out-of-Distribution Generalization

Recall from § 2, where gradual unfreezing improved OOD performance for PEFT with transformers. Here, we first validate that gradual unfreezing applied to the early period of training in our controlled setting (training from scratch) could also help OOD generalization. By examining three different datasets and two model architectures in Table 1, progressive parameter unfreezing (i.e. gradual unfreezing) does not influence ID results by a large margin (mostly minor degradation, but can also positively impact the ID results). However, gradual unfreezing has a non-negligible positive impact on the OOD results. This observation is also applicable to different learning rates, although the default (larger) learning rate empirically is better for both ID and OOD for CIFAR datasets. Here, we provide evidence that gradual unfreezing can improve OOD performance when training from scratch even if it was proposed for transfer learning [], and validate the usability of gradual unfreezing as an intervention for our study.

**Table 1: Classification results on various datasets and model architectures. RN18 is ResNet-18. The default learning rate ($lr_{d}$) cases are the same as described in §4, 0.1*$lr_{d}$ indicates the learning rate is 1/10th of the default learning rate. GU: indicates that gradual unfreezing is applied to the early period of training, here we observe the results that are close to the ID results, but with better OOD results.**
|  | MNIST RN18 | CIFAR10 RN18 | CIFAR100 RN18 | CIFAR10 VGG11 |
| --- | --- | --- | --- | --- |
| Method | ID / OOD | ID / OOD | ID / OOD | ID / OOD |
| $lr_{d}$ | 99.06/33.36 | 93.32/72.36 | 71.07/45.10 | 88.62/71.63 |
| $lr_{d}$ + GU | 98.98/63.99 | 93.26/72.95 | 71.03/46.34 | 88.53/72.16 |
| 0.1*$lr_{d}$ | 99.26/58.46 | 91.66/71.14 | 69.95/44.59 | 86.93/68.41 |
| 0.1*$lr_{d}$ + GU | 99.18/62.51 | 91.51/71.26 | 70.67/46.03 | 87.01/69.45 |
|  |  |  |  |  |

## 6 Early Period of Training Impacts Out-of-Distribution Generalization

### 6.1 Evidence of Impact on Out-of-Distribution Generalization

[, ] among others show that for ID generalization, there is a “critical learning period” of the neural network. ID generalization degrades with the delaying application of regularization, as well as the removal of regularization.

Using gradual unfreezing, we experimented with different unfreezing steps $k$ (ranging from 1 to equally dividing the total training steps among the number of trainable parameter blocks) to measure its impact on both ID and OOD test results. Indeed (as in Figure 1), it is possible that withholding trainable parameters can influence the OOD generalization as early as after training on a single batch of data. The effect is especially prominent for simpler datasets like MNIST.

Prolonging the unfreezing interval of parameters during training initially results in minimal change in the ID test performance with a larger learning rate, subsequently leading to quick deterioration of the ID results. The deterioration of ID results over unfreezing intervals aligns with trends observed in the early stages of training using other interventions, as reported in prior work [, ], while the effects on OOD results serve as evidence that the early period of training can impact OOD generalization.

Gradual unfreezing casts interesting trends on the OOD generalization and shows the trade-off between ID and OOD generalizations. Before the quick deterioration of ID results, there is a short window where OOD results could be improved. As soon as the ID results start decreasing rapidly, the OOD results first increase, then rapidly decrease for CIFAR10/100 with ResNet, but not in MNIST. There seems to be a range of $k$ in the early period of training, such that the ID results decayed minimally, but with better OOD results, and we will examine this in more detail in § 6.4.

Figure: Figure 1: Change in ID and OOD evaluation results when unfreezing parameters at different times (compared to standard training). Sub-figures (a)-(d) are with the $lr_{d}$, and Sub-figures (e)-(h) are with 0.1*$lr_{d}$. The shortest unfreezing time $k$ is 1 (training with 1 batch of data). The x-axis is in the log scale.

### 6.2 Learning Dynamics in the Early Period of Training

Observing Figure 2, by freezing the number of trainable parameters at a time (and gradually unfreezing them), we can induce higher Fisher Information and larger $S^{\rho}_{avg}$, $S^{\rho}_{worst}$ at the beginning of training compared to the standard training procedure, although there is an anomaly in the $\operatorname{\texttt{tr}(F)}$ of CIFAR100 with ResNet. In general, the longer we withhold parameters, the higher the level of sharpness and $\operatorname{\texttt{tr}(F)}$ we can sustain, unfreezing parameters reduce these metrics.

While there are variations between $S^{\rho}_{avg}$, $S^{\rho}_{worst}$ and $\operatorname{\texttt{tr}(F)}$ they are all sensitive to the early period of training and interventions. $S^{\rho}_{avg}$ shows more consistent trends across datasets and architectures compared to the other two metrics. Due to the randomness in estimating $S^{\rho}_{avg}$, $S^{\rho}_{worst}$, and $\operatorname{\texttt{tr}(F)}$, it is also evident that a single, absolute largest value of these metrics during the early period of training may not be a consistent indication of OOD generalization (or ID generalization, in fact). This indicates that the discussion for a high or low value of $S^{\rho}_{avg}$, $S^{\rho}_{worst}$, or $\operatorname{\texttt{tr}(F)}$ during the early period of training should be relative rather than absolute.

Empirically, our findings differed from prior work on ID generalization (such as []) that demonstrates the ‘explosion’ of $\operatorname{\texttt{tr}(F)}$ during the early period of training (due to using a small learning rate) is harmful. Here, a higher $\operatorname{\texttt{tr}(F)}$ induced by parameter freezing does not hurt generalization, in both ID and OOD. When considering the trainable parameters as a variable, having initial higher sharpness or $\operatorname{\texttt{tr}(F)}$ can be advantageous up to a certain time frame during training. More interestingly, some sub-figures (such as Figure 2 (f) and (j)) show a rapid increase of sharpness after about 100 training steps (see the $k$=750 curve). The $\operatorname{\texttt{tr}(F)}$ and sharpness transition from a transient phase to a relatively stabilized value when withholding parameters for long. Introducing new trainable parameters induces a quick drop in the corresponding metrics. As a result, the optimization trajectory and learning mechanism change when manipulating the trainable parameters through unfreezing during the early period of training (such as the network resorting to higher rank features at the initial learning period, more details in Appendix E.1).

Overall, these results suggest that while a lower sharpness or $\operatorname{\texttt{tr}(F)}$ during the early period of training may be good for ID generalization, when factoring in the trainable parameters (as also shown empirically in Table 1), a lower initial sharpness or $\operatorname{\texttt{tr}(F)}$ could lead to worse OOD generalization. This observation applies strictly to the very early period of training, and the eventual reduction of sharpness or $\operatorname{\texttt{tr}(F)}$ after the initial period is still desirable (evident in Figure 1, Figure 2, and work like SAM [, , ]). Importantly, while sharpness and $\operatorname{\texttt{tr}(F)}$ are effective metrics to study the early period of training, we need to look at their relative trends for OOD generalization (in fact, also in ID).

Figure: Figure 2: Unfreezing parameters at different times affects the learning dynamics in the early period of training (with $lr_{d}$). Sub-figures (a)-(d): The $\operatorname{\texttt{tr}(F)}$ when unfreezing parameters at k= [250,750] steps versus standard training. Sub-figures (e)-(h): The average-case $L2$ sharpness (Eqn. 3) when unfreezing parameters at k=[250,750] steps versus standard training. Sub-figures (i)-(l): The worst-case $L2$ sharpness (Eqn. 4) when unfreezing parameters at k=[250,750] steps for ResNet and k=[250,5000] for VGG, versus standard training. The y-axis is in the log scale and is normalized between 0 and 1000 for visualization. All figures are with the default learning rates to the first 2000 training steps, we also observe similar trends with 0.1*$lr_{d}$ (see Appendix E.4).

### 6.3 Final Solutions and Out-of-Distribution Results

Changing the learning dynamics in the early period of training inevitably results in different final solutions. For example with CIFAR10 on ResNet, the largest eigenvalue of the Hessian ($\lambda_{max}$, calculated on a subset of training data) of final solutions shows a negative correlation with OOD results (see Appendix C). However, such a negative correlation is not consistent nor always statistically significant across different setups. Our results complement the findings in [], which serve as additional evidence of the need for developing new robust metrics and further thorough investigation for OOD generalization.

### 6.4 Learning Dynamics Could Signify the Time Period to Remove Interventions

While the value of $\operatorname{\texttt{tr}(F)}$ or sharpness during the initial period of learning (or at the end of learning) may not be indicative of good OOD generalization in general, they could signify the time to release parameters for minimal ID decay and better OOD results. When unfreezing parameters (Figure 2), we observed that in many cases, those metrics experience a “transient” period (the first 50-100 steps, characterized by rapid growth or drop of the sharpness or $\operatorname{\texttt{tr}(F)}$), followed by a ‘stabilization’ phase where the rate of change in metric values slows down (unless parameters are released).

Notice that the best range of $k$ for overall best ID and OOD results is after the stabilization of sharpness and $\operatorname{\texttt{tr}(F)}$ (in Figure 1 and Figure 2), but not for too long. For instance in Figure 1 (b), although the OOD results improved significantly, the ID results deteriorated drastically after 800-1000 training steps (at least 1 point drop for CIFAR10/100). This observation leads to the conjecture that the best time to remove the intervention (i.e. unfreezing parameters in our case) while keeping reasonable ID results (less than 0.5 points decrease in accuracy) and achieving better OOD results must satisfy two constraints: 1) after the initial rapid change of sharpness or $\operatorname{\texttt{tr}(F)}$ (the transient phase), and 2) not too far into the stabilization phase.

The second constraint is self-evident in Figure 1, as a larger $k$ hurts both ID and OOD results. To quickly validate the first constraint, we use MNIST to pick the earliest ending step of the transient period among three metrics ($S^{\rho}_{worst}$, $S^{\rho}_{avg}$ and $\operatorname{\texttt{tr}(F)}$). Then we experiment with 10 different $k$ values each consecutively (10 steps apart) that are smaller or greater than that $k$ value. For the smaller $k$s, we get 98.93/52.72 as the median ID/OOD results, for the larger $k$s, we get 98.91/53.54 as the median ID/OOD results, which validates the constraints. The stabilization of examined metrics indicates when to introduce new trainable parameters.

Next, we use a heuristic algorithm that satisfies the above-mentioned constraints to determine the stabilization time of the three metrics (we first detect a significant change in metrics, then detect the stabilization point of the metrics, the algorithm is in Appendix D). The OOD results are then compared with with ten random $k$ values per dataset ($k\leq 800$) to determine the winning rate (i.e., the percentage of times where the value picked by the algorithm is better than a sampled value), and the results are in Table 2.

Empirically, using such an algorithm is better than doing a random hyperparameter search the majority of the time. In most cases, the degradation of ID accuracy is within 0.5 points, except when using the VGG network. Nonetheless, this further validates that the stabilization of $S^{\rho}_{worst}$, $S^{\rho}_{avg}$ and $\operatorname{\texttt{tr}(F)}$ could signify the removal of interventions (in our case gradual unfreezing) to trade some ID performance for OOD. While $\operatorname{\texttt{tr}(F)}$ shows better results, there isn’t a clear winning metric for intervention removal due to: 1) metrics exhibiting high noise during training; and 2) the determined stabilization points from different metrics collide or are very close to each other. We will defer the exploration of more sophisticated algorithms for future work. However, it’s worth noting the existence of an optimal time that effectively balances good ID and ODD results.

**Table 2: Results using the heuristic algorithm (Appendix D) to determine the $k$ for gradual unfreezing (GU), best OOD results are bolded. The algorithm can determine the same value of $k$ in different metrics in multiple cases (hence the same results). WR stands for winning rate (OOD). See Appendix D for visualization of $k$ overlay on Figure 2.**
|  | MNIST RN18 | CIFAR10 RN18 | CIFAR100 RN18 | WR | CIFAR10 VGG11 | WR |
| --- | --- | --- | --- | --- | --- | --- |
| Method | ID / OOD | ID / OOD | ID / OOD | - | ID / OOD |  |
| Standard | 99.06/33.36 | 93.32/72.36 | 71.07/45.10 | - | 88.62/71.63 | - |
| GU${}_{S^{\rho}_{worst}}$ | 98.78/52.48 | 93.06/72.75 | 70.68/45.19 | 60% | 87.69/71.47 | 40% |
| GU${}_{S^{\rho}_{avg}}$ | 98.78/52.48 | 93.02/72.58 | 70.67/45.35 | 60% | 87.71/72.37 | 100% |
| GU${}_{\operatorname{\texttt{tr}(F)}}$ | 98.91/54.12 | 93.02/73.56 | 70.78/45.82 | 83% | 88.40/71.86 | 60% |
|  |  |  |  |  |  |  |

## 7 Generality of Findings on the Early Period of Training

### 7.1 Higher Initial Sharpness via Learning Rate

Figure: Figure 3: The $S^{\rho}_{avg}$ profile during the early period of training, the learning rate switch happens at 200 steps (all parameters trainable).
Refer to caption: x1.png

Using gradual unfreezing is a specific case where high sharpness at the initial learning period could benefit OOD generalizations. A critical question is: is there another way to intervene in the early period of training with higher sharpness, that also positively impacts OOD generalization?

Recall that a higher learning rate typically results in lower sharpness (and a lower learning rate results in higher sharpness, as indicated in  []). Based on our findings in the previous sections, we hypothesize that using a lower learning rate at the initial period of learning, then switching to a higher learning rate later (high sharpness to low sharpness), may help OOD generalization (surprisingly, this is exactly a simple form of learning rate warm-up!).

To validate, we use the CIFAR10 dataset on ResNet18 with two learning rates: we initially use 1/10th of the default learning rate, then increase the learning rate to the default value after $k$ steps (i.e. low-to-high, denoted as lr${}_{l2h}$, and the reverse is lr${}_{h2l}$. $k$ is determined using random hyperparameter search). The results are in Table 3 (an example sharpness profile is in Figure 3, with more details in Appendix E.2); indeed, lr${}_{l2h}$ can provide better OOD results (and without degrading ID results). Further, with the learning rate switching step $k$ determined using $\operatorname{\texttt{tr}(F)}$, $S^{\rho}_{avg}$ and $S^{\rho}_{worst}$, the OOD results are 72.57 and 72.94 (same $k$ for both sharpness metrics) respectively.

**Table 3: Effects of switch learning rates (increasing or decreasing) in the early period of training.**
| Method | ID/OOD | Method | ID/OOD |
| --- | --- | --- | --- |
| $lr_{d}$ | 93.32/72.36 | $lr_{l2h}$ | 93.35/72.89 |
| 0.1*$lr_{d}$ | 91.66/71.14 | $lr_{h2l}$ | 91.58/71.18 |

### 7.2 Empirical Validations in Transformers

For our complex experimental setup (described in §4), Figure 4 shows the learning dynamics of XLM-R with MNLI in the early period, withholding trainable parameters increases the sharpness and $\operatorname{\texttt{tr}(F)}$, note that in Figure 4 (c) the $S^{\rho}_{worst}$ value is negative, withholding trainable parameters still increase the $S^{\rho}_{worst}$ during training based on Eqn. 4 (the learning dynamics for the SQuAD dataset is in Appendix E.3).

Figure: Figure 4: Learning Dynamics of XLM-R with LoRA training with MNLI, y-axis for figures are in the log scale, the original value sharpness value in (c) is negative where we take the absolute value before visualization. All values are normalized between 0 and 1000 for visualization.

Similarly, the results using $\operatorname{\texttt{tr}(F)}$ to determine the $k$ for unfreezing is in Table 4 ($k$ values are in Appendix D, results with $S^{\rho}_{worst}$ / $S^{\rho}_{avg}$ are in Table 7 in the Appendix) and the winning rate is 80%. The ID results are not sacrificed in this experimental setting, hence further pointing towards that the stabilization of sharpness and $\operatorname{\texttt{tr}(F)}$ could signify ‘when’ to remove intervention in the early period of training for better OOD generalization.

**Table 4: Cross-lingual transfer results of standard training and using $\operatorname{\texttt{tr}(F)}$ to determine unfreezing interval $k$ for gradual unfreezing (GU). WR stands for winning rate, averaged over 10 randomly sampled $k$ per training dataset. EM is the exact match score.**
|  | XQuAD |  | MLQA |  | XNLI | WR |
| --- | --- | --- | --- | --- | --- | --- |
| Method | F1- En/X-ling | EM- En/X-ling | F1- X-ling | EM- X-ling | Acc- En/X-ling |  |
| Standard | 82.96/68.72 | 71.39/52.64 | 56.27 | 40.93 | 83.17/71.84 | - |
| GU${}_{\operatorname{\texttt{tr}(F)}}$ | 83.77/70.70 | 72.33/54.40 | 58.47 | 42.31 | 83.36/72.49 | 80% |

## 8 Conclusions

In this work, we investigate the early period of training and its impact on out-of-distribution generalization. We demonstrate that using gradual unfreezing to modulate the number of trainable parameters during the early period of training can affect out-of-distribution generalization in various settings, including transformers, and reveal that the number of trainable parameters at a time is an important factor that was missing in the previous literature. We observe different patterns to previous work in in-distribution generalization, where higher sharpness and $\operatorname{\texttt{tr}(F)}$ during the early period of training may be beneficial for out-of-distribution generalization. We reveal that metric values during the early period of training may not be indicative of the out-of-distribution generalization, however, they could signify “when” to remove interventions such as gradual unfreezing (or to increase the learning rate, §7.1) for better results.

The significance of effective training and fine-tuning, along with the growing emphasis on research into techniques that involve modifying only partial parameters of the final model, such as freezing parameters (e.g.
Adapters [, , ] or pruning) or network expansions [, , ] cannot be overstated. Our empirical investigations highlight the need to develop a more theoretical understanding of the early period of training and OOD generalization, as well as the creation of new theoretical metrics for better indication of OOD generalization.

## 9 Acknowledgement

This work was funded by the German Federal Ministry of Education and Research (BMBF) under the promotional reference 13N15897 (MISRIK).

## Appendix A Gradual Unfreezing

Following the notations and algorithm in [], let $\textrm{FORWARD}(*)$ be the standard forward pass, and $\textrm{BACKWARD}(*)$ calculates gradients and performs updates for trainable parameters. The modified gradual unfreezing algorithm is in Algorithm 1.

In our experiments, we partition the blocks by their natural namespaces, as the following:

ResNet18: The definition block follows the standard implementation of ResNet, with an input convolution layer and a batch norm group together as the additional block. The model parameters are partitioned into 5 blocks, and a classification head.

VGG11: The definition block follows the standard implementation of VGG, with 8 blocks in total. The classification head consists of 3 linear layers with a ReLU function in between.

Figure: Algorithm 1 Gradual Unfreezing

XLM-RoBERTa + LoRA: The experiment follows  []. Each parameter block consists of 2 sets of LoRA adapters added to the query and value of the backbone transformer from the same layer. The LoRA parameters are partitioned into 12 blocks, and a classification head, where the classification head and the last layer of LoRA adapters are trainable initially.

## Appendix B Hyperparameters

For all our experiments, the hyperparameters are listed in Table 5. We use the default hyperparameters for the AdamW optimizer, except for the learning rate. All other hyperparameters for the transformer experiments follow [], and we use the HuggingFace PEFT [] implementations of LoRA.

For calculating $S^{\rho}_{worst}$ and $S^{\rho}_{avg}$, we use $L2$ norm and $\rho=0.01$ with 15 examples. We follow the setup in [] and use the implementation with 2048 data points from the training data (un-augmented when calculating sharpness metrics) for all experiments. We use a batch size of 256, except for SQuAD (the batch size is 32) for calculating all the metrics. The sharpness and $\operatorname{\texttt{tr}(F)}$ are recorded every 10 batches (steps) for all datasets.

**Table 5: Hyperparameters used in all experiments.**
|  | MNIST | CIFAR10 | CIFAR10 | CIFAR100 | SQuAD | MNLI |
| --- | --- | --- | --- | --- | --- | --- |
|  | RN18 | RN18 | VGG11 | RN18 | XLM-R | XLM-R |
| optimizer | SGD | SGD | SGD | SGD | AdamW | AdamW |
| lr scheduler | const. | const. | const. | const. | linear | linear |
| $lr_{d}$ | 0.01 | 0.1 | 0.15 | 0.01 | 0.0005 | 0.0005 |
| batch size | 128 | 128 | 128 | 128 | 32 | 128 |
| training epochs | 10 | 200 | 200 | 200 | 15 | 15 |
| weight decay | 0.01 | 0 | 0 | 0.0005 | 0.01 | 0.01 |
| momentum | 0.9 | 0 | 0 | 0.9 | - | - |
| LoRA r | - | - | - | - | 8 | 8 |
| LoRA alpha | - | - | - | - | 8 | 8 |
| LoRA dropout | - | - | - | - | 0.2 | 0.2 |

## Appendix C Properties of the Final Solutions and OOD Results

We plot the final solution’s $\lambda_{max}$ (largest eigenvalue of training data feature), $S^{\rho}_{worst}$ and $S^{\rho}_{avg}$ against the OOD test results in Figure 5 respectively.

While in general the sharpness measures and OOD have negative correlations (i.e. the smaller sharpness values the better, especially $S^{\rho}_{worst}$ has a consistent negative correlation), they are not always statistically significant (e.g., for MNIST). The learning rate has a big impact on the final solutions’ sharpness. Furthermore, such as in Figure 5 (c), we can even attain slightly positive correlations. Our results complement the findings in [], which serve as evidence pointing towards the need for developing robust new metrics and thorough investigation for OOD generalization.

Figure: Figure 5: Final feature $\lambda_{max}$, $S^{\rho}_{worst}$, and $S^{\rho}_{avg}$ versus the OOD test results (coloured by learning rate), labelled with Kendall’s $\tau$ and p-value.

## Appendix D Algorithm to Determine the Unfreeze Time

To verify our hypothesis, we use a simple algorithm with heuristic to determine the unfreezing time $k$, Algorithm 2 presents the flow, $\tau$ is 3 or 8 and $\epsilon$ is 0.02 (i.e. the percentage of change in the signal is within 2%). The algorithm takes $t_{\Delta_{\hat{S}}}$ as the input, which is the index marking the end of the rapid increase of the signal using a similar logic.

Figure: Algorithm 2 Find Stabilization

Using the heuristic algorithm, we determine the value $k$ for experiments in Table 6, where we observe the determined $k$ are very close to each other except for VGG with CIFAR10 and XLM-R with SQuAD. All the $k$ value are shown visually in Figure 6, overlaying on top of the learning dynamics.

**Table 6: Different $k$ determined by Algorithm 2.**
| Metric | MNIST RN18 | CIFAR10 RN18 | CIFAR100 RN18 | CIFAR10 VGG11 | SQuAD XLM-R | MNLI XLM-R |
| --- | --- | --- | --- | --- | --- | --- |
| $S^{\rho}_{worst}$ | 270 | 230 | 260 | 960 | 810 | 780 |
| $S^{\rho}_{avg}$ | 270 | 270 | 250 | 1010 | 1090 | 720 |
| $\operatorname{\texttt{tr}(F)}$ | 210 | 260 | 230 | 250 | 1310 | 790 |

Table 7 shows the complete results for our complex experimental setup (described in §4). While all results are better than the standard training, empirically, $\operatorname{\texttt{tr}(F)}$ is a metric that gives a better winning rate compared to a random hyperparameter search.

**Table 7: Cross-lingual transfer results of standard training and using all 3 metrics to determine the unfreezing interval $k$ for gradual unfreezing (GU). WR stands for winning rate, averaged over 10 randomly sampled $k$ per training dataset.**
|  | XQuAD |  | MLQA |  | XNLI | WR |
| --- | --- | --- | --- | --- | --- | --- |
| Method | F1- En/X-ling | EM- En/X-ling | F1- X-ling | EM- X-ling | Avg- En/X-ling |  |
| Standard | 82.96/68.72 | 71.39/52.64 | 56.27 | 40.93 | 83.17/71.84 | - |
| GU${}_{S^{\rho}_{worst}}$ | 83.78/70.09 | 72.10/54.17 | 57.86 | 42.02 | 82.83/72.13 | 45% |
| GU${}_{S^{\rho}_{avg}}$ | 83.84/70.00 | 72.12/53.69 | 58.10 | 42.03 | 83.03/72.27 | 40% |
| GU${}_{\operatorname{\texttt{tr}(F)}}$ | 83.77/70.70 | 72.33/54.40 | 58.47 | 42.31 | 83.36/72.49 | 80% |

Figure: Figure 6: Learning dynamics with determined value of $k$ (vertical line) on the figure.

## Appendix E Additional Learning Dynamics

### E.1 Ranks

Figure 7 shows the evolution of feature ranks before the classification head for the first 2000 training steps.
We observe standard training typically starts with a lower feature rank, as the training progresses the feature rank gradually increases.
When withholding parameters from training, feature ranks are high at the initial period of learning, and as we gradually release parameters, the feature ranks reduce compared to the initial value.

Figure: Figure 7: Change of feature ranks before the classification head. The sudden decrease in feature ranks is due to unfreezing the trainable parameters.

### E.2 Change Learning Rate from Low to High

We provide an example of $S^{\rho}_{avg}$ during training when switching from a low learning rate to a high learning rate at $k=200$ in Figure 8. From Figure 8, we can see the $S^{\rho}_{avg}$ is high initially, then quickly drops when the learning rate increases at step k. More interestingly, switching the learning rate induces similar effects as release parameters from frozen.

Figure: Figure 8: Learning Dynamics of CIFAR10 on ResNet18, when training with 0.1*lr${}_{d}$ at the initial 200 steps, then switched to lr${}_{d}$ after 200 steps using standard training.

### E.3 SQuAD

In Figure 9, we present the learning dynamics for XLM-R with SQuAD in the early period of learning. The learning dynamics show a similar trend as the SQuAD dataset, the $S^{\rho}_{worst}$ value is also negative, and withholding trainable parameters increases the $S^{\rho}_{worst}$ during training based on Eqn. 4 in our main paper.

Figure: Figure 9: Learning Dynamics of XLM-R with LoRA training with SQuAD, y-axis for figures are in the log scale, the original value sharpness value in sub-figure (c) is negative where we take the absolute value before visualization. All values are normalized between 0 and 1000 for visualization.

### E.4 1/10th of the Default Learning Rate

Figure 10 shows the learning dynamics in the early period of training using 1/10th of the default learning rate (i.e. 0.1*lr${}_{d}$). The trends are similar to using the default learning rate.

Figure: Figure 10: Unfreezing parameters at different times affect the learning dynamics in the early period of training. Sub-figures (a)-(c): The $\operatorname{\texttt{tr}(F)}$ when unfreezing parameters at k= [250,750] steps versus standard training. Sub-figures (d)-(f): The average-case $L2$ sharpness (Eqn. 3) when unfreezing parameters at k= [250,750] steps versus standard training. Sub-figures (g)-(i): The worst-case $L2$ sharpness (Eqn. 4) when unfreezing parameters at k= [250,750] steps versus standard training. The y-axis is in the log scale.
Refer to caption: x2.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: see the weaknesses section.
soundness: 2
strengths: - The Discovery of the early learning dynamics and its impact on OOD generalization is an interesting finding, and the empirical evaluations seem to show that they indeed exist.
- The evaluations are done across various tasks from images and language, and via different datasets which provides generality to the observed results.
- Both FI and sharpness are well-studied topics which have been connected to generalization from both theoretical and empirical standpoints, hence, well suited for the development of the proposed method to improve OOD generalization.
summary: This work investigates the impact of interventions by weight freezing during early stage of training on out of distribution generalization and reports empirical evidence on 3 different OOD tasks based on covariate shift using image and language data.
They further hypothesize that using Fisher information and sharpness measures one can detect learning dynamic phase changes to leverage this phenomenon as an effective learning algorithm.
Their analysis and results provides some evidence of the. usefulness of the proposed algorithm.
weaknesses: The proposed algorithm to improve OOD generalization is --as it is also mentioned by the authors-- heuristic and does not provide deeper understanding of how these changes in learning dynamics occur or is connected to generalization, beyond what has been discovered already in the literature.The gains from the algorithm that are presented in tables 1 and 2 also seem to be marginal, and it is unclear how well they perform compared to SoTA OOD, as no real baseline has been used in this work to demonstrate how the proposed algorithm compares to latest advances in OOD generalization. Furthermore, results are provided only partially, and it is unclear how the proposed algorithm performs on domain adaptation tasks. Please provide comparisons to SoTA, and extend results to domain adaptation using the proposed algorithm. Further deepen the analysis and hypothesize why this phenomenon occurs.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Is the proposed heuristic method computationally expensive? (Because you need to train an extra time and calculate the relevant metric)
soundness: 3
strengths: 1. The finding of the significant impact of the early stage training on OOD performances is novel.
2. The experimental results include comprehensive OOD tasks for both vision and language domains.
3. The empirical finding that low sharpness doesn't necessarily improve OOD generalization in Figure 5 is meaningful, challenging the common belief that flat loss landscapes could benefit generalization.
summary: This paper investigates how gradual unfreezing the parameters during the early training stage affects OOD performances, and empirically unveils that gradual unfreezing leads to a time-sensitive trade-off between OOD and ID performance. The authors provide an explanation of the role of gradual unfreezing on OOD generalization: gradual unfreezing could help better to align early mini-batch gradients to the full-batch gradient, thus preventing potential overfitting to the mini-batches and eliminating spurious features. Moreover, they find that gradual unfreezing increases the sharpness and Fisher Information of the model parameters. The authors also propose a heuristic algorithm to find the best training step to start gradual unfreezing, gaining some empirical improvements on OOD tasks.
weaknesses: 1. Why would starting GU at the end of a dramatic change in sharpness help OOD? Since sharpness cannot reflect OOD performance, why is it used as an indicator of the timing of start GU?
2. Section 6.2 is insufficient to validate claim 2) of the hypothesis, since there lacks a clear correlation between the changes in the sharpness and the OOD performance across the training steps. It is best to plot the "sharpness-training step" curve and the "OOD acc-GU start step" curve in a single graph for all tasks and models to more clearly show the correlation between the changes in the sharpness and the best interval to start GU. Otherwise, it is hard to say that it is always best to start GU at the end of a sharp change in sharpness just by the simple MNIST experiment in line 440.
3. Table 1 and 2 merely shows the winning rate, rather than show the concrete number of improvements, which could be less convincing.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: **Q1** In the conclusion, you claim that the paper contributes to a deeper understanding of the early period of training. Could you clarify what this understanding entails? It is unclear from how I read the paper. What does the empirical performance tells us about what happens during the early training phase within the network?

**Q2** After designing the final method, it would be interesting to apply it to larger scale, real-world applications cases (without much tuning). For instance, how does the method perform for increasing robustness on a dataset like ImageNet-C or ImageNet-R?

**Q3** The title references “adaptation”, but this is not part of the investigation. Could you clarify?
soundness: 2
strengths: The paper is generally easy to follow. The problem setup is clear, the proposed method conceptually simple to implement. A variety of datasets are included in the study. The results show possibilities to increased OOD performance on these datasets with the right gradual unfreezing schedule.
summary: The authors investigate whether gradual increase of the parameters in a neural network during early training helps with boosting OOD performance. To this end, the authors propose simple strategy of gradually training more parameters in a network, starting with the earlier layers and blocks until the full model is trained. This introduces a parameter, $k$, which determines the #steps until the next block in unfrozen. ID and OOD performance is tracked with this strategy on a variety of language and vision datasets under different distribution shifts. With an optimal $k$, the authors show boost in OOD performance. In the second half of the paper, the authors discuss how $k$ can be optimally chosen based on Fisher Information or sharpness, and results are presented under this heuristic.
weaknesses: ### Summary

The paper has two big issues:

First, the flow of writing: The first half of the paper addresses a possibility to increase the OOD performance depending on optimal selection of the unfreezing schedule, with little control experiments. These results appear strong. The second half of the paper then discusses how to select the unfreezing schedule, which sets the results much more in perspective. A better way, in my opinions, would be to upfront show the final method and its empirical performance: what is the proposed selection scheme for $k$, and what is the resulting performance on ID and OOD test cases across dataset and models?

Second, missing controls and baselines. The paper claims to contribute to the “understanding” of early training dynamics, but really only shows one special case of addressing this problem. The authors claim previous methods focus on ID performance and failed to address OOD improvements, but none of the methods were implemented and ran as baselines. The unfreezing schedule is also not well motivated and arbitrarily presented. What about randomly unfreezing? What about doing this not per block, but truly randomly across network parameters? What about unfreezing top to bottom? How do the methods stack up against established strategies of 

I feel like addressing any subset of these questions would greatly enhance readability, quality of investigation, and impact of the paper. Please find an additional list of weaknesses to address below (if possible, please ref their labels W1,... etc) during the rebuttal:

### Major Weaknesses

**W1.** The results presented in Figure 2 and 3 are interesting, but not sufficient to back up the authors’ claims about improvement of OOD performance. Namely, all metrics are reported directly on the test sets, which in practice are not observable, and this is not appropriately discussed in the paper. It is unclear how the optimal k in Fig. 2 and 3 would be chosen. In a typical experiments, it would be selected based on ID validation performance, and for quite a few of the presented settings, this would result in much smaller gains (or even decrease in performance).

**W2.** it is good that experiments were run across multiple seeds, but the authors should reflect these in the plots and update all plots with appropriate error bars (e.g. SEM or 95% CIs), e.g. in Fig. 2-4, and also in Fig. 5. In the tables, the WR should also be equipped with error bars.

**W3.** Is it actually required to unfreeze the model step by step? There are little controls against this proposal. What happens if blocks are unfrozen randomly? Or from last to first layer, instead vice-versa? More control experiments here would strengthen the claims made.

**W4.** The authors cite other strategies for adapting early changes to network training for improved ID performance. How do these strategies stack up against the proposed method for OOD improvements? Adding some of these comparisons would strengthen the proposed method.

**W5.** Section 3.1 is not well written. Not all variables and symbols are defined (e.g. P_w in Eq. (1)), and some of the sentences are broken. There are also multiple statements about the Fisher matrix, like “A larger Fisher information…”, “tr(F) correlates well with the full Fisher information…” which are all imprecise given that F is a matrix. I would suggest to rewrite 3.1 to improve clarity.

**W6.** The study is limited in the number of models that are investigated. It would corroborate the paper story if the findings could be demonstrated for a larger set of model architectures.

**W7.** Significance: The performance improvements in Table 1 are really slim, except for MNIST. On CIFAR, we get 72.36 vs. 73.56 and 45.10 vs. 45.82%. It would really help to contextualise this against other ways of making model training more robust. What happens when models are trained with robust pre-training techniques? Here, improvements are typically much bigger than 1-2% points.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the above ‘weakness’ section and provide clarification if needed.

Other questions are: 

1. In Figure 5, why the k used in different subfigures different?
2. Figure 5 again, it seems that the pattern of ‘rapid change followed by stablization phase’ does not hold on the first plot of subfigure (b)(c).

Minor:

1. For a first-time reader, it is not clear what is the ‘removal of interventions’ at row 24 of the abstract, maybe consider just using ‘gradual unfreezing’ instead.
soundness: 3
strengths: 1. The paper is clear and easy to follow.
2. The experiments span various datasets and training setups.
3. The study of covariate shift early training dynamics, to the reviewer's knowledge, is new.
summary: This paper studies the early learning dynamics of OOD generalization under covariate shift. The authors leveraged gradual unfreezing (i.e. gradually unfreeze blocks of parameters during training) as a tool and observed that OOD generalization can be improved if gradual unfreezing is applied at the beginning of training. In addition, they also observe a general trade-off between ID and OOD performance when varying the time interval of unfreezing. They further use Fisher information and loss sharpness as proxies to design heuristics that allow for capturing the optimal time interval of gradual unfreezing.
weaknesses: 1. Lack of explanation on why the approach should work: while the reviewer acknowledges that (from its title) this paper is an empirical study, they found it hard to evaluate the value of findings in the paper without more explanation or intuition, especially the author proposes a heuristics-driven approach to improve OOD.
    1. explanation of gradual unfreezing: while the author provided intuition on page 6 about why gradual unfreezing would work, the reviewer found it not convincing: most of the experiments in Kumar et al. [1], if the reviewer remembers correctly, train linear probe then fine-tuning till convergence, but this paper considers early-training with small amount of training steps. In addition, Kumar et al. [1] consider transferring from a pre-trained model (that’s why the feature can be distorted). In contrast, in this paper, the authors train from scratch for part of the experiments. It is hard to imagine the benefits of first aligning the classification head (or top layers) while keeping the bottom layers (feature extractor) randomly fixed. 
    2. On phases of the change of the metrics: Besides, the dynamics of metrics in Figure 5 are very different, where in some cases they are decreasing but in other cases they are increasing. 
2. Related works: the reviewer found it slightly over-claimed to say that ‘the impact of the early learning period on OOD generalization remains unexplored’ throughout the paper. How the (early) learning dynamics influence the OOD generalization has been widely explored in spurious correlation, which is another major OOD generalization problem. It would be better if the authors could discuss the connection to these related works and rephrase their claims in the paper. Some examples are [2][3].
3. Robustness of the proposed heuristics: how the proposed heuristics would be able to transfer to real-world dataset is unknown. While the authors conduct experiments on multiple datasets, it is still hard to evaluate because of the lack of explanations of some components in this work (see 1.)
4. Ablation study: while we see that in Figures 2 and 3 there exists a sweet spot where ID performance remains stable while OOD performance improves, all the runs use gradual unfreezing at the beginning. A helpful ablation study shows that gradual unfreezing at the middle period or ending period of training does not exhibit this observation. That is, starting training the whole network for some epochs and then freeze all parameters and gradually unfreeze.

[1] Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution. Kumar et al. ICLR 2022  
[2] Identifying Spurious Biases Early in Training through the Lens of Simplicity Bias. Yang et al. AISTATS 2024  
[3] Complexity Matters: Feature Learning in the Presence of Spurious Correlations. Qiu et al. ICML 2024

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Could additional baselines from existing OOD algorithms be included to better evaluate the impact of the proposed method?

2. How does the gradual unfreezing strategy compare to other parameter-freezing approaches during early training? Could the similarities and differences be further clarified?

3. Have similar changes in Fisher information or sharpness been observed in other OOD algorithms, and if so, how do they relate to the findings in this paper?
soundness: 2
strengths: 1. This work is the first to highlight the impact of the early training period on OOD generalization, offering a novel perspective.

2. The use of relative values of Fisher information and sharpness as metrics for OOD generalization is interesting.

3. The gradual unfreezing method is tested across a wide range of transfer learning tasks and backbone networks, demonstrating its general applicability.
summary: This paper investigates the connection between early training dynamics and out-of-distribution (OOD) generalization. By using the trace of Fisher information and sharpness as indicators to examine the effects of interventions, the authors propose a gradual unfreezing strategy that enhances OOD generalization during the early training phase without impacting in-distribution (ID) performance.
weaknesses: 1. The paper lacks comparisons to relevant baselines in some transfer learning tasks, such as Office-Home and DomainNet. Many studies in domain generalization (DG) or OOD research focus on these datasets, but the absolute improvements in OOD accuracy presented here seem modest compared to existing methods.

2. Previous work, such as [1], has shown that freezing parameters improves OOD performance. It is unclear whether the gradual unfreezing method can be considered a variation or a weaker form of such approaches. The discussion in Section 5.2 touches on this but remains somewhat unclear.

3. From a causal standpoint, it is difficult to establish a causal relationship between sharpness and OOD performance rather than mere correlation. It is possible that gradual unfreezing simultaneously affects both sharpness and OOD performance. Causal intervention strategies might be needed to investigate this in more detail, for example, by ensuring sharpness remains constant during optimization after freezing and observing the potential drop in OOD performance.

[1] Kumar A, Raghunathan A, Jones R, et al. Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution[C]//International Conference on Learning Representations. 2022.

Meta-review: This paper investigates the training dynamic of deep neural networks and how it impacts OOD generalization. In particular, this is the first paper that focuses on the impact of early training stages. They achieve this by gradually unfreezing, which is a method that existed before. The authors maintain that their key contribution is not the method but correlating early training dynamics with OOD generalization. The reviewers did not converge on an acceptance decision and no reviewer was willing to champion the paper. The main criticism, also acknowledged by the authors, is that the paper does not explore sufficiently how to improve OOD generalization by acting on the early phases of training. This would be important in a future resubmission to establish that the correlation the authors observed is indeed actionable to improve robustness. The reviewers mention several avenue for improvement, as: 
* why unfreezing increases OOD generalization? what happens exactly and analytically, and in a larger scale, that contributes to such a phenomena?
* how can we systematically achieve predictive OOD generalization (using unfreezing)?
* why (analytically, and in a larger scale) OOD increases with unfreezing? what is the main factor, and to what extend can it increase? is there a bound to this improvement? are there other contributing factor? the unfreezing by itself involves in other dimensions such as optimization, learning capacity, architectural change, learning dynamics and so on. All those aspects, and perhaps more, need to be carefully studied, as the concept of unfreezing by itself does not tell us much about OOD generalization.

Further, reviewer pcei raised several points that the authors ignored in their reply,

Decision: Reject