<figure><img src="image_1.png" /></figure>

HAL
open science

# A Minimax Approach for Robustifying Statistical Learning Algorithms

Xinying Zou, Samir M. Perlaza, Iñaki Esnaola

► To cite this version:

Xinying Zou, Samir M. Perlaza, Iñaki Esnaola. A Minimax Approach for Robustifying Statistical Learning Algorithms. 9600, Inria. 2025. hal-05393845

HAL Id: hal-05393845

https://inria.hal.science/hal-05393845v1

Submitted on 2 Dec 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_2.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
informatics mathematics
Inria

# A Minimax Approach for Robustifying Statistical Learning Algorithms

Xinying Zou, Samir M. Perlaza, Iñaki Esnaola

## RESEARCH
REPORT

### N° 9600

December 2025

Project-Team NEO
<figure><img src="image_3.png" /></figure>
<figure><img src="image_4.png" /></figure>

# A Minimax Approach for Robustifying Statistical Learning Algorithms

Xinying Zou, Samir M. Perlaza, Iñaki Esnaola

Project-Team NEO

Research Report n° 9600 — December 2025 — 31 pages

**Abstract:** In this paper, we address the challenge of learning models that stay reliable under distribution shifts when only finite training data is available. A novel training-dependent minimax problem is proposed to design learning algorithms that are robust to the worst-case data-generating distribution. For the ambiguity sets, we construct Kullback-Leibler (KL) divergence neighborhoods on both model and data distributions. We obtain an analytical solution to this minimax problem, referred to as the robust learner. We show that the robust learner follows a Gibbs distribution, in which the prior can be chosen as any baseline learning algorithm to be robustified. Such a robust learner minimizes the worst-case generalization gap within a KL-divergence neighbourhood of unseen new data and it provides a smaller generalization error compared with the baseline learning algorithm Q. Under certain conditions, the robust learner also guarantees a smaller expected testing error than Q. We also provide a training-dependent PAC-Bayes bound on the robust learner's performance on unseen data. Closed-form expressions for generalization error and expected loss of the robust learner are given in terms of mutual information, lautum information, and KL-divergence. As a by-product, we show that the proposed minimax problem admits a two-player zero-sum game formulation, for which a unique Nash equilibrium exists. This enhances our understanding of learning algorithm robustification. Numerical experiments validate the applicability of the results and the benefits of robustification.

**Key-words:** Statistical Learning, Minimax, and Gibbs Algorithm

RESEARCH CENTRE
SOPHIA ANTIPOLIS – MÉDITERRANÉE

2004 route des Lucioles - BP 93
06902 Sophia Antipolis Cedex

<footnote>Xinying Zou and Samir M. Perlaza are with INRIA, Centre Inria d'Université Côte d'Azur, Sophia Antipolis 06902, France. Iñaki Esnaola is with the School of Electrical and Electronic Engineering at The University of Sheffield, Sheffield S1 3JD, UK. Samir M. Perlaza and Iñaki Esnaola are also with the Electrical and Computer Engineering Department at Princeton University, Princeton N.J. 08544, USA. Samir M. Perlaza is also with the GAATI Laboratory at the Université de la Polynésie Française, Faaa 98702, French Polynesia.</footnote>
Une approche minimax pour la robustification
des algorithmes d'apprentissage statistique

**Résumé** : Dans cet article, nous abordons le défi consistant à apprendre des modèles qui restent fiables en présence de décalages de distribution lorsque seules des données d'apprentissage finies sont disponibles. Nous proposons un nouveau problème minimax dépendant de l'entraînement afin de concevoir des algorithmes d'apprentissage robustes à la pire distribution génératrice de données. Pour les ensembles d'ambiguïté, nous construisons des voisinages en divergence de Kullback-Leibler (KL) à la fois sur les distributions de modèles et sur les distributions de données. Nous obtenons une solution analytique à ce problème minimax, appelée l'apprenant robuste. Nous montrons que l'apprenant robuste suit une distribution de Gibbs, dans laquelle le prior peut être choisi comme n'importe quel algorithme d'apprentissage de base à robustifier. Un tel apprenant robuste minimise l'écart de généralisation au pire cas au sein d'un voisinage, en divergence de KL, de nouvelles données non vues et fournit une erreur de généralisation plus faible que l'algorithme d'apprentissage de base Q. Sous certaines conditions, l'apprenant robuste garantit également une erreur de test espérée plus faible que Q. Nous fournissons également une borne PAC-Bayes dépendant de l'entraînement sur les performances de l'apprenant robuste sur des données non vues. Des expressions en forme fermée pour l'erreur de généralisation et la perte espérée de l'apprenant robuste sont données en fonction de l'information mutuelle, de la lautum information et de la divergence de KL. Comme corollaire, nous montrons que le problème minimax proposé admet une formulation en jeu à somme nulle à deux joueurs, pour laquelle il existe un unique équilibre de Nash. Cela améliore notre compréhension de la robustification des algorithmes d'apprentissage. Des expériences numériques valident l'applicabilité des résultats et les bénéfices de la robustification.

**Mots-clés** : Apprentissage Statistique, Minimax et algorithme de Gibbs
Contents

<table><tr><td>1</td><td>Introduction</td><td>4</td></tr><tr><td>1.1</td><td>Related Works</td><td>5</td></tr><tr><td>1.2</td><td>Contributions</td><td>5</td></tr><tr><td>2</td><td>Preliminaries</td><td>6</td></tr><tr><td>2.1</td><td>Notation</td><td>6</td></tr><tr><td>2.2</td><td>Supervised Machine Learning Problem</td><td>6</td></tr><tr><td>3</td><td>Main Results</td><td>8</td></tr><tr><td>3.1</td><td>A Robust Learner</td><td>8</td></tr><tr><td>3.2</td><td>Performance Analysis</td><td>10</td></tr><tr><td>3.3</td><td>Connections to Zero-Sum Games</td><td>12</td></tr><tr><td>4</td><td>Numerical Experiment</td><td>13</td></tr><tr><td>5</td><td>Conclusion and Future Work</td><td>19</td></tr><tr><td></td><td>Appendices</td><td>20</td></tr><tr><td>A</td><td>Preliminary Results</td><td>20</td></tr><tr><td>B</td><td>The Worst-Case Data-Generating Probability Measure</td><td>22</td></tr><tr><td>C</td><td>Proof of Theorem 1</td><td>23</td></tr><tr><td>D</td><td>Proof of Theorem 2</td><td>23</td></tr><tr><td>E</td><td>Proof of Lemma 1</td><td>24</td></tr><tr><td>F</td><td>Proof of Theorem 3</td><td>25</td></tr><tr><td>G</td><td>Proof of Theorem 4</td><td>25</td></tr><tr><td>H</td><td>Proof of Theorem 5</td><td>28</td></tr><tr><td>I</td><td>Numerical Experiments</td><td>28</td></tr></table>
# 1 Introduction

Machine learning algorithms are expected to generalize to patterns that are not present in the training data, which are likely coming from a testing distribution different from the training. However, often the learner has access only to a finite training dataset, and for such out-of-distribution settings, empirical risk performance alone is a poor proxy for performance on new data. This happens in two scenarios: either the learning algorithm is selecting from a rich hypothesis class that fits the training well but with the potential of overfitting on new data, or the fitting on the training data is poor already.

These challenges motivate us to design a learning rule that (i) incorporates inductive bias over models and (ii) is robust to the worst-case testing data.

In this paper, we consider the worst-case data distribution to be the one within a Kullback-Leibler (KL) divergence neighborhood of a reference measure. Formally, given a training dataset $z$; a reference measure $Q$ over the models $\theta \in \mathcal{M}$; a reference measure $P_S$ over the data $(x, y) \in \mathcal{X} \times \mathcal{Y}$; and a loss function $\ell: \mathcal{X} \times \mathcal{Y} \times \mathcal{M} \to [0, +\infty]$, consider the following minimax problem

$$ \min_{P \in \mathcal{Q}} \max_{P_{Z|\Theta} \in \mathcal{P}} \int \left( \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) - L(z, \theta) \right) dP(\theta), \quad (1a) $$

with

$$ \mathcal{Q} \triangleq \{ P \in \Delta_Q(\mathcal{M}) : D(P \| Q) \le \gamma_1 \}; \quad (1b) $$

and

$$ \mathcal{P} \triangleq \{ P_{Z|\Theta} \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) : \forall \theta \in \mathcal{M}, D(P_{Z|\Theta=\theta} \| P_S) \le \gamma_2 \}. \quad (1c) $$

Here, $L(z, \theta)$ is the empirical risk induced by model $\theta$ w.r.t the dataset $z$ defined in (5); $D(P \| Q)$ is the Kullback-Leibler (KL) divergence of $P$ w.r.t $Q$ (Appendix A); the sets $\Delta_Q(\mathcal{M})$ and $\Delta_{P_S}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})$ denote the sets of probability measures that are absolutely continuous w.r.t $Q$ and $P_S$, respectively; and for some given $\gamma_1, \gamma_2 > 0$. The inner maximization lets “nature" choose, for each model $\theta$, an adversarial testing data-generating probability measure within a KL-divergence ball centred at $P_S$ with radius $\gamma_2$; the outer minimization lets the learner choose a probability measure over the models within a KL-divergence ball centred at $Q$ with radius $\gamma_1$. The objective in (1a) explicitly features the difference between expected testing risk and expected empirical risk, thereby directly targeting the generalization gap. The references $P_S$ and $Q$ play complementary roles. The measure $P_S$ encodes prior information about the target (testing) data while $Q$ captures inductive bias. In particular, if $Q$ is the posterior produced by an existing algorithm, the solution to (1) can be seen as “robustifying" such a baseline $Q$, adapting it to distributional neighborhood of $P_S$ while controlling the deviation from $Q$. For instance, $Q$ can be the Gibbs algorithm that minimizes the KL-regularized empirical risk minimizer (see equation (11)) [1] or simply an initial distribution over the weights of neural networks.
## 1.1 Related Works

Distributionally robust optimization (DRO) has emerged as a powerful framework for decision-making under distributional uncertainty, where a learner competes against an adversary selecting the worst-case distribution from an ambiguity set. Early studies focus on moment-based ambiguity sets [2,3]. More recently, a popular choice of ambiguity set is defined as a neighborhood in the space of probability measures via statistical distances, such as f-divergence [4–9]. In [10], an ambiguity set is defined as a Wasserstein ball around the empirical measure of the training data and it was proved that the associated minimax problems often reduce to tractable convex programs and that finite-sample performance guarantees can be derived via concentration inequalities. In [11], strong duality for Wasserstein DRO, characterized the structure of worst-case distributions and connected these formulations for robust optimization. The choice of ambiguity set is essential to any DRO framework. In this paper, by adopting double KL-divergence neighborhoods, we are able to provide a closed-form solution for (1) with explicit formulation of the worst-case data-generating probability measure [12]. Note that in related works, such as [5, 13, 14], the minimization is often over the set of models in $\mathbb{R}^d$ and only upper and lower bounds for the distributionally robust risks are provided [4]. However, in this work, the minimization is over probability measures on the set of models and we derive a closed-form solution along with performance guarantees for general learning algorithms.

The concept of a worst-case data source has been widely adopted in statistical learning. Adversarial training [15,16] is a natural approach to improve the robustness of machine learning algorithms. By training on data that is designed in an adversarial minimax fashion, learning algorithms are forced to accommodate instances of data that are intentionally selected to strain its generalization capabilities [17]. The concept of a worst-case probability measure is widely adopted in robust learning to formulate a minimax optimization problem [18–20]. Domain adaptation [21] approaches to this problem, show that learning an invariant representation and achieving a small training error does not guarantee generalization on the testing domain. In [22], the authors introduce an expansion function for quantifying the difficulty of an out-of-distribution (OOD) generalization problem and prove novel OOD generalization error bounds in this setting. More recently, in [23], a framework that establishes information-theoretical OOD generalization error bounds is introduced.

## 1.2 Contributions

We construct a training-dependent minimax problem with double KL-divergence neighborhoods and show that it admits a unique analytical solution following a Gibbs distribution. Such a robust learner provides a simple framework to improve robustness of any baseline learning algorithm w.r.t unseen data. Such a robust learner minimizes the worst-case generalization gap over a KL-divergence neighbourhood of the target test distribution and, as a result, achieves a smaller
worst-case generalization error than the baseline learning algorithm $Q$. Under additional regularity conditions, the robust learner also guarantees a smaller expected test error than $Q$. A training-dependent PAC-Bayes bound is provided for such a robust learner. Moreover, by connecting to existing results, a closed-form generalization decomposition into mutual and lautum information quantifies the robust learner's generalization capability, and an explicit expression involving KL-divergences characterizes the expected loss difference between the robust learner and any learning algorithm on the same support. Furthermore, we show that the proposed minimax problem admits a two-player zero-sum game formulation, for which a unique Nash equilibrium exists. Finally, we support our theoretical guarantees with numerical experiments on MNIST, CIFAR-10 and CIFAR-10-C. We demonstrate the applicability of the robust learner on neural networks and propose a simple last-layer mixture algorithm to improve robustness.

# 2 Preliminaries

## 2.1 Notation

Sets are denoted by calligraphic uppercase letters, except sets of probability measures. Given a set $\mathcal{M}$ and a sigma-field $\mathcal{F}$ on $\mathcal{M}$, the set of all probability measures that can be defined on the measurable space $(\mathcal{M}, \mathcal{F})$ is denoted by $\Delta(\mathcal{M}, \mathcal{F})$ or simply $\Delta(\mathcal{M})$, when the sigma-field is fixed in the analysis. If $\mathcal{M} \subseteq \mathbb{R}^d$, then the Borel sigma-field on $\mathcal{M}$ is denoted by $\mathcal{B}(\mathcal{M})$. The subset of measures in $\Delta(\mathcal{M})$ that are absolutely continuous with a probability measure $Q$ is denoted by $\Delta_Q(\mathcal{M})$. Given a measure $P \in \Delta_Q(\mathcal{M})$, the Radon-Nikodym derivative of $P$ with respect to $Q$ is denoted by $\frac{dP}{dQ}$. Given a set $\mathcal{N}$, the set of all probability measures defined on $(\mathcal{M}, \mathcal{F})$ conditioned on an element of $\mathcal{N}$ is denoted by $\Delta(\mathcal{M}|\mathcal{N})$. More specifically, given a measure $P_{M|N} \in \Delta(\mathcal{M}|\mathcal{N})$ and $n \in \mathcal{N}$, the measure $P_{M|N=n}$ is in $\Delta(\mathcal{M})$.

## 2.2 Supervised Machine Learning Problem

Let $\mathcal{M}, \mathcal{X}$ and $\mathcal{Y}$, with $\mathcal{M} \subseteq \mathbb{R}^d$, be sets of models, patterns, and labels, respectively. A pair $(x, y) \in \mathcal{X} \times \mathcal{Y}$ is referred to as a labelled pattern or as a data point. Let the probability measure

$$P_Z \in \Delta(\mathcal{X} \times \mathcal{Y}) \qquad (2)$$

be the one from which data points are independently sampled. While the probability measure $P_Z$ is unknown in a practical setting, a training dataset $z \in (\mathcal{X} \times \mathcal{Y})^n$ is available in the form of a tuple of $n$ data points of the form:

$$z \triangleq ((x_1, y_1), (x_2, y_2), \dots, (x_n, y_n)) \in (\mathcal{X} \times \mathcal{Y})^n. \qquad (3)$$

Let the function $f: \mathcal{M} \times \mathcal{X} \to \mathcal{Y}$ be such that the label assigned to the pattern $x$ according to the model $\theta \in \mathcal{M}$ is $f(\theta, x)$. For instance, the function $f(\theta, \cdot):$
$\mathcal{X} \to \mathcal{Y}$ can represent a neural network whose weights are vectorized into $\theta$. Let the function $\hat{\ell}: \mathcal{Y} \times \mathcal{Y} \to [0, +\infty]$, be the loss function. Such a function is often of the form $\hat{\ell}(\mu, \nu) = |\mu - \nu|^p$, with $p \ge 1$; or $\hat{\ell}(\mu, \nu) = 1_{\{\mu=\nu\}}$. In general, given a labelled pattern $(x, y) \in \mathcal{X} \times \mathcal{Y}$, the loss induced by a model $\theta \in \mathcal{M}$ is $\hat{\ell}(f(\theta, x), y)$. In the following, the loss function $\hat{\ell}$ is assumed to be a general nonnegative function and for all $y \in \mathcal{Y}$, $\hat{\ell}(y, y) = 0$. For the ease of notation, consider the function $\ell: \mathcal{X} \times \mathcal{Y} \times \mathcal{M} \to [0, +\infty]$ such that

$$ \ell(\theta, x, y) = \hat{\ell}(f(\theta, x), y). \quad (4) $$

The empirical risk induced by the model $\theta \in \mathcal{M}$, with respect to the dataset $z$ in (3), is determined by the function $\mathbf{L}: (\mathcal{X} \times \mathcal{Y})^n \times \mathcal{M} \to [0, +\infty]$, which satisfies

$$ \mathbf{L}(z, \theta) = \frac{1}{n} \sum_{i=1}^{n} \ell(\theta, x_i, y_i), \quad (5) $$

where the function $\ell$ is defined in (4). In a broader sense, when data pairs are sampled identically according to a probability measure $P \in \Delta(\mathcal{X} \times \mathcal{Y})$, for a fixed model $\theta$, the expectation of the loss $\ell$ in (4) is denoted using the functional $\mathbf{R}_{\theta}: \Delta(\mathcal{X} \times \mathcal{Y}) \to [0, +\infty]$, which satisfies

$$ \mathbf{R}_{\theta}(P) = \int \ell(\theta, x, y) dP(x, y). \quad (6) $$

Alternatively, given a dataset $z$, which is kept fixed, the expectation of $\mathbf{L}(z, \theta)$ under the assumption that the model $\theta$ is drawn from a probability measure $P \in \Delta(\mathcal{M})$ is denoted using the functional $\mathbf{R}_z: \Delta(\mathcal{M}) \to \mathbb{R}$, which satisfies

$$ \mathbf{R}_z(P) = \int \mathbf{L}(z, \theta) dP(\theta). \quad (7) $$

Furthermore, the generalization error of a learning algorithm $P_{\Theta|Z} \in \Delta(\mathcal{M} | (\mathcal{X} \times \mathcal{Y})^n)$ is defined as the expectation of the generalization gap. Formally, let the functional $\bar{\bar{G}}: \Delta(\mathcal{M} | (\mathcal{X} \times \mathcal{Y})^n) \times \Delta(\mathcal{X} \times \mathcal{Y}) \times \Delta(\mathcal{X} \times \mathcal{Y})^n \to \mathbb{R}$ be such that

$$ \bar{\bar{G}}(P_{\Theta|Z}, \tilde{P}_Z, P_Z) = \iint (\mathbf{R}_{\theta}(\tilde{P}_Z) - \mathbf{L}(z, \theta)) dP_{\Theta|Z=z}(\theta) dP_Z(z), \quad (8) $$

where $\mathbf{R}_{\theta}(\tilde{P}_Z)$ is the population risk with $\mathbf{R}_{\theta}$ defined in (6) for some $\tilde{P}_Z \in \Delta(\mathcal{X} \times \mathcal{Y})$ and the measure $P_Z$ is a product measure formed by $P_Z$ in (2).

Using the notion of type (empirical measure) [24], the empirical risk $\mathbf{L}(z, \theta)$ in (5) can be written in terms of the functional $\mathbf{R}_{\theta}$ in (6). A type is defined as follows.

**Definition 1 (Type).** The type induced by the dataset $z$ in (3) is a probability measure in $\Delta(\mathcal{X} \times \mathcal{Y})$, denoted by $P_Z$, such that for all measurable subsets $A$ of $\mathcal{X} \times \mathcal{Y}$,

$$ P_Z(A) = \frac{1}{n} \sum_{t=1}^{n} 1_{\{(x_t, y_t) \in A\}}. \quad (9) $$
The empirical risk $L(z, \theta)$ in (5) then satisfies $L(z, \theta) = R_{\theta}(P_z)$. See for instance, [12, Lemma 6.1].

In statistical machine learning, a learning algorithm can be characterized by a conditional probability measure formally defined as follows.

**Definition 2.** A conditional probability measure $P_{\Theta|Z} \in \Delta(\mathcal{M}|(\mathcal{X} \times \mathcal{Y})^n)$ is said to represent a supervised machine learning algorithm if it satisfies (i) for all $z \in (\mathcal{X} \times \mathcal{Y})^n$, the set-function $P_{\Theta|Z=z}: \mathcal{B}(\mathcal{M}) \to [0, 1]$ is a probability measure in $\Delta(\mathcal{M})$; (ii) for all $A \in \mathcal{B}(\mathcal{M})$, the function $h_A: (\mathcal{X} \times \mathcal{Y})^n \to [0, 1]$ that satisfies $h_A(z) = P_{\Theta|Z=z}(A)$ is measurable with respect to the measurable spaces $((\mathcal{X} \times \mathcal{Y})^n, \mathcal{L}_n)$ and $([0, 1], \mathcal{B}([0, 1]))$, where $\mathcal{L}_n$ is a $\sigma$-algebra on $(\mathcal{X} \times \mathcal{Y})^n$.

Recall that the unique solution to the following ERM-RER (Empirical Risk Minimization with Relative Entropy Regularization) problem [1]:

$$ \min_{P \in \Delta_Q(\mathcal{M})} \int L(z, \theta) dP(\theta) + \lambda D(P \| Q) \quad (10) $$

is a Gibbs probability measure that satisfies

$$ \frac{dP_{\Theta|Z=z}^{(Q, \lambda)}}{dQ}(\theta) = \frac{\exp\left(\frac{-1}{\lambda}L(z, \theta)\right)}{\int \exp\left(\frac{-1}{\lambda}L(z, \nu)\right)dQ(\nu)}, \quad (11) $$

where the probability measure $Q \in \Delta(\mathcal{M})$ is a reference measure.

# 3 Main Results

## 3.1 A Robust Learner

The solution to the robust optimization problem posed in (1) is presented analytically in Theorem 1, where the choice of parameters is determined by the sets

$$ \mathcal{J}_{P_S, \theta} \triangleq \left\{ t \in (0, +\infty) : \log \left( \int \exp \left( \frac{1}{t} \ell(\theta, x, y) \right) dP_S(x, y) \right) < +\infty \right\}; \quad (12) $$

and

$$ \mathcal{K}_{Q,z} \triangleq \left\{ s \in (0, +\infty) : \log \left( \int \exp \left( \frac{1}{s} L(z, \theta) \right) dQ(\theta) \right) < +\infty \right\}. \quad (13) $$

The following assumption is adopted for deriving a closed-form solution for (1).

**Assumption 1.** The set of models $\mathcal{M}$ is chosen such that given a probability measure $P_S \in \Delta(\mathcal{X} \times \mathcal{Y})$, for all $\theta \in \mathcal{M}$, the set $\mathcal{J}_{P_S, \theta}$ in (12) is not empty.
The assumption is similar to the one imposed in [14] for studying a distributionally robust optimization problem with KL-divergence neighbourhood constraint. However, the minimization problem in [14] is over a vector space, which is different from the setting in this work. In fact, the condition in Assumption 1 is always satisfied in the case in which the function $\ell$ is bounded almost surely with respect to $P_S$ as shown in [12, Example 1] and [1, Lemma 1].

Under Assumption 1, the following theorem shows that the minimax optimization problem in (1) has a unique solution and is described by a Gibbs distribution.

**Theorem 1.** The minimax optimization problem in (1) has a unique solution, denoted by $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \in \Delta_Q(\mathcal{M})$, which satisfies for all $\theta \in \mathcal{M}$,

$$ \frac{d\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}}{dQ}(\theta) = \frac{\exp\left(\frac{-1}{\eta(z)}\left(\mathbf{R}_{\theta}\left(P_{\hat{Z}|\Theta=\theta}^{(P_S,\beta(\theta))}\right) - \mathbf{L}(z,\theta)\right)\right)}{\int \exp\left(\frac{-1}{\eta(z)}\left(\mathbf{R}_{\theta}\left(P_{\hat{Z}|\Theta=\nu}^{(P_S,\beta(\nu))}\right) - \mathbf{L}(z,\nu)\right)\right)dQ(\nu)}, \quad (14) $$

where for all $(x, y) \in \mathcal{X} \times \mathcal{Y}$, it holds that

$$ \frac{dP_{\hat{Z}|\Theta=\theta}^{(P_S,\beta(\theta))}}{dP_S}(x, y) = \frac{\exp\left(\frac{1}{\beta(\theta)}\ell(\theta, x, y)\right)}{\int \exp\left(\frac{1}{\beta(\theta)}\ell(\theta, x, y)\right)dP_S(x, y)}, \quad (15) $$

where the function $\beta: \mathcal{M} \to \mathbb{R}^+$ is such that $D(P_{\hat{Z}|\Theta=\theta}^{(P_S,\beta(\theta))} \| P_S) = \gamma_2$; the function $\eta: (\mathcal{X} \times \mathcal{Y})^n \to \mathbb{R}$ is such that $D(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| Q) = \gamma_1$; the function $\mathbf{L}$ is defined in (5) and the functional $\mathbf{R}_{\theta}$ is defined in (6).

*Proof:* The proof is presented in Appendix C.

<figure><img src="image_5.png" /></figure>

It is worth noting that despite the fact that the formulation in (1) aims to provide a robust learning with respect to variations in the data-generating probability measure, the solution $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ in (14) depends on the choice of the dataset $z$. Moreover, the solution $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ is parameterized by $P_S, Q, \gamma_1$ and $\gamma_2$, which jointly provide a robust learning framework that tradeoffs inductive bias via $Q$ in (1b) and robustness to the data distribution shift via $\mathcal{P}$ in (1c). The probability measure $P_{\hat{Z}|\Theta=\theta}^{(P_S,\beta(\theta))}$ in (15) is known as the Worst-Case Data-Generating (WCDG) probability measure [9, 12] (Appendix B).

The choice of $P_S$ incorporates prior information about the underlying data-generating distribution. For example, $P_S$ can be chosen as the empirical measure (Definition 1) induced by a dataset sample. When the sample size is large enough, it might be expected that the underlying data-generating distribution is within the neighbourhood of $P_S$ ($\mathcal{P}$ in (1c)).

The measure $Q$ represents the prior information of the inductive bias over the models, or (ii) the posterior produced by an existing learning algorithm whose robustness one wishes to improve. Accordingly, the parameters $\gamma_1$ and $\gamma_2$ balance inductive bias and robustness for the problem at hand. Further performance guarantees for $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ are presented in the next section.
Remarkably, $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ is itself a Gibbs probability measure, ensuring convenient sampling and regularity properties [1, 25, 26]. The practical impact of this observation is further corroborated by the fact that long-run distribution of SGD resembles a Gibbs distribution with temperature equal to the method's step-size and energy function determined by the problem's objective, see for instance [27] and [28].

## 3.2 Performance Analysis

When $Q$ is chosen as a learning algorithm, the probability measure $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ in (14) can be interpreted as a robust learner that improves upon the performance of $Q$ by adapting to new data generated in the neighborhood of $P_S$ as in (1c). We analyze the robust learner's benefits via the expected-loss improvement, PAC-Bayes guarantees, and information-theoretic decompositions in this section.

By the minimax formulation, the robust learner $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ in (14) minimizes the worst-case generalization gap. Therefore, for all data-generating distributions within the KL-neighbourhood of $P_S$ ($P$ in ((1c))), the robust learner guarantees a smaller worst-case generalization gap than the baseline learner $Q$. Moreover, when the ground-truth distribution $\tilde{P}_Z$ lies in the neighbourhood $P$, the robust learner also guarantees a smaller standard generalization error in (22) w.r.t $Q$. Moreover, a natural choice of $Q$ is the empirical risk minimizer of $L(z, \theta)$, in this case, the robust learner provides distributionally robust guarantee for test performance on $P$, which is consistent with the improvement on testing data in the numerical experiments in Section 4.

More importantly, since we have the flexibility on choosing $P_S$ and $\gamma_2$ for the neighbourhood $P$, the unknown data-generating distribution is expected to be within $P$ as stated earlier. Hence, the robust learner can be designed to improve the performance of the baseline learner $Q$ against adversarial data shifts. On the other hand, without any constraints on the choices of the neighbourhoods $P$ and $Q$, when the robust learner is observed to empirically out-perform the baseline learner $Q$ in expected loss, we prove that such an improvement is lower bounded by a Jeffreys divergence. This complementary result is formally stated in the following theorem.

**Theorem 2.** Given a testing dataset $\hat{z} \in (\mathcal{X} \times \mathcal{Y})^n$, if the robust learner $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ in (14) satisfies $R_{\hat{z}}(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}) \le \hat{M}$, where the functional $R_{\hat{z}}$ is defined in (7) and $\hat{M} \in [\inf_{\theta \in \mathcal{M}} L(\hat{z}, \theta), \int L(\hat{z}, \theta)dQ(\theta)]$, then

$$ R_{\hat{z}}(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}) - R_{\hat{z}}(Q) \\ \le -\frac{1}{\lambda} \left( D(\hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})} \| Q) + D(Q \| \hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})}) \right), \quad (16) $$

where

$$ \frac{d\hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})}}{dQ}(\theta) = \frac{\exp(-\lambda L(\hat{z}, \theta))}{\int \exp(-\lambda L(\hat{z}, \mu))dQ(\mu)}; \quad (17) $$
and $\lambda$ is chosen such that $\int L(\hat{z}, \theta) d\hat{Q}_{\Theta|Z=\hat{z}}^{(Q, \hat{M})}(\theta) = \hat{M}$.

*Proof:* The proof follows from applying an analogue of Pythagora's theorems for KL-divergence [29, Theorem 2.2] by noticing that the probability measure $\hat{Q}_{\Theta|Z=\hat{z}}^{(Q, \hat{M})}$ in (17) is the I-projection of $Q$ on the convex set $\{P \in \Delta(M) : \int L(\hat{z}, \theta) dP(\theta) \le \hat{M}\}$. The complete proof is presented in Appendix D. ■

Next, we present a training-data dependent PAC-Bayes bound for the robust learning algorithm $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$. To that end, we first introduce the log-partition function below, which plays a central role in the description of the result.

**Definition 3.** Let the log-partition function of the Gibbs probability measure $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ in (14) be $\tilde{\mathcal{K}}_{Q,z}: \mathbb{R} \to \mathbb{R} \cup \{+\infty\}$ such that

$$ \tilde{\mathcal{K}}_{Q,z}(t) = \log\left(\int \exp\left(t\left(\mathbb{R}_{\theta}(P_{\hat{z}|\Theta=\theta}^{(P_S, \beta(\theta))}) - \mathbb{L}(z, \theta)\right)\right) dQ(\theta)\right). \quad (18) $$

**Lemma 1.** Consider the robust algorithm $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ in (14), let $P_S \in \Delta(\mathcal{X} \times \mathcal{Y})^n$ be the product measure formed by $P_S$. For all $\delta \in (0, 1)$, with probability at least $1-\delta$ over the random draw of a pair $(\hat{z}, \theta) \sim P_S \otimes \tilde{P}_{\Theta|Z=z}^{(P_S, Q)} \in \Delta((\mathcal{X} \times \mathcal{Y})^n \times \mathcal{M})$, it holds that

$$ \mathbb{L}(\hat{z}, \theta) \le -\frac{\beta(\theta)}{\eta(z)} \left( D(P_z \| P_{\hat{z}|\Theta=\theta}^{(P_S, \beta(\theta))}) - D(P_z \| P_S) + \gamma_2 \right) \\ - \tilde{\mathcal{K}}_{Q,z} \left( \frac{-1}{\eta(z)} \right) + \log \left( \frac{A}{\delta} \right), \quad (19) $$

where $A = \int \int \exp(\mathbb{L}(\hat{z}, \theta)) dQ(\theta) dP_S(\hat{z})$, the probability measure $P_z$ is the type induced by training dataset $z$ defined in (9); the probability measure $P_{\hat{z}|\Theta=\theta}^{(P_S, \beta(\theta))}$ is defined in (15); and the function $\tilde{\mathcal{K}}_{Q,z}$ is the log-partition function defined in (18).

*Proof:* The proof is presented in Appendix E. The intermediate step of deriving (19) has been proved in [30, Theorem 1]. ■

This is a training-dependent PAC-Bayes inequality that characterizes testing error through $P_S$ and $\gamma_2$. Note that the data-generating measure $P_S$ need not coincide with the distribution from which the training dataset $z$ is sampled. Hence, this framework applies to both in-distribution and out-of-distribution settings. Since $P_{\hat{z}|\Theta=\theta}^{(P_S, \beta(\theta))}$ is the worst data-generating distribution for $P_S$, when the testing distribution $\tilde{P}_Z$ is absolutely continuous with respect to $P_S$, for all learning algorithms $P \in \Delta(M)$, we provide an upper bound on the generalization error $\overline{\mathcal{G}}(P_{\Theta|Z}, \tilde{P}_Z, P_Z)$ in (8) given by

$$ \int \left( \int \ell(\theta, x, y) dP_{\hat{z}|\Theta=\theta}^{(P_S, \beta(\theta))}(x, y) - \mathbb{L}(z, \theta) \right) dP(\theta) \ge \overline{\mathcal{G}}(P_{\Theta|Z}, \tilde{P}_Z, P_Z). \quad (20) $$

The following theorem connects the lower bound on the choice of $\gamma_1$ in (1b) with the upper bound of the training error induced by the robust learner.
**Theorem 3.** For a given training dataset $z$ in (3), if the learning algorithm $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ in (14) satisfies $R_z(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}) \le M$, for some $M > 0$, then the choice of radius $\gamma_1$ in (1b) must satisfy $\gamma_1 \ge D(\hat{Q}_{\Theta|Z=z}^{(Q,M)} \|Q\})$, where the probability measure $\hat{Q}_{\Theta|Z=z}^{(Q,M)}$ satisfies

$$ \frac{d\hat{Q}_{\Theta|Z=z}^{(Q,M)}}{dQ}(\theta) = \frac{\exp(-\lambda L(z, \theta))}{\int \exp(-\lambda L(z, \mu))dQ(\mu)}, \quad (21) $$

and $\lambda$ is chosen such that $\int L(z, \theta) d\hat{Q}_{\Theta|Z=z}^{(Q,M)}(\theta) = M$.

*Proof:* The proof is presented in Appendix F. ■

**Corollary 1.** The generalization error defined in (8) of the robust learner satisfies

$$ \bar{\bar{G}}(\tilde{P}_{\Theta|Z}^{(P_S,Q)}, P_Z, P_Z) = I(\tilde{P}_{\Theta|Z}^{(P_S,Q)}; P_Z) + L(\tilde{P}_{\Theta|Z}^{(P_S,Q)}; P_Z). \quad (22) $$

Moreover, for all datasets $\hat{z}$ and for all learning algorithms $P \in \Delta_Q(\mathcal{M})$, it holds that

$$ \begin{aligned} R_{\hat{z}}(P) = R_{\hat{z}}(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}) &+ \lambda \left( D(P \| P_{\Theta|Z=\hat{z}}^{(Q,\lambda)}) - D(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| Q) + D(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| Q) - D(P \| Q) \right), \end{aligned} \quad (23) $$

where the Gibbs measure $P_{\Theta|Z=\hat{z}}^{(Q,\lambda)}$ is defined in (11).

*Proof:* The generalization error expression in (22) follows from [31]. The expression in (23) follows from [1, Theorem 37]. ■

## 3.3 Connections to Zero-Sum Games

The robust learning algorithm $\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$ in (14) admits a two-player zero-sum game interpretation. This observation is characterized in the following theorem.

**Theorem 4.** Assume the set $\mathcal{M} \subset \mathbb{R}^d$ is closed, the loss function $\ell$ in (4) is bounded and jointly continuous in $\mathcal{M} \times \mathcal{X} \times \mathcal{Y}$, then the minimax optimization problem in (1) satisfies

$$ \begin{aligned} & \min_{P \in \mathcal{Q}} \max_{P_Z|\Theta \in \mathcal{P}} \int \left( \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) - L(z, \theta) \right) dP(\theta) \\ & = \max_{P_Z|\Theta \in \mathcal{P}} \min_{P \in \mathcal{Q}} \int \left( \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) - L(z, \theta) \right) dP(\theta). \end{aligned} \quad (24) $$

where the ambiguity sets $\mathcal{Q}$ and $\mathcal{P}$ are defined in (1b) and (1c), respectively.

*Proof:* The proof is presented in Appendix G. ■

Additionally, the solution to the corresponding maximin DRO follows the WCDG probability measure as stated in the following theorem.
**Theorem 5.** For a given training dataset z of the form in (3); probability measures $P_S \in \Delta(\mathcal{X} \times \mathcal{Y})$, $Q \in \Delta(\mathcal{M})$; and positive reals $\gamma_1$ and $\gamma_2$, the following maximin optimization problem

$$ \max_{P_Z|\Theta \in \mathcal{P}} \min_{P \in \mathcal{Q}} \int \left( \int l(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) - L(z, \theta) \right) dP(\theta) \quad (25) $$

has a unique solution denoted by $P_{\tilde{Z}|\Theta}^{(P_S, \beta(\cdot))} \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})$, which satisfies (15). The sets $Q$ and $\mathcal{P}$ are defined in (1b) and (1c), respectively.

*Proof:* The proof is presented in Appendix H.

<figure><img src="image_6.png" /></figure>

The problem in (24) can be viewed as a two-player zero-sum game between the learner (the leader) and nature (the follower). The learner plays first by selecting a learning algorithm from the action set $Q$ in (1b) and nature selects the Worst-Case Data-Generating probability measure in (15) from the action set $\mathcal{P}$ in (1c) for this particular learning algorithm. Since the two players choose their actions to minimize and maximize their payoff, respectively, the game is zero-sum: the “loss” of the learner is the “gain” of nature. The actions played by both players are probability measures. From Section 3.1, it follows that both players are committed to always choosing the same action, i.e. the action of each player is a single and fixed strategy. Remarkably, $(\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}, P_{\tilde{Z}|\Theta}^{(P_S, \beta(\cdot))})$ is a Nash equilibrium, as introduced in [32]. The existence of such a Nash equilibrium implies that neither the learner nor nature can unilaterally deviate from their equilibrium strategy to improve their respective objectives (i.e., reduce or increase the generalization error). In a Bayesian or mixed-strategy framework, the strategy of the learner corresponds to selecting a probability measure over the model $\theta$, rather than a single fixed parameter vector. In a learning setting, this Nash equilibrium implies that if the adversarial probability measure of nature $P_{\tilde{Z}|\Theta}^{(P_S, \beta(\cdot))}$ is fixed, then no unilateral reallocation of probability mass among the parameters of the learner can reduce its generalization error beyond that attained at the equilibrium point. Conversely, if the probability measure of the learner $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ is fixed, then nature cannot exploit it further by adapting its probability measure for the data-generating mechanism. Hence, at equilibrium, neither player gains by unilaterally changing their strategy, which frames the stability and robustness of the solution in a Bayesian or mixed-strategy learning context. Interestingly, such an equilibrium is also a Stackelberg equilibrium [33] since one player acting first does not alter the equilibrium outcome of the game.

# 4 Numerical Experiment

In this section, we provide numerical evidence supporting the analysis of the robust learner $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ in (14). Importantly, no additional hyperparameter tuning is needed for the robust learner to show improved performance, for example, $\gamma_1$ in (1b) is chosen from the beginning and kept fixed for all corrupted testing data.
**MNIST.** We evaluate binary classification on MNIST between the digits 6 and 7 as in [34]. The detailed experiment settings are presented in Section I. First, we compare the test performance of the robust learner $\tilde{P}_{\Theta|Z=z}^{(P_S, P_{\Theta|Z=z_1}^{(Q,\lambda)})}$ with that of the ERM-RER minimizer $P_{\Theta|Z=z_1}^{(Q,\lambda)}$ defined in (11), which is also used as the prior for the robust learner. The datasets $z$ and $z_1$ are training samples that may or may not coincide. In this first comparison (Figure 1a), both $P_{\Theta|Z=z_1}^{(Q,\lambda)}$ and $\tilde{P}_{\Theta|Z=z_1}^{(P_S, P_{\Theta|Z=z_1}^{(Q,\lambda)})}$ are based on the same training dataset $z_1$. In Figure 1b, we compare $\tilde{P}_{\Theta|Z=z_1}^{(P_S, Q)}$ and $P_{\Theta|Z=z_1}^{(Q, \lambda)}$, where both learning algorithms use the counting measure as the reference $Q$. All probability measures over data are estimated in the 2D space using Gaussian Kernel Density Estimation (KDE) with fixed bandwidth [35]. The test-side reference $P_S$ for the robust learner is approximated by a KDE fitted on a random subsample of the *clean* testing data, which encodes prior knowledge about the target (testing) distribution. We create *noisy* testing data by adding i.i.d. Gaussian noise to the 2D testing data $\mathbf{X}$ as follows, $\mathbf{X}^{\text{noisy}} = \mathbf{X} + \epsilon$, $\epsilon \sim \mathcal{N}(0, \sigma^2 I)$. For each noise level $\sigma$, we form a testing dataset $z_2$ uniformly sampled from $\mathbf{X}^{\text{noisy}}$. We estimated the KL divergence between the noisy and clean testing densities, denoted by $D(P_{\text{noisy test}}||P_{\text{test}})$. Importantly, for each $\sigma$, let $\gamma_2 = D(P_{\text{noisy test}}||P_{\text{test}})$. The simulation is repeated 50 times, where the datasets $z_1$ and $z_2$ are uniformly resampled at each iteration. We present the averaged expected testing error across runs from high-SNR region to low-SNR region as $\sigma$ grows. When the reference $Q$ is simply the uniform (discrete) distribution, Figure 1b shows there exists a minimum for the average expected testing error induced by the robust learner.
<figure><img src="image_7.png" /></figure>

(a) Average of $R_{z_2}(P_{\Theta|Z=z_1}^{(Q,\lambda)})$ and $R_{z_2}(\tilde{P}_{\Theta|Z=z_1}^{(P_S,P_{\Theta|Z=z_1}^{(Q,\lambda)})})$, where the functional $R_{z_2}$ is defined in (7).

<figure><img src="image_8.png" /></figure>

(b) Average of $R_{z_2}(P_{\Theta|Z=z_1}^{(Q,\lambda)})$ and $R_{z_2}(\tilde{P}_{\Theta|Z=z_1}^{(P_S,Q)})$, where the functional $R_{z_2}$ is defined in (7).

<caption>Figure 1: Average Testing Error: average of the expected empirical risks on testing datasets computed over fifty different random selections of training dataset $z_1$ and testing dataset $z_2$.</caption>

Figure 1a shows that the robust learner consistently achieves lower average testing error compared with the baseline ERM-RER Gibbs in most SNR regions, which matches the theoretical guarantee in Theorem 2. The reason behind is a concentration in the distribution over $\theta$ caused by the robust learner. Specifically, for all $\theta \in \mathcal{M}$, recall from (15), $\beta(\theta) \in \mathbb{R}$ is chosen such that $D(\tilde{P}_{\Theta|\theta}^{(P_S,\beta(\theta))} || P_S) = \gamma_2$ and here $\gamma_2$ increases with $\sigma$. In [9, Lemma 6], for a
fixed $\theta \in \mathcal{M}$, it was proved that

$$ \frac{d}{d\beta} D(P_{\hat{Z}|\Theta=\theta}^{(P_S, \beta)}|P_S) = \frac{1}{\beta} \frac{d}{d\beta} R_{\theta}(P_{\hat{Z}|\Theta=\theta}^{(P_S, \beta)}) \le 0, \quad (26) $$

which shows $R_{\theta}(P_{\hat{Z}|\Theta=\theta}^{(P_S, \beta)}) = O(\gamma_2)$ decreases with $\gamma_2$. Equation (14) shows that the posterior weights for the models assigned by $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ decay exponentially with respect to the increase in $\gamma_2$, which corresponds to the increase in domain shift. This is observed from the heatmaps in figure 2 in which $\theta$ concentrates as SNR decreases ($\gamma_2$ increases). Moreover, from the constraint in (12), in practice for models that such a $\beta(\theta)$ does not exist, the robust learner excludes them from the support. In summary, the robust learner assigns weight to the model based on its sensitivity to domain shift, its posterior concentrates exponentially on the set of models stable to perturbations around $P_S$.

The performance in low SNR could due to the under-adaptation from the fixed $\gamma_1$. Future work will focus on adapting $\gamma_1$ under heavy shift, trying different loss functions and hyperparameter tuning on $\gamma_2$ for large drifts.

<figure><img src="image_9.png" /><figcaption>Figure 2: Heatmap $\theta$ for MNIST.</figcaption></figure>

**CIFAR10.** We evaluate the robust learner $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ on CIFAR-10 and CIFAR-10-C [36], which applies 19 corruption types at five severity levels. Each corruption family has a deterministic scale parameter per severity. At a high level, the learner forms a posterior $\tilde{P}_{\Theta|Z=z}^{(P_S, Q)}$ over last-layer heads that (i) stay $\gamma_1$ close to a baseline distribution $Q$ over last layer heads and (ii) is robust to the worst-case testing data distribution in a KL ball of radius $\gamma_2$ centered at a test-side prior $P_S$. We train a standard ResNet-18 on CIFAR-10 and freeze the backbone after training. Let $x \in \mathbb{R}^d$ denote the penultimate features and let the classifier be a linear head ($W, b_0$) with $W \in \mathbb{R}^{d \times C}$, $C = 10$. To obtain a low-dimensional model subspace, we collect $T$ snapshots of the last layer weights $\{W_t\}_{t=1}^T$ during training. We compute the mean $W_0 = \frac{1}{T} \sum_{t=1}^T W_t$ and build
<figure><img src="image_10.png" /><figcaption>Figure 3: Robust Learner, Baseline Gibbs and Original Resnet-18 on CIFAR-10-C.</figcaption></figure>
<caption>**Algorithm 1 Robust Last-Layer Mixture**</caption>

**Require:** Prior $Q \in \Delta(\mathcal{M})$ over $\{\theta_g\}_{g=1}^G$, $\theta_g \in \mathbb{R}^K$, test-side reference $P_S \in \Delta(\mathcal{X} \times \mathcal{Y})$, radius $\gamma_1, \gamma_2 > 0$, basis $B \in \mathbb{R}^{K \times dC}$, mean $\mu = \text{vec}(W_0) \in \mathbb{R}^{dC}$, bias $b_0$, source split $(X, y)$ for empirical loss, sample set $S = \{(x_s, y_s)\}$ to estimate $P_S$.

**Ensure:** Class probabilities $\bar{p}(x)$ and predictions $\hat{y}$

1: Initialize mixture weights $w_g \propto Q_g \quad \triangleright$ start from baseline prior

2: **for** $g = 1$ to $G$ **do**

3: $W_g \leftarrow \text{unvec}(\mu + B^\top \theta_g)$ \quad $\triangleright$ reshape last-layer weights back to $d \times C$

4: $Z_g^{\text{src}} \leftarrow X W_g + b_0$; $p_g^{\text{src}} \leftarrow \text{softmax}(Z_g^{\text{src}})$

5: $\ell_g^{\text{emp}} \leftarrow \text{CE}(p_g^{\text{src}}, y)$ \quad $\triangleright$ empirical CE on source split

6: $\ell_g^{\text{worst}}, \beta_g \leftarrow \text{WCDG-Loss}(\theta_g, S, \gamma_2)$

7: $s_g \leftarrow \ell_g^{\text{emp}} - \ell_g^{\text{worst}}$ \quad $\triangleright$ $s_g = (L(z, \theta_g) - R_{\theta_g}(\hat{P}_{WCDG})) \text{ matches (14)}$

8: **end for**

9: $\eta \leftarrow \text{SOLVEETA}(Q, \{s_g\}_{g=1}^G, \gamma_1)$ \quad $\triangleright$ choose $\eta$ such that $D(\hat{P}\|Q) = \gamma_1$

10: $w_g \propto Q_g \exp(s_g/\eta)$; normalize $\sum_g w_g = 1$

11: **return** $\bar{p}(\cdot) \leftarrow \sum_g w_g \text{softmax}(\cdot W_g + b_0)$; $\hat{y} \leftarrow \text{arg max}_c \bar{p}_c$

a orthonormal basis $B \in \mathbb{R}^{d \times K}$ (via SVD on $\{W_t - W_0\}$), then, the parameterized last-layer weights are now $W(\theta) = W_0 + B\theta$, $\theta \in \mathbb{R}^K$. The subspace of $\theta$ is constructed by samples drawn from a Gaussian distribution $N(0, \tau^2 I)$. Let baseline $Q$ follow the Gibbs distribution defined in (11), where $L(z, \theta)$ is the empirical cross-entropy on a clean source split. The reference $P_S$ is chosen from the testing data of CIFAR-10 and estimated via KDE. For each CIFAR-10-C corruption and severity, we estimate the KL-divergence $D(P_{\text{corrupt}}\|P_S)$ and let $\gamma_2 = D(P_{\text{corrupt}}\|P_S)$, where $P_{\text{corrupt}}$ is estimated via KDE on corresponding corrupted samples. Finally, the robust learner assigns weights to $\theta$ following (14). We report top-1 accuracy, negative log-likelihood (NLL) and expected calibration error (ECE) on CIFAR-10-C comparing the original ResNet-18, $Q$ baseline mixture ResNet-18 and robust learner mixture ResNet-18. This simple implementation of the robust learner is summarized in the Algorithm above, where the robust learner is denoted by $\tilde{P}$ and $P^{(P_S, \beta(\theta))}$ is denoted by $\hat{P}_{WCDG}$. Figure 3 shows the robust learner yields consistently lower expected calibration error and lower negative log-likelihood while preserving accuracy. The robust learner is proved to reduce over-confidence on corrupted inputs and produces probability estimates that stay better calibrated under corruptions. Without fine-tuning the weights and re-training of the neural network, the matching accuracy is expected and the robust learner successfully improve reliability under distribution shifts. Nonetheless, in future work, it is possible to boost the accuracy by tuning $\gamma_1, \gamma_2$ and involving second-to-last layer mixture to enrich $\theta$ subspace.
# 5 Conclusion and Future Work

In this work, we introduce a robust learner as the unique analytical solution to a minimax problem with double KL-divergence constraints. The robust learner follows a Gibbs distribution and minimizes the worst-case generalization gap over a KL-divergence neighbourhood of the target test distribution and can achieve a smaller worst-case generalization error than the baseline learning algorithm $Q$. Under suitable conditions, it further guarantees a smaller expected test error than $Q$. A training-dependent PAC-Bayes guarantee and characterizations of performance metrics such as generalization error are provided in terms of mutual information, lautum information and KL-divergence. Furthermore, the proposed minimax problem admits a two-player zero-sum game interpretation along with a unique Nash equilibrium. Finally, we validate our theoretical results on MNIST and CIFAR-10-C. We introduce a simple algorithm on neural network showing only a last-layer mixture by the robust learner strictly improves reliability and preserves accuracy under distribution shift without re-training and fine-tuning. For future work, the robust learner incorporates high flexibility through its hyperparameters, which are not fully analyzed and optimized.
# Appendices

## A Preliminary Results

**Definition 4.** Given a model $\theta \in \mathcal{M}$, the function $\ell$ in (4), is said to be separable with respect to the probability measure $P \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) (\mathcal{X} \times \mathcal{Y})$, if there exist a positive real $c > 0$ and two subsets $\mathcal{A}$ and $\mathcal{B}$ of $\mathcal{X} \times \mathcal{Y}$ that are nonnegligible with respect to $P$, and for all $((x_1, y_1), (x_2, y_2)) \in \mathcal{A} \times \mathcal{B}$,

$$ \ell(\theta, x_1, y_1) < c < \ell(\theta, x_2, y_2) < +\infty. $$

**Definition 5 (Kullback-Leibler divergence).** Given two $\sigma$-finite measures $P_1$ and $P_2$ on the same measurable space, such that $P_2$ is absolutely continuous with respect to $P_1$, the relative entropy of $P_2$ with respect to $P_1$ is

$$ D(P_2\|P_1) = \int \frac{dP_2}{dP_1}(x) \log \left( \frac{dP_2}{dP_1}(x) \right) dP_1(x), \quad (27) $$

where the function $\frac{dP_2}{dP_1}$ is the Radon-Nikodym derivative of $P_2$ with respect to $P_1$.

**Lemma 2.** For all probability measures $P \in \Delta(\mathcal{M})$, it holds that

$$ \begin{aligned} & \max_{P_{Z|\Theta} \in \mathcal{P}} \int \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) dP(\theta) \\ & = \int \max_{P_{Z|\Theta=\theta} \in \mathcal{P}_{\theta}} \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) dP(\theta), \end{aligned} \quad (28) $$

where the sets are $\mathcal{P} = \{P_{Z|\Theta} \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) : \forall \theta, D(P_{Z|\Theta=\theta}\|P_S) \le \gamma_2\}$; and for all $\theta \in \mathcal{M}$, $\mathcal{P}_{\theta} = \{P_{Z|\Theta=\theta} \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y}) : D(P_{Z|\Theta=\theta}\|P_S) \le \gamma_2\}$.

*Proof:* Let

$$ f(\theta, P_{Z|\Theta=\theta}) \triangleq \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}, \quad (29) $$

$$ g(\theta) \triangleq \max_{P_{Z|\Theta=\theta} \in \mathcal{P}_{\theta}} f(\theta, P_{Z|\Theta=\theta}). \quad (30) $$

Followed from [9, Lemma 3, Lemma 5], $g(\theta)$ is a nonnegative function. Moreover, followed from [37, Proposition 2.1.5], $g(\theta)$ is Borel measurable. Therefore, by the definition of the integral [38, Definition 1.5.3], it holds that

$$ \begin{aligned} & \int g(\theta) dP(\theta) \\ & = \sup\{\int s(\theta) dP(\theta) : s \text{ is a simple function and } 0 \le s(\theta) \le g(\theta)\}. \end{aligned} \quad (31) $$

The simple function $s$ is given by

$$ s(\theta) = \sum_{i=1}^{r} x_i \mathbb{1}_{\{\theta \in A_i\}}, \quad (32) $$
where $x_i > 0$ and $A_i, \dots, A_r$ are disjoint measurable sets in $\mathcal{M}$. Since for all $\theta \in \mathcal{M}$, $s(\theta) \le g(\theta)$, it holds that for $i = 1, \dots, r$,

$$x_i \le \inf_{\theta \in A_i} g(\theta) = \inf_{\theta \in A_i} \max_{P_{Z|\Theta=\theta} \in \mathcal{P}_{\theta}} f(\theta, P_{Z|\Theta=\theta}). \quad (33)$$

By Theorem 6, for each $\theta \in \mathcal{M}$ there is a unique solution to

$$\max_{P_{Z|\Theta=\theta} \in \mathcal{P}_{\theta}} f(\theta, P_{Z|\Theta=\theta}).$$

Hence, for each $A_i$, $i = 1, \dots, r$, there exists $\theta^* \in A_i$, $f(\theta^*, P_{Z|\Theta=\theta^*}) \ge x_i$. Therefore,

$$\int f(\theta, P_{Z|\Theta=\theta}) dP(\theta) = \sum_{i=1}^{r} \int_{A_i} f(\theta^*, P_{Z|\Theta=\theta^*}) dP(\theta^*) \quad (34)$$

$$\ge \sum_{i=1}^{r} \int_{A_i} x_i dP(\theta^*) \quad (35)$$

$$= \int s(\theta) dP(\theta). \quad (36)$$

Thus,

$$\max_{P_{Z|\Theta} \in \mathcal{P}} \int f(\theta, P_{Z|\Theta=\theta}) dP(\theta) \\ \ge \max_{P_{Z|\Theta} \in \mathcal{P}} \int s(\theta) dP(\theta) \quad (37)$$

$$= \int g(\theta) dP(\theta). \quad (38)$$

On the otherhand, for all $P_{Z|\Theta=\theta} \in \mathcal{P}_{\theta}$, it holds that

$$\int f(\theta, P_{Z|\Theta=\theta}) dP(\theta) \le \int g(\theta) dP(\theta). \quad (39)$$

Hence,

$$\max_{P_{Z|\Theta} \in \mathcal{P}} \int f(\theta, P_{Z|\Theta=\theta}) dP(\theta) \le \int g(\theta) dP(\theta). \quad (40)$$

Combining (38) and (40) leads to

$$\int \max_{P_{Z|\Theta=\theta} \in \mathcal{P}_{\theta}} f(\theta, P_{Z|\Theta=\theta}) dP(\theta) = \max_{P_{Z|\Theta} \in \mathcal{P}} \int f(\theta, P_{Z|\Theta=\theta}) dP(\theta), \quad (41)$$

which proves the two maximization problems have the same optimal value. ■

**Lemma 3.** For a given training dataset $z$ in (3) and a probability measure $Q \in \Delta(\mathcal{M})$, consider the following minimization problem

$$\min_{\hat{Q} \in \Delta_Q(\mathcal{M})} D(\hat{Q} \| Q) \quad (42a)$$

$$\text{s. t.} \quad \int \mathbf{L}(z, \theta) d\hat{Q}(\theta) \le M, \quad (42b)$$
where $M \in [\inf_{\theta \in \mathcal{M}} L(z, \theta), \int L(z, \theta)dQ(\theta)]$. The unique solution to (42) denoted by $\hat{Q}_{\Theta|Z=z}^{(Q, \lambda)}$ satisfies

$$ \frac{d\hat{Q}_{\Theta|Z=z}^{(Q,M)}}{dQ}(\theta) = \frac{\exp(-\lambda L(z, \theta))}{\int \exp(-\lambda L(z, \theta))dQ(\theta)}, \quad (43) $$

where $\lambda$ is chosen such that $\int L(z, \theta) d\hat{Q}_{\Theta|Z=z}^{(Q,M)}(\theta) = M$.

*Proof:* The proof of this lemma and the existence of $\lambda$ follow the same technique as the proof for [1, Theorem 3]. The uniqueness of the solution follows from the strict convexity of the KL-divergence [1, Theorem 2]. ■

## B The Worst-Case Data-Generating Probability Measure

The Worst-Case Data-Generating (WCDG) probability measure [12], which is parameterized by a real $t > 0$ and a probability measure $P_S$ in $\Delta(\mathcal{X} \times \mathcal{Y})$, is specific for a given model $\theta$ and thus, denoted by $P_{\hat{Z}|\Theta=\theta}^{(P_S,t)}$. Note that $P_{\hat{Z}|\Theta=\theta}^{(P_S,t)}$ is a measure in $\Delta_{P_S}(\mathcal{X} \times \mathcal{Y})$. An explicit expression for $P_{\hat{Z}|\Theta=\theta}^{(P_S,t)}$ requires defining the function $J_{P_S, \theta}: \mathbb{R} \to \mathbb{R}$ such that

$$ J_{P_S, \theta}(t) = \log \left( \int \exp(t\ell(\theta, x, y)) dP_S(x, y) \right), \quad (44) $$

where the function $\ell$ is defined in (4).

The probability measure $P_{\hat{Z}|\Theta=\theta}^{(P_S,t)}$ satisfies for all $\theta \in \mathcal{M}$,

$$ \frac{dP_{\hat{Z}|\Theta=\theta}^{(P_S,t)}}{dP_S}(x, y) = \exp\left(\frac{1}{t}\ell(\theta, x, y) - J_{P_S, \theta}\left(\frac{1}{t}\right)\right), \quad (45) $$

where the function $J_{P_S, \theta}$ is defined in (44). The existence of $P_{\hat{Z}|\Theta=\theta}^{(P_S,t)}$ is ensured as long as $t \in \mathcal{J}_{P_S, \theta}$, where

$$ \mathcal{J}_{P_S, \theta} \triangleq \left\{ t \in (0, +\infty) : J_{P_S, \theta}\left(\frac{1}{t}\right) < +\infty \right\}. \quad (46) $$

Under the assumption that the loss function $\ell$ is separable (4), the WCDG probability measure has been shown to be the unique solution to two expected loss maximization problems. We restate this result in the following theorem.

**Theorem 6.** [9, Theorem 1] The WCDG probability measure $P_{\hat{Z}|\Theta=\theta}^{(P_S,t)}$ in (45) is the unique solution to the following two optimization problems,

$$ \max_{P \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y})} \int \ell(\theta, x, y) dP(x, y) - tD(P \| P_S); \quad (47) $$
and

$$ \max_{P \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y})} \int \ell(\boldsymbol{\theta}, x, y) dP(x, y) \quad (48a) $$

$$ \text{s.t.} \quad D(P \| P_S) \le \gamma, \quad (48b) $$

where $t$ is chosen such that $D(P_{\hat{Z}|\Theta=\theta}^{(P_S,t)} \| P_S) = \gamma$; and the function $\ell$ is defined in (4).

## C Proof of Theorem 1

*Proof:* Let

$$ \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}) = \{P \in \Delta_{P_S}(\mathcal{X} \times \mathcal{Y}) : D(P \| P_S) \le \gamma_2\}. \quad (49) $$

The following holds:

$$ \begin{aligned} & \min_{\tilde{P} \in \mathcal{P}} \max_{P \in \mathcal{P}} \int \left( \int \ell(\boldsymbol{\theta}, x, y) dP_{\tilde{Z}|\Theta=\theta}(x, y) - \mathbf{L}(\boldsymbol{z}, \boldsymbol{\theta}) \right) d\tilde{P}(\boldsymbol{\theta}) \\ &= \min_{\tilde{P} \in \mathcal{P}} \int \max_{P \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y})} \left( \int \ell(\boldsymbol{\theta}, x, y) dP(x, y) - \mathbf{L}(\boldsymbol{z}, \boldsymbol{\theta}) \right) d\tilde{P}(\boldsymbol{\theta}) \end{aligned} \quad (50) $$

$$ = \min_{\tilde{P} \in \mathcal{P}} \int \left( \mathbf{R}_{\boldsymbol{\theta}} \left( P_{\tilde{Z}|\Theta=\boldsymbol{\theta}}^{(P_S, \beta(\boldsymbol{\theta}))} \right) - \mathbf{L}(\boldsymbol{z}, \boldsymbol{\theta}) \right) d\tilde{P}(\boldsymbol{\theta}), \quad (51) $$

where the equality in (50) follows from Lemma 2 in Appendix A; and the equality in (51) follows from Theorem 6. We conclude the proof by solving the minimization problem in (51), which follows the same technique as the proof for [12, Theorem 1] and leads to the solution in (14). ■

## D Proof of Theorem 2

*Proof:* Consider the set of probability measures

$$ \mathcal{S}(Q, +\infty) = \{P \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) (\mathcal{M}) : D(P \| Q) < +\infty\} \quad (52) $$

and a convex set of probability measures

$$ \mathcal{I} = \{P \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) (\mathcal{M}) : \int \mathbf{L}(\hat{\boldsymbol{z}}, \boldsymbol{\theta}) dP(\boldsymbol{\theta}) \le \hat{M}\}. \quad (53) $$

By Lemma 3 and the definition of $I$-projection by [29, Equation (1.5)], the probability measure $\hat{Q}_{\Theta|Z=\hat{z}}^{(Q, \hat{M})} \in \mathcal{I} \cap \mathcal{S}(Q, +\infty)$ is the $I$-projection of $Q$ on $\mathcal{I}$. Therefore, by [29, Theorem 2.2], for all $P \in \mathcal{I}$, it holds that

$$ D(P \| Q) \ge D(P \| \hat{Q}_{\Theta|Z=\hat{z}}^{(Q, \hat{M})}) + D(\hat{Q}_{\Theta|Z=\hat{z}}^{(Q, \hat{M})} \| Q). \quad (54) $$
By [1, Theorem 37], it holds that

$$ 
\begin{aligned}
& R_{\hat{z}}(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}) - R_{\hat{z}}(Q) \\
&= \frac{1}{\lambda} \left( D\left(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| \hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})}\right) - D\left(Q \| \hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})}\right) \right. \\
& \qquad \left. - D\left(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| Q\right) \right)
\end{aligned} 
\quad (55) $$

$$ 
\begin{aligned}
&\le \frac{1}{\lambda} \left( D\left(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| Q\right) - D\left(\hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})} \| Q\right) \right. \\
& \qquad \left. - D\left(Q \| \hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})}\right) - D\left(\tilde{P}_{\Theta|Z=z}^{(P_S,Q)} \| Q\right) \right)
\end{aligned} 
\quad (56) $$

$$ = \frac{1}{\lambda} \left( - D\left(\hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})} \| Q\right) - D\left(Q \| \hat{Q}_{\Theta|Z=\hat{z}}^{(Q,\hat{M})}\right) \right), \quad (57) $$

where the inequality in (56) follows from (54). This completes the proof. ■

## E Proof of Lemma 1

By Markov's inequality, for all $\delta \in (0, 1)$, with probability at least $1 - \delta$ over the random draw of a pair $(\hat{z}, \theta) \sim P_S \otimes \tilde{P}_{\Theta|Z=z}^{(P_S,Q)}$, it holds that

$$ \log (L(\hat{z}, \theta)) \le \log \left( \iint L(\hat{z}, \theta) d\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}(\theta) dP_S(\hat{z}) \right) + \log \left( \frac{1}{\delta} \right). \quad (58) $$

It holds that

$$ 
\begin{aligned}
& \log \left( \exp \left( L(\hat{z}, \theta) - \log \left( \frac{d\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}}{dQ}(\theta) \right) \right) \right) \\
& \le \log \left( \iint \exp \left( L(\hat{z}, \theta) - \log \left( \frac{d\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}}{dQ}(\theta) \right) \right) d\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}(\theta) dP_S(\hat{z}) \right) \\
& \quad + \log \left( \frac{1}{\delta} \right),
\end{aligned} 
\quad (59) $$
which is equivalent to

$$
\begin{align*}
& L(\hat{z}, \theta) \\
&\le \log \left( \frac{d\tilde{P}_{\Theta|Z=z}^{(P_S,Q)}}{dQ}(\theta) \right) + \log \left( \iint \exp(L(\hat{z}, \theta)) dQ(\theta) dP_S(\hat{z}) \right) \\
&\qquad + \log \left( \frac{1}{\delta} \right) \\
&= -\frac{1}{\eta} (R_{\theta}(P_{\hat{z}|\Theta=\theta}^{(P_S,\beta(\theta))}) - L(z, \theta)) - \log \left( \iint \exp\left(-\frac{1}{\eta(z)}(R_{\theta}(P_{\hat{z}|\Theta=\theta}^{(P_S,\beta(\theta))}) - L(z, \theta))\right) dQ(\theta) \right) \\
&\qquad + \log \left( \iint \exp(L(\hat{z}, \theta)) dQ(\theta) dP_S(\hat{z}) \right) + \log \left( \frac{1}{\delta} \right) \\
&= -\frac{\beta(\theta)}{\eta(z)} \left( D(P_z \| P_{\hat{z}|\Theta=\theta}^{(P_S,\beta(\theta))}) - D(P_z \| P_S) + \gamma_2 \right) - \tilde{K}_{Q,z} \left( \frac{-1}{\eta(z)} \right) \\
&\qquad + \log \left( \iint \exp(L(\hat{z}, \theta)) dQ(\theta) dP_S(\hat{z}) \right) + \log \left( \frac{1}{\delta} \right) \\
&= -\frac{\beta(\theta)}{\eta(z)} \left( D(P_z \| P_{\hat{z}|\Theta=\theta}^{(P_S,\beta(\theta))}) - D(P_z \| P_S) + \gamma_2 \right) - \tilde{K}_{Q,z} \left( \frac{-1}{\eta} \right) \\
&\qquad + \log \left( \frac{A}{\delta} \right),
\end{align*}
$$

where

$$A = \iint \exp(L(\hat{z}, \theta)) dQ(\theta) dP_S(\hat{z}). \qquad (60)$$

This completes the proof.

## F Proof of Theorem 3

*Proof:* Theorem 3 is a direct result from the Lemma 3 in Appendix A. Therefore, for all $P \in \{P \in \Delta(M) : \int L(z, \theta) dP(\theta) \le M\}$, it holds that $D(P||Q) \ge D(\hat{Q}_{\Theta|Z=z}^{(Q,M)}||Q)$. This completes the proof. ■

## G Proof of Theorem 4

The proof is completed by verifying the conditions and directly applying Sion's minimax theorem. Recall Sion's minimax theorem [39] as follows,

*Corollary 3.3 in [39]*: Let M and N be convex spaces one of which is compact, and f a function on $M \times N$, quasi-concave/convex and upper semicontinuous - lower semicontinuous. Then $\sup_{f} \inf_{f} f$.

The first part of the proof consists of showing the set $Q$ in (1b) is compact. For the accuracy of the notation, denote $Q$ as $\Delta_{Q,\gamma_1}(M)$ and denote $P$ in (1c) as $\Delta_{P_S,\gamma_2}(M \times Y|M)$. Firstly, we show the set $\Delta_{Q,\gamma_1}(M)$ is relatively

RR n° 9600
compact (the set is compact on its closure), which, based on the Prokhorov's Theorem [38, Chapter 7], is equivalent to proving the set $\Delta_{Q,\gamma_1}(\mathcal{M})$ is tight.

In order to show the set $\Delta_{Q,\gamma_1}(\mathcal{M}) = \{P \in \Delta_Q(\mathcal{M}) : D(P\|Q) \le \gamma_1\}$ is tight, based on [38, Definition 7.2.3], it is equivalent to prove that for all $\epsilon > 0$, there exists a compact set $\mathcal{K} \subseteq \mathcal{M}$ such that

$$ \forall P \in \Delta_{Q,\gamma_1}(\mathcal{M}), P(\mathcal{K}) \ge 1 - \epsilon. \quad (61) $$

For all $P \in \Delta_{Q,\gamma_1}(\mathcal{M})$, let $f(\theta) \triangleq \frac{dP}{dQ}(\theta)$ and $M > e$, it holds that

$$ \int_{\{\theta \in \mathcal{M}: f(\theta) \ge M\}} \left( f(\theta) \log (f(\theta)) - f(\theta) + 1 \right) dQ(\theta) \\ \ge \int_{\{\theta \in \mathcal{M}: f(\theta) \ge M\}} \left( f(\theta) \log M - f(\theta) + 1 \right) dQ(\theta) \quad (62) $$

$$ = (\log(M) - 1) \int_{\{\theta \in \mathcal{M}: f(\theta) \ge M\}} f(\theta) dQ(\theta) + Q(\{\theta \in \mathcal{M}: f(\theta) \ge M\}) \quad (63) $$

$$ \ge (\log(M) - 1) P(\{\theta \in \mathcal{M}: f(\theta) \ge M\}). \quad (64) $$

For all $P \in \Delta_{Q,\gamma_1}(\mathcal{M})$, it holds that

$$ D(P\|Q) = \int \left( f(\theta) \log (f(\theta)) - f(\theta) + 1 \right) dQ(\theta) \le \gamma_1. \quad (65) $$

Therefore,

$$ \gamma_1 \ge \int_{\{\theta \in \mathcal{M}: f(\theta) \ge M\}} \left( f(\theta) \log (f(\theta)) - f(\theta) + 1 \right) dQ(\theta) \\ + \int_{\{\theta \in \mathcal{M}: f(\theta) < M\}} \left( f(\theta) \log (f(\theta)) - f(\theta) + 1 \right) dQ(\theta) \quad (66) $$

$$ \ge \int_{\{\theta \in \mathcal{M}: f(\theta) \ge M\}} \left( f(\theta) \log (f(\theta)) - f(\theta) + 1 \right) dQ(\theta) \quad (67) $$

$$ \ge (\log(M) - 1) P(\{\theta \in \mathcal{M}: f(\theta) \ge M\}), \quad (68) $$

where the inequality in (67) follows from $f(\theta) \log (f(\theta)) - f(\theta) + 1 \ge 0$.

Hence, for all $P \in \Delta_{Q,\gamma_1}(\mathcal{M})$,

$$ P(\{\theta \in \mathcal{M}: f(\theta) \ge M\}) \le \frac{\gamma_1}{\log(M) - 1}, \quad (69) $$

which implies that for all $M \ge \exp(\frac{\gamma_1}{\epsilon} + 1)$, $P(\{\theta \in \mathcal{M}: f(\theta) \ge M\}) \le \epsilon$. Under the assumption that $\mathcal{M} \subseteq \mathbb{R}^d$ is closed, $\mathcal{M}$ is a Polish space [37, Proposition 8.1.2]. Moreover, $Q \in \Delta(\mathcal{M})$ is a probability measure, therefore, for all $\delta > 0$, there exists a compact subset $\mathcal{K}_\delta \subseteq \mathcal{M}$ such that $Q(\mathcal{K}_\delta) \ge 1 - \delta$ [37, Proposition 8.1.12].
Note that for all $P \in \Delta_{Q,\gamma_1}(\mathcal{M})$,

$$P(\mathcal{K}_{\delta}^{c}) = \int_{\mathcal{K}_{\delta}^{c}} f(\boldsymbol{\theta})dQ((\boldsymbol{\theta})) \qquad (70)$$

$$= \int_{\{\boldsymbol{\theta} \in \mathcal{M}: f(\boldsymbol{\theta}) < M\} \cap \mathcal{K}_{\delta}^{c}} f(\boldsymbol{\theta})dQ(\boldsymbol{\theta}) + \int_{\{\boldsymbol{\theta} \in \mathcal{M}: f(\boldsymbol{\theta}) \ge M\} \cap \mathcal{K}_{\delta}^{c}} f(\boldsymbol{\theta})dQ(\boldsymbol{\theta}) \quad (71)$$

$$\le \int_{\{\boldsymbol{\theta} \in \mathcal{M}: f(\boldsymbol{\theta}) < M\} \cap \mathcal{K}_{\delta}^{c}} f(\boldsymbol{\theta})dQ(\boldsymbol{\theta}) + \int_{\{\boldsymbol{\theta} \in \mathcal{M}: f(\boldsymbol{\theta}) \ge M\}} f(\boldsymbol{\theta})dQ(\boldsymbol{\theta}) \quad (72)$$

$$\le M \cdot Q(\mathcal{K}_{\delta}^{c}) + P(\{\boldsymbol{\theta} \in \mathcal{M}: f(\boldsymbol{\theta}) \ge M\}) \quad (73)$$

$$\le M \cdot Q(\mathcal{K}_{\delta}^{c}) + \frac{\gamma_1}{\log(M) - 1}, \quad (74)$$

where the inequality in (74) follows from (69). Since the limit of RHS of (69) is 0 when $M \to +\infty$, let $\delta$ and $M$ be such that $M \cdot \delta \le \frac{\epsilon}{2}$ and $\frac{\gamma_1}{\log(M)-1} \le \frac{\epsilon}{2}$, then $P(\mathcal{K}_{\delta}^{c}) \le \epsilon$, therefore, for all $P \in \Delta_{Q,\gamma_1}(\mathcal{M})$, it holds that $P(\mathcal{K}_{\delta}) \ge 1 - \epsilon$, which completes the proof of the tightness of the set $\Delta_{Q,\gamma_1}(\mathcal{M})$ and based on the Prokhorov's Theorem [38, Chapter 7], the set $\Delta_{Q,\gamma_1}(\mathcal{M})$ is relatively compact.

The second part of the proof consists of proving the quasi-concave/convex and upper semicontinuous - lower semicontinuous properties on the space $\Delta_{P_S,\gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) \times \Delta_{Q,\gamma_1}(\mathcal{M})$. The lower semicontinuity of the map $P \to D(P||Q)$ [29, Equation (2.5)] leads to $\Delta_{Q,\gamma_1}(\mathcal{M})$ being a closed subset of $\Delta_Q(\mathcal{M})$. Since a relatively compact set that is also closed is compact, we complete the proof of compactness of the set $\Delta_{Q,\gamma_1}(\mathcal{M})$.

Furthermore, the convexity of the sets $\Delta_{P_S,\gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})$ and $\Delta_{Q,\gamma_1}(\mathcal{M})$ are verified due to the convexity of relative entropy [1, Theorem 2].

Let $F: \Delta_{P_S,\gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M}) \times \Delta_{Q,\gamma_1}(\mathcal{M}) \to \mathbb{R}$ be such that

$$F(P_{Z|\Theta}, P) = \int \left( \int \ell(\theta, x, y)dP_{Z|\Theta=\theta}(x, y) - L(z, \theta) \right) dP(\theta). \quad (75)$$

For the semicontinuity, we continue the proof by noticing that the objective function $F$ in (75) is affine in $P_{Z|\Theta}$ for fixed $P$, and linear in $P$ for fixed $P_{Z|\Theta}$. For every fixed $P \in \Delta_Q(\mathcal{M})$, for each $\theta$, the map $P_{Z|\Theta=\theta} \to \int \ell(\theta, x, y)dP_{Z|\Theta=\theta}(x, y)$ is continuous by the Portmanteau theorem for bounded continuous integrals [37]. Since $|\int \ell(\theta, x, y)dP_{Z|\Theta=\theta}(x, y)| \le \|\ell\|_{\infty}$ uniformly, dominated convergence in $\theta$ yields that the map $P_{Z|\Theta} \to F(P_{Z|\Theta}, P)$ is continuous (hence upper semicontinuous) on $\Delta_{P_S,\gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})$ [37]. For every fixed $P_{Z|\Theta} \in \Delta_{P_S,\gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})$, the term $\int \ell(\theta, x, y)dP_{Z|\Theta=\theta}(x, y) - L(z, \theta)$ is bounded and continuous. Therefore, the map $P \to F(P_{Z|\Theta}, P)$ is continuous (hence lower semicontinuous) on $\Delta_Q(\mathcal{M})$ [37].

Finally, we have verified all the conditions for Sion's minimax result to hold and hence, the proof is completed.
## H Proof of Theorem 5

*Proof:* Following the same arguments as in the proof Theorem 1, the solution to the inner minimization problem denoted by $\tilde{P}_{\Theta|Z=z}^{(Q,\lambda)} \in \Delta_Q(\mathcal{M})$, satisfies

$$ \frac{d\tilde{P}_{\Theta|Z=z}^{(Q,\lambda)}}{dQ}(\theta) = \frac{\exp\left(-\frac{1}{\lambda}\left(\int \ell(\theta, x, y)dP_{Z|\Theta=\theta}(x, y) - \mathbf{L}(z, \theta)\right)\right)}{\int \exp\left(-\frac{1}{\lambda}\left(\int \ell(\theta, x, y)dP_{Z|\Theta=\theta}(x, y) - \mathbf{L}(z, \theta)\right)\right)dQ(\theta)}, (76) $$

where a positive real $\lambda$ is chosen such that $D(\tilde{P}_{\Theta|Z=z}^{(Q,\lambda)} \| Q) = \gamma_1$. This leads to

$$ \begin{aligned} & \max_{P_{Z|\Theta} \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})} \min_{P \in \Delta_Q(\mathcal{M})} \int \left( \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) - \mathbf{L}(z, \theta) \right) dP(\theta) \\ & = \max_{P_{Z|\Theta} \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})} \int \int \ell(\theta, x, y) dP_{Z|\Theta=\theta}(x, y) d\tilde{P}_{\Theta|Z=z}^{(Q,\lambda)}. \end{aligned} \quad (77) $$

Following the same arguments as in the proof of Theorem 1, the solution to the above maximization problem denoted by $P_{Z|\Theta}^{(P_S, \beta(\theta))} \in \Delta_{P_S, \gamma_2}(\mathcal{X} \times \mathcal{Y}|\mathcal{M})$, satisfies that for all $\theta \in \text{supp } Q$,

$$ \frac{dP_{\hat{Z}|\Theta=\theta}^{(P_S, \beta(\theta))}}{dP_S}(x, y) = \frac{\exp\left(\frac{1}{\beta(\theta)}\ell(\theta, x, y)\right)}{\int \exp\left(\frac{1}{\beta(\theta)}\ell(\theta, x, y)\right)dP_S(x, y)}, \quad (78) $$

where the function $\beta: \mathcal{M} \to \mathbb{R}^+$ is such that $D(P_{\hat{Z}|\Theta=\theta}^{(P_S, \beta(\theta))} \| P_S) = \gamma_2$. The uniqueness of the solution is verified due to the strict concavity of the objective function with respect to the optimization value, which completes the proof. ■

## I Numerical Experiments

**Data and representation.** The MNIST dataset consists of 60,000 training and 10,000 testing images. For the classification between number 6 and 7, there are 12,183 available training and 1,986 testing images. Each image is 28 × 28 grayscale, represented as a matrix in $[0, 1]^{28 \times 28}$. To obtain a two-dimensional representation, we apply Uniform Manifold Approximation and Projection (UMAP) [40] as a preprocessing step.

**Settings.** We set $\mathcal{M} = [-50, 50]^2$; The set of patterns $\mathcal{X}$ is the 2D embedding from UMAP, and the labels space is $\mathcal{Y} = \{6, 7\}$. The prediction rule is the linear separator

$$ f(\theta, x) = \begin{cases} 6 & \text{if } x\theta > 0, \\ 7 & \text{if } x\theta < 0, \end{cases} \quad (79) $$

with the 0-1 loss function $\ell(\theta, x, y) = 1_{\{f(\theta, x) \neq y\}}$.
References

[1] S. M. Perlaza, G. Bisson, I. Esnaola, A. Jean-Marie, and S. Rini, “Empirical risk minimization with relative entropy regularization,” *IEEE Transactions on Information Theory*, 2024.

[2] H. E. Scarf, K. Arrow, and S. Karlin, “A min-max solution of an inventory problem,” Rand Corporation Santa Monica, Tech. Rep., 1957.

[3] E. Delage and Y. Ye, “Distributionally robust optimization under moment uncertainty with application to data-driven problems,” *Operations research*, vol. 58, no. 3, pp. 595–612, 2010.

[4] J. C. Duchi and H. Namkoong, “Learning models with uniform performance via distributionally robust optimization,” *The Annals of Statistics*, vol. 49, no. 3, pp. 1378–1406, 2021.

[5] A. Shapiro, “Distributionally robust stochastic programming,” *SIAM Journal on Optimization*, vol. 27, no. 4, pp. 2258–2275, 2017.

[6] F. Farokhi, “Distributionally robust optimization with noisy data for discrete uncertainties using total variation distance,” *IEEE Control Systems Letters*, vol. 7, pp. 1494–1499, 2023.

[7] A. Dixit, M. Ahmadi, and J. W. Burdick, “Distributionally robust model predictive control with total variation distance,” *IEEE Control Systems Letters*, vol. 6, pp. 3325–3330, 2022.

[8] F. Daunas, I. Esnaola, S. M. Perlaza, and H. V. Poor, “Equivalence of the empirical risk minimization to regularization on the family of f-divergences,” in *Proceedings of the IEEE International Symposium on Information Theory (ISIT)*, Athens, Greece, Jun. 2024.

[9] X. Zou, S. M. Perlaza, I. Esnaola, E. Altman, and H. V. Poor, “The worst-case data-generating probability measure in statistical learning,” *IEEE Journal on Selected Areas in Information Theory*, vol. 5, pp. 175–189, 2024.

[10] P. Mohajerin Esfahani and D. Kuhn, “Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations,” *Mathematical Programming*, vol. 171, no. 1, pp. 115–166, 2018.

[11] R. Gao and A. Kleywegt, “Distributionally robust stochastic optimization with wasserstein distance,” *Mathematics of Operations Research*, vol. 48, no. 2, pp. 603–655, 2023.

[12] X. Zou, S. M. Perlaza, I. Esnaola, and E. Altman, “Generalization analysis of machine learning algorithms via the worst-case data-generating probability measure,” in *Proceedings of the AAAI Conference on Artificial Intelligence*, Vancouver, Canada, Feb. 2024.
[13] R. Dudley, *Uniform Central Limit Theorems*. Cambridge University Press, 2014.

[14] Z. Hu and L. J. Hong, “Kullback-Leibler divergence constrained distribu-
tionally robust optimization,” *Optimization Online*, 2012.

[15] F. Tramer and D. Boneh, “Adversarial training and robustness for mul-
tiple perturbations,” in *International Conference on Neural Information
Processing Systems*, 2019, pp. 5866–5876.

[16] H. Xu, P. He, J. Ren, Y. Wan, Z. Liu, H. Liu, and J. Tang, “Probabilistic categorical adversarial attack and adversarial training,” in *International Conference on Machine Learning*. PMLR, 2023, pp. 38428–38442.

[17] M. Zhao, L. Zhang, J. Ye, H. Lu, B. Yin, and X. Wang, “Adversarial training: A survey,” *arXiv preprint arXiv:2410.15042*, 2024.

[18] J. Lee and M. Raginsky, “Minimax Statistical Learning with Wasserstein distances,” in *International Conference on Neural Information and Processing Systems*, vol. 31, 2018, pp. 2687–2696.

[19] A. Agarwal and T. Zhang, “Minimax regret optimization for robust ma-
chine learning under distribution shift,” in *Conference on Learning Theory*
PMLR, 2022, pp. 2704–2729.

[20] J. Dong, H.-L. Hsu, Q. Gao, V. Tarokh, and M. Pajic, “Variational adver-
sarial training towards policies with improved robustness,” in *International
Conference on Artificial Intelligence and Statistics*, 2025.

[21] H. Zhao, R. T. Des Combes, K. Zhang, and G. Gordon, “On learning in-
variant representations for domain adaptation,” in *International conference
on machine learning*. PMLR, 2019, pp. 7523–7532.

[22] H. Ye, C. Xie, T. Cai, R. Li, Z. Li, and L. Wang, “Towards a theoretical framework of out-of-distribution generalization,” in *International Conference on Neural Information Processing Systems*, vol. 34, 2021, pp. 23519–23531.

[23] W. Liu, G. Yu, L. Wang, and R. Liao, “An information-theoretic framework for out-of-distribution generalization,” *arXiv preprint arXiv:2403.19895*, 2024.

[24] I. Csiszár, “The method of types,” *IEEE Transactions on Information Theory*, vol. 44, no. 6, pp. 2505–2523, Oct. 1998.

[25] S. Geman and D. Geman, “Stochastic relaxation, gibbs distributions and the bayesian restoration of images,” *Journal of applied statistics*, vol. 20, no. 5-6, pp. 25–62, 1993.

[26] G. Casella and E. I. George, “Explaining the gibbs sampler,” *The American Statistician*, vol. 46, no. 3, pp. 167–174, 1992.
[27] W. Azizian, F. Iutzeler, J. Malick, and P. Mertikopoulos, “What is the long-run distribution of stochastic gradient descent? A large deviations analysis,” *arXiv preprint arXiv:2406.09241*, 2024.

[28] —, “The global convergence time of stochastic gradient descent in non-convex landscapes: Sharp estimates via large deviations,” *arXiv preprint arXiv:2503.16398*, 2025.

[29] I. Csiszár, “I-divergence geometry of probability distributions and minimization problems,” *The annals of probability*, pp. 146–158, 1975.

[30] O. Rivasplata, I. Kuzborskij, C. Szepesvári, and J. Shawe-Taylor, “PAC-Bayes analysis beyond the usual bounds,” in *International Conference on Neural Information Processing Systems*, vol. 33, 2020, pp. 16833–16845.

[31] G. Aminian, Y. Bu, L. Toni, M. Rodrigues, and G. Wornell, “An exact characterization of the generalization error for the Gibbs algorithm,” *Advances in Neural Information Processing Systems*, vol. 34, pp. 8106–8118, Dec. 2021.

[32] J. F. Nash Jr, “Equilibrium points in n-person games,” *Proceedings of the national academy of sciences*, vol. 36, no. 1, pp. 48–49, 1950.

[33] V. H. Stackelberg and A. T. Peacock, *Theory of The Market Economy*. William Hodge & Company Limited, 1952.

[34] F. Daunas, I. Esnaola, S. M. Perlaza, and H. V. Poor, “Asymmetry of the relative entropy in the regularization of empirical risk minimization,” *IEEE Transactions on Information Theory*, 2025.

[35] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg et al., “Scikit-learn: Machine learning in python,” *the Journal of machine Learning research*, vol. 12, pp. 2825–2830, 2011.

[36] D. Hendrycks and T. Dietterich, “Benchmarking neural network robustness to common corruptions and perturbations,” *International Conference on Learning Representations*, 2019.

[37] D. L. Cohn, *Measure theory*. Springer, 2013, vol. 2.

[38] R. B. Ash and C. A. Doleans-Dade, *Probability and Measure Theory*, 2nd ed. Burlington, MA, USA: Harcourt Academic Press, 2000.

[39] M. Sion, “On general minimax theorems,” *Pacific Journal of Mathematics*, 1957.

[40] L. McInnes, J. Healy, and J. Melville, “Umap: Uniform manifold approximation and projection for dimension reduction,” *arXiv preprint arXiv:1802.03426*, 2018.
informatics mathematics
Inria

RESEARCH CENTRE
SOPHIA ANTIPOLIS – MÉDITERRANÉE

2004 route des Lucioles - BP 93
06902 Sophia Antipolis Cedex

Publisher
Inria
Domaine de Voluceau - Rocquencourt
BP 105 - 78153 Le Chesnay Cedex
inria.fr

ISSN 0249-6399