Higher-Order Asymptotics of Test-Time Adaptation for Batch Normalization Statistics
This work addresses the reliability and performance of Batch Normalization layers in adapting to changing data distributions, which is an incremental improvement for machine learning models facing distribution shifts.
This study tackled the problem of test-time adaptation for Batch Normalization statistics under distribution shift by developing a higher-order asymptotic framework, resulting in an optimal weighting parameter that minimizes mean-squared error and providing refined density and tail probability estimates for the adapted statistic.
This study develops a higher-order asymptotic framework for test-time adaptation (TTA) of Batch Normalization (BN) statistics under distribution shift by integrating classical Edgeworth expansion and saddlepoint approximation techniques with a novel one-step M-estimation perspective. By analyzing the statistical discrepancy between training and test distributions, we derive an Edgeworth expansion for the normalized difference in BN means and obtain an optimal weighting parameter that minimizes the mean-squared error of the adapted statistic. Reinterpreting BN TTA as a one-step M-estimator allows us to derive higher-order local asymptotic normality results, which incorporate skewness and other higher moments into the estimator's behavior. Moreover, we quantify the trade-offs among bias, variance, and skewness in the adaptation process and establish a corresponding generalization bound on the model risk. The refined saddlepoint approximations further deliver uniformly accurate density and tail probability estimates for the BN TTA statistic. These theoretical insights provide a comprehensive understanding of how higher-order corrections and robust one-step updating can enhance the reliability and performance of BN layers in adapting to changing data distributions.