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

2024-03-22

AI Summary

This work identifies the initial training dynamics of neural networks as a critical control window for improving out-of-distribution (OOD) generalization—particularly under covariate shift. Addressing the common trade-off between OOD robustness and in-distribution (ID) performance, we propose a progressive unfreezing strategy: during early training, the number of trainable parameters is dynamically constrained, and the timing of parameter unfreezing is adaptively determined using the Fisher information trace and loss landscape sharpness as practical criteria. To our knowledge, this is the first work to establish these two metrics as empirically grounded, actionable unfreezing indicators. The method achieves Pareto improvements in both ID and OOD performance. Experiments across image and text benchmarks demonstrate substantial gains in OOD accuracy while preserving or even enhancing ID accuracy—without increasing model capacity or inference cost.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

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.

Problem

• Neural Network Learning
• Parameter Unfreezing
• Fisher Information Optimization

Innovation

• Fisher Information
• Parameter Unfreezing Strategy
• Generalization Enhancement