📊 statistics

MMD-Balls as Credal Sets: A PAC-Bayesian Framework for Epistemic Uncertainty in Test-Time Adaptation

This paper proposes a PAC-Bayesian framework for test-time adaptation that interprets Maximum Mean Discrepancy (MMD) balls as credal sets to derive generalization bounds explicitly dependent on distribution shift magnitude, thereby providing a principled method for quantifying epistemic uncertainty and guiding adaptation decisions.

Original authors: Ahanaf Hasan Ariq

Published 2026-05-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ahanaf Hasan Ariq

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef who has spent years perfecting a recipe for a specific type of soup using ingredients from a local farm (the Source). You know exactly how that soup tastes, and you are confident in your predictions about how it will turn out.

Now, imagine you are suddenly deployed to a new kitchen in a different city (the Target). The local farm there grows slightly different vegetables. The soup might still taste good, or it might taste terrible. The problem is, you don't know how different the new vegetables are, and you don't have a way to measure the risk of serving a bad bowl of soup.

This paper proposes a new "safety meter" for machine learning models in this exact situation. It's called Test-Time Adaptation (TTA), where a model tries to adjust itself on the fly using new, unlabeled data. But until now, these models had no way to say, "Hey, the ingredients are too different; I shouldn't trust my prediction."

Here is how the paper solves this, using simple analogies:

1. The "Uncertainty Bubble" (The MMD-Ball)

The authors introduce a concept called MMD (Maximum Mean Discrepancy). Think of this as a high-tech ruler that measures how different two groups of ingredients are.

The Idea: Instead of just saying "The new farm is different," the paper draws a bubble (a ball) around the original farm's ingredients.
The Size of the Bubble: The size of this bubble represents how much the new ingredients could possibly differ. If the bubble is small, the new farm is very similar to the old one. If the bubble is huge, the new farm could be growing something totally wild.
The Credal Set: The paper calls this bubble a "Credal Set." In plain English, this is a "set of all possible realities." It says, "We don't know exactly what the new ingredients are, but we know they are somewhere inside this bubble."

2. The "Safety Guarantee" (The PAC-Bayesian Bound)

The paper creates a mathematical formula (a bound) that tells you exactly how bad your soup could taste based on the size of that bubble.

The Formula: Total Risk = (How hard the recipe is) + (How much you changed the recipe) + (How big the bubble is)
The Magic: The "How big the bubble is" part is the Epistemic Uncertainty. This is the paper's big breakthrough. It separates:
- Aleatoric Uncertainty: "This recipe is just naturally hard to get right." (Inherent noise).
- Epistemic Uncertainty: "I don't know what the new ingredients are, so I'm guessing." (Ignorance due to shift).

If the bubble gets too big, the formula tells the model: "Your risk is too high. Stop guessing. Either ask for help or refuse to serve the soup."

3. The "Geometry of Taste" (Geodesic Preservation)

The paper also explains why using this specific "bubble" method is better than other methods.

The Problem with Other Methods: Some current methods try to fix the soup by just making the flavor "flatter" or more generic (Entropy Minimization). The authors argue this is like smashing all the vegetables into a mush. It might taste okay for common veggies, but it destroys the unique shape of rare, delicate herbs (rare classes).
The Paper's Solution: Their method acts like a "geometric guardian." It ensures that even if the ingredients change, the shape and structure of the flavor profile stay intact. It preserves the "local geometry," meaning rare and unique flavors don't get crushed just because they are rare.

4. The Decision Maker

Finally, the paper gives a clear rule for when to adapt:

Small Bubble: The new ingredients are close enough. You can trust your current prediction. No need to panic.
Medium Bubble: The ingredients are different, but you can adjust your recipe (adapt) to stay safe.
Huge Bubble: The ingredients are too different. The "Upper Risk" is too high. The model should say, "I don't know what's going on here," and stop making predictions rather than guessing wrong.

Summary

In short, this paper builds a trustworthy safety net for AI. It uses a mathematical "bubble" to measure how much the world has changed since the AI was trained. It then calculates a strict "worst-case scenario" for how wrong the AI might be. If the bubble is too big, the AI knows to stop and say, "I'm not sure," rather than confidently giving a wrong answer. This is crucial for safety-critical tasks like self-driving cars or medical diagnosis, where a silent mistake can be dangerous.

Technical Summary: MMD-Balls as Credal Sets for Epistemic Uncertainty in Test-Time Adaptation

1. Problem Statement

The paper addresses a critical gap in Test-Time Adaptation (TTA): while existing methods (e.g., entropy minimization, regularization) improve model accuracy under distribution shifts, they lack formal guarantees connecting the magnitude of the shift to prediction reliability. In safety-critical domains (autonomous driving, medical imaging), models may silently degrade when encountering target distributions ( $P_t$ ) that differ from the source training distribution ( $P_s$ ).

Current uncertainty quantification methods often conflate aleatoric uncertainty (inherent data noise) with epistemic uncertainty (uncertainty due to limited knowledge of the data-generating process). Furthermore, they fail to provide formal bounds on how risk degrades as a function of distributional shift. The authors formalize the core question: Can we bound target risk as a function of distribution shift to provide actionable epistemic uncertainty?

2. Methodology and Theoretical Framework

The paper proposes a unified framework integrating PAC-Bayesian theory, Kernel Mean Embeddings (MMD), and Walley's Imprecise Probability theory.

2.1 Core Definitions

RKHS-Lipschitz Loss (Assumption 1): The conditional expected loss function $L(w, \cdot)$ is assumed to belong to a Reproducing Kernel Hilbert Space (RKHS) $\mathcal{H}$ with a bounded norm $\|L(w, \cdot)\|_\mathcal{H} \leq L_H$ . This allows the use of kernel properties to bound risk differences.
MMD-Induced Credal Set: The authors define a set of plausible target distributions as an MMD-ball around the source:
$C_\varepsilon(P_s) = \{Q \in \mathcal{P}(\mathcal{X}) : \text{MMD}(P_s, Q) \leq \varepsilon\}$
Interpreted through Walley's theory, this set represents all distributions indistinguishable from the source at resolution $\varepsilon$ , forming a credal set that quantifies epistemic uncertainty.

2.2 Theoretical Derivations

PAC-Bayesian Bound with MMD Penalty (Theorem 1):
The authors derive a generalization bound for the target risk $R_{P_t}(\rho)$ that explicitly includes an MMD-dependent shift penalty:
$R_{P_t}(\rho) \leq \hat{R}_{P_s}(\rho) + \sqrt{\frac{KL(\rho\|\pi) + \log(2\sqrt{n}/\delta)}{2n}} + L_H \cdot \text{MMD}(P_s, P_t)$
This decomposes risk into: (i) Aleatoric uncertainty (empirical risk), (ii) Estimation uncertainty (PAC-Bayesian complexity), and (iii) Epistemic shift penalty (linear in MMD).
Finite-Sample Version (Theorem 3):
Since population MMD is unknown, the paper provides a computable bound using the unbiased MMD estimator $\widehat{\text{MMD}}_u$ and a concentration term $\varepsilon_{m,n}$ derived from sub-Gaussian concentration inequalities. The bound holds with probability $1-\delta$ and achieves the minimax-optimal rate $O(1/\sqrt{n})$ .
Risk Decomposition (Corollary 9):
By defining lower ( $\underline{R}_\varepsilon$ ) and upper ( $\overline{R}_\varepsilon$ ) risks over the credal set, the framework yields an imprecision interval. The width of this interval separates epistemic uncertainty (driven by the shift penalty $L_H \cdot \varepsilon$ and complexity) from aleatoric uncertainty (captured by the empirical risk).
Geometric Preservation (Proposition 10):
The paper establishes bounds on geodesic distortion in the RKHS. It shows that under bounded feature map assumptions, the distortion of local feature geometry between source and target neighborhoods is controlled by the MMD. This provides a theoretical explanation for why kernel-guided adaptation preserves the structure of rare classes better than entropy-based methods, which may collapse rare-class regions.

3. Key Contributions

The paper claims four primary theoretical contributions:

A PAC-Bayesian Bound with MMD Shift Penalty: A generalization bound (Theorem 1) and its finite-sample counterpart (Theorem 3) that explicitly parameterizes risk degradation by the distribution shift magnitude.
Credal Set Interpretation: A novel interpretation of MMD-balls as credal sets (Proposition 7), providing worst-case risk guarantees over the entire set of plausible distributions.
Uncertainty Decomposition: A formal lower-upper risk decomposition (Corollary 9) that separates epistemic uncertainty (distribution shift) from aleatoric uncertainty (task difficulty), addressing the call by Hüllermeier & Waegeman (2021) for principled uncertainty separation.
Geodesic Preservation: Bounds explaining why kernel-guided adaptation protects local feature geometry (Proposition 10), offering a theoretical basis for robustness against rare-class collapse.

4. Results and Implications

Decision Criterion for Adaptation: The imprecision interval $[\underline{R}_\varepsilon, \overline{R}_\varepsilon]$ serves as a decision criterion. Adaptation is warranted when the upper risk exceeds a tolerance threshold $r_{max}$ while the lower risk does not. If the lower risk already exceeds $r_{max}$ , adaptation is deemed futile.
Hypothesis Testing Connection: The width $\varepsilon$ of the credal set can be calibrated using the asymptotic null distribution of the MMD two-sample test, linking classical statistical testing to epistemic uncertainty quantification.
Rare-Class Robustness: The geometric preservation bounds suggest that controlling MMD prevents the distortion of local feature manifolds, offering a formal argument for the robustness of kernel-guided TTA methods compared to entropy minimization.

5. Significance and Claims

The paper positions itself as the first to formulate TTA uncertainty through MMD-induced credal sets, unifying PAC-Bayesian generalization, kernel methods, and imprecise probability.

Trustworthy Deployment: The framework enables systems to trigger abstention or fallback mechanisms when epistemic uncertainty (measured by the credal width) exceeds a tolerable threshold, aligning with the goals of Epistemic Intelligence.
Formal Guarantees: Unlike existing TTA methods that rely on heuristics, this approach provides formal, data-dependent bounds connecting shift magnitude to risk.
Relation to Conformal Prediction: The authors note that while conformal prediction handles predictive uncertainty (marginal coverage), their credal width $\varepsilon$ handles distributional epistemic uncertainty. They propose a future combination where $\varepsilon$ adapts the coverage level of conformal predictors, though they do not present experimental results for this combination.

Limitations Acknowledged:
The paper acknowledges that Assumption 1 (RKHS-Lipschitz loss) is non-trivial for deep networks with softmax outputs, though universality of RBF kernels offers informal support. It also notes that Assumption 2 (bounded feature map) holds approximately in specific regimes (e.g., NTK, spectral normalization) but requires further architectural-specific analysis. The MMD convergence rate, while minimax-optimal, may be loose in practice without adaptive kernel selection.