MMD-Balls as Credal Sets: A PAC-Bayesian Framework for Epistemic Uncertainty in Test-Time Adaptation
For machine learning practitioners deploying models under distribution shift, this work provides formal guarantees linking shift magnitude to prediction reliability, enabling principled uncertainty quantification and adaptation decisions.
The paper develops a PAC-Bayesian framework for test-time adaptation that provides generalization bounds parameterized by the maximum mean discrepancy (MMD) between source and target distributions, interpreting MMD-balls as credal sets to quantify epistemic uncertainty. The framework yields a uniform worst-case risk bound and a decision criterion for when adaptation is warranted.
Test-time adaptation (TTA) methods improve model performance under distribution shift but lack formal guarantees connecting shift magnitude to prediction reliability. We develop a PAC-Bayesian framework yielding generalization bounds explicitly parameterized by the maximum mean discrepancy (MMD) between source and target distributions. Our principal contribution is interpreting MMD-balls around the source distribution as credal sets in Walley's imprecise probability theory, yielding natural epistemic uncertainty quantification. We establish: (i) a PAC-Bayesian bound with an MMD-dependent shift penalty under an RKHS-Lipschitz loss assumption; (ii) a finite-sample version via MMD concentration; (iii) a uniform worst-case risk bound over all distributions in the credal set, with a lower-upper risk decomposition; and (iv) geodesic preservation bounds explaining why kernel-guided adaptation protects local feature geometry. The credal set interpretation separates epistemic from aleatoric uncertainty and provides a principled decision criterion for when adaptation is warranted.