Higher-Order Asymptotics of Test-Time Adaptation for Batch Normalization Statistics

TL;DR

This paper introduces a higher-order asymptotic framework for test-time adaptation of Batch Normalization statistics, combining Edgeworth expansion and saddlepoint techniques to improve understanding and performance under distribution shifts.

Contribution

It develops a novel higher-order asymptotic analysis for BN test-time adaptation, integrating classical methods with a new one-step M-estimation perspective.

Findings

Derived an Edgeworth expansion for BN mean differences.

Identified an optimal weighting parameter minimizing MSE.

Quantified bias, variance, and skewness trade-offs in adaptation.

Abstract

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.