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

TL;DR

This work investigates how the very early phase of neural network training shapes out-of-distribution generalization under covariate shift. Using gradual unfreezing to vary the number of trainable parameters during the initial training period, the authors monitor the Fisher Information trace $tr(F)$ and sharpness measures $S^{rho}_{avg}$ and $S^{rho}_{worst}$ to identify when to remove the intervention. They find that early, time-sensitive interventions can yield Pareto improvements in ID and OOD performance across both image and language tasks, and that $tr(F)$ and sharpness metrics can guide the timing of intervention removal, though they do not consistently predict final OOD performance themselves. The results point to a new direction in understanding and leveraging early learning dynamics to improve robust generalization with minimal added complexity, suggesting both practical training strategies and theoretical questions for future work.

Abstract

Prior research shows that differences in the early period of neural network training significantly impact the performance of in-distribution (ID) data of tasks. Yet, the implications of early learning dynamics on out-of-distribution (OOD) generalization remain poorly understood, primarily due to the complexities and limitations of existing analytical techniques. In this work, we investigate the relationship between learning dynamics, OOD generalization under covariate shift and the early period of neural network training. We utilize the trace of Fisher Information and sharpness, focusing on gradual unfreezing (i.e., progressively unfreezing parameters during training) as our methodology for investigation. Through a series of empirical experiments, we show that 1) changing the number of trainable parameters during the early period of training via gradual unfreezing can significantly improve OOD results; 2) the trace of Fisher Information and sharpness can be used as indicators for the removal of gradual unfreezing during the early period of training for better OOD generalization. Our experiments on both image and text data show that the early period of training is a general phenomenon that can provide Pareto improvements in ID and OOD performance with minimal complexity. Our work represents a first step towards understanding how early learning dynamics affect neural network OOD generalization under covariate shift and suggests a new avenue to improve and study this problem.

Figures (18)

• Figure 1: (Left) Interventions during the early period of training are applied for a much shorter time. (Right) Impact of intervention in the early period of training on OOD performance across diverse settings (CIFAR10, cifar10cifarc; Office-Home, officehome/cvpr/VenkateswaraECP17; XNLI, conneau-etal-2018-xnli). $^*$ indicates optimal OOD results (§\ref{['sec:kidood']}).
• Figure 2: Changes in ID and OOD (noise-corrupted images) evaluation results when unfreezing parameters at different times (i.e., $k$) highlight the early training period's impact on OOD generalization. $\Delta_{acc}$ is calculated by subtracting gradual unfreezing results from standard training. The x-axis is in the log scale. Each data point on the plot is obtained by averaging over 6 runs for MNIST and 4 runs for CIFAR datasets (a total of 166 experiments per subfigure).
• Figure 3: Changes in ID and OOD evaluation results when unfreezing parameters at different times (i.e., $k$) for domain shift (vision) and language shift (text) with pre-trained transformers. $\Delta$ is calculated by subtracting the gradual unfreezing results from standard training, averaged over 4 runs (168 experiments in total for subfigure (a) due to single source-domain training and 68 experiments for subfigure (b)). The x-axis is in the log scale.
• Figure 4: Gradient similarity (mini-batch vs full-data) for the classification head of a ResNet18 trained with CIFAR10. The mini-batch gradient is more similar to the full-data gradient in the early period of training when gradual unfreezing is applied (K=100) compared to standard training (K=0).
• Figure 5: Learning dynamics with three metrics: $\mathop{\mathrm{\texttt{tr}(F)}}\nolimits$, $S^{\rho}{avg}$, and $S^{\rho}{worst}$. Unfreezing parameters at different steps impact early learning dynamics. The y-axis is log-scaled and normalized between 0 and 1000 for clarity.