Dimension-Free PAC-Bayes Generalization for Deep Neural Networks via Dirichlet Processes

Let f_\theta: \mathcal{X} \to \mathcal{Y} be an L-layer neural network parameterized by \theta = (W_1, \dots, W_L) \in \mathbb{R}^D:

where \sigma is a 1-Lipschitz activation function (such as ReLU). Let \hat{\theta} be an empirical risk minimizer, and consider a Gaussian posterior perturbation \theta = \hat{\theta} + \epsilon where \epsilon \sim \mathcal{N}(0, \sigma_{\mathrm{post}}^2 I_D).

By first-order Taylor expansion and the chain rule across L layers:

where \Sigma'_j are diagonal derivative matrices satisfying \|\Sigma'_j\|_{\mathrm{op}} \le 1. If each layer has nominal width W with balanced weight norms \|W_l\|_{\mathrm{op}} \sim \mathcal{O}(1), the spectral radius of the compositional operator scales as M_L = \sup_x \|\nabla_\theta f_\theta(x)\|_{\mathrm{op}} \ge c \cdot W^{L-1} for some constant c > 0.

Under the hypothesis that the posterior preserves empirical risk within margin \epsilon_0, we must have:

Now, computing the Kullback–Leibler divergence between two multivariate Gaussians in \mathbb{R}^D with \pi_0 = \mathcal{N}(0, \sigma_0^2 I_D) and \pi_n = \mathcal{N}(\hat{\theta}, \sigma_{\mathrm{post}}^2 I_D):

Substituting the upper bound \sigma_{\mathrm{post}} \le \frac{\epsilon_0}{c \cdot W^{L-1}}:

Since \frac{\sigma_{\mathrm{post}}^2}{\sigma_0^2} \ge 0, we obtain:

which establishes the claim.

Let (\mathcal{Z}, d_{\mathcal{Z}}) be a compact metric space. By the Banach–Alaoglu and Portmanteau theorems, the space \mathcal{M}_1(\mathcal{Z}) of Borel probability measures on \mathcal{Z}, equipped with the weak topology (the initial topology generated by bounded continuous test functions C_b(\mathcal{Z})), is a compact, metrizable Polish space [billingsley2013convergence]. A compatible metric generating this topology is the Wasserstein-1 metric W_1(P, Q) = \sup_{\|g\|_{\mathrm{Lip}} \le 1} |\int g \, d(P - Q)| or the Prokhorov metric. The Borel \sigma-algebra \mathcal{B}(\mathcal{M}_1(\mathcal{Z})) generated by this weak topology coincides with the cylindrical \sigma-algebra generated by the evaluation projections P \mapsto P(B) for all B \in \mathcal{B}(\mathcal{Z}).

To prove that a deterministic mapping \mathcal{A}: \mathcal{M}_1(\mathcal{Z}) \to \Theta \subseteq \mathbb{R}^D is (\mathcal{B}(\mathcal{M}_1(\mathcal{Z})), \mathcal{B}_\Theta)-measurable, it suffices to prove that \mathcal{A} is continuous with respect to the weak topology, since every continuous map between metric spaces is Borel measurable.

paragraph41.5ex plus 0.5ex minus .2ex-1emPart (i): Regularized Empirical Risk Minimization. Define the regularized risk functional J: \Theta \times \mathcal{M}_1(\mathcal{Z}) \to \mathbb{R} by:

For every fixed P \in \mathcal{M}_1(\mathcal{Z}), since \ell \ge 0 and the regularizer \theta \mapsto \frac{\lambda}{2}\|\theta\|_2^2 is \lambda-strongly convex on the closed convex set \Theta, the functional J(\cdot, P) is strictly convex and coercive: \lim_{\|\theta\| \to \infty} J(\theta, P) = \infty. Consequently, the minimizer \mathcal{A}_{\mathrm{ERM}}(P) = \arg\min_{\theta \in \Theta} J(\theta, P) exists and is unique for every P \in \mathcal{M}_1(\mathcal{Z}).

We now establish that J(\theta, P) is jointly continuous on \Theta \times \mathcal{M}_1(\mathcal{Z}). Let (\theta_m, P_m) \to (\theta, P) where \theta_m \to \theta in Euclidean norm on \mathbb{R}^D and P_m \xrightarrow{w} P weakly. Decomposing the difference:

Because \mathcal{Z} is compact and (z, \theta) \mapsto \ell(f_\theta(x), y) is continuous, the family of functions \{\ell(f_{\theta_m}(\cdot), \cdot)\} converges uniformly to \ell(f_\theta(\cdot), \cdot) on \mathcal{Z} as \theta_m \to \theta by the Heine–Cantor theorem. Thus:

which bounds the first integral uniformly over all probability measures P_m:

The second term vanishes because z \mapsto \ell(f_\theta(x), y) is a bounded continuous test function on compact \mathcal{Z} and P_m \xrightarrow{w} P weakly. The third term vanishes by continuity of the Euclidean norm. Hence J is jointly continuous.

By Berge’s Maximum Theorem [berge1963topological], the argmin correspondence P \mapsto \arg\min_{\theta \in \Theta} J(\theta, P) is upper hemicontinuous. Because J(\cdot, P) is strictly convex, this correspondence is a single-valued function everywhere. An upper hemicontinuous single-valued correspondence into a Hausdorff space is continuous. Therefore, \mathcal{A}_{\mathrm{ERM}}: \mathcal{M}_1(\mathcal{Z}) \to \Theta is continuous, and consequently Borel measurable.

paragraph41.5ex plus 0.5ex minus .2ex-1emPart (ii): T-Step Gradient Descent. Let \theta_0 \in \Theta be fixed. Define the gradient integration operator G: \Theta \times \mathcal{M}_1(\mathcal{Z}) \to \mathbb{R}^D by:

We show that G is jointly continuous. Let (\theta_m, P_m) \to (\theta, P). By the triangle inequality:

By uniform continuity of the gradient on compact subsets of \mathcal{Z} \times \Theta, \sup_{z \in \mathcal{Z}} \|\nabla_\theta \ell(f_{\theta_m}(x), y) - \nabla_\theta \ell(f_\theta(x), y)\|_2 \to 0, vanishing the first term. The second term vanishes coordinate-wise because each component z \mapsto \frac{\partial}{\partial \theta_j} \ell(f_\theta(x), y) is a bounded continuous test function under weak convergence P_m \xrightarrow{w} P. Hence G(\theta, P) is jointly continuous.

We proceed by induction on t \in \{0, 1, \dots, T\}. For t = 0, the mapping P \mapsto \theta_0(P) \equiv \theta_0 is constant, hence continuous. Now assume \theta_t: \mathcal{M}_1(\mathcal{Z}) \to \Theta is continuous. The recurrence relation is:

Since \theta_{t+1} is the composition of the continuous map P \mapsto (\theta_t(P), P) with the continuous map (\theta, P) \mapsto \theta - \eta G(\theta, P), it follows that \theta_{t+1}: \mathcal{M}_1(\mathcal{Z}) \to \Theta is continuous. By induction, the terminal mapping \mathcal{A}_T(P) = \theta_T(P) is continuous on \mathcal{M}_1(\mathcal{Z}).

Since every continuous map between metric spaces is Borel measurable, \mathcal{A}_T is strictly (\mathcal{B}(\mathcal{M}_1(\mathcal{Z})), \mathcal{B}_\Theta)-measurable, which completes the proof.

Let (\Omega_1, \Sigma_1) and (\Omega_2, \Sigma_2) be standard Borel measurable spaces. Let \mu, \nu be two probability measures on (\Omega_1, \Sigma_1) such that \mu \ll \nu. Let T: \Omega_1 \to \Omega_2 be a (\Sigma_1, \Sigma_2)-measurable mapping. Let T_\# \mu = \mu \circ T^{-1} and T_\# \nu = \nu \circ T^{-1} denote the respective push-forward measures on (\Omega_2, \Sigma_2).

Define the sub-\sigma-algebra \mathcal{G} = T^{-1}(\Sigma_2) \subseteq \Sigma_1. The restriction of \mu and \nu to \mathcal{G} satisfies \mu|_{\mathcal{G}} = T_\# \mu and \nu|_{\mathcal{G}} = T_\# \nu via the isomorphism between \mathcal{G} and \Sigma_2.

By the Radon–Nikodym theorem, let X = \frac{d\mu}{d\nu} \in L^1(\Omega_1, \Sigma_1, \nu) be the density of \mu with respect to \nu. For any event E \in \Sigma_2:

Therefore, the Radon–Nikodym derivative of the push-forward measures is the conditional expectation:

Now, consider the relative entropy function f(t) = t \ln t, which is strictly convex on (0, \infty) with f(1) = 0. By the conditional Jensen’s inequality for convex functions:

Setting \Omega_1 = \mathcal{M}_1(\mathcal{Z}), \Omega_2 = \Theta, T = \mathcal{A}, \mu = \mathcal{Q}_n, and \nu = \mathcal{P}_0 completes the proof.

Let \mathcal{P} = \{A_1, \dots, A_K\} be a partition of \mathcal{Z}. Under \mathcal{P}_0 = \mathcal{DP}(\alpha, F_0), the random vector \mathbf{p} = (P(A_1), \dots, P(A_K)) is distributed as \mathrm{Dir}(\mathbf{a}) with \mathbf{a} = (a_1, \dots, a_K) and a_k = \alpha F_0(A_k), where a_0 = \sum_{k=1}^K a_k = \alpha.

Under the posterior \mathcal{Q}_n = \mathcal{DP}(\alpha + n, \frac{\alpha F_0 + n \mathbb{P}_n}{\alpha + n}), the parameter vector updates to \mathbf{b} = (b_1, \dots, b_K) with b_k = a_k + n_k, where n_k = \sum_{i=1}^n \mathbf{1}\{z_i \in A_k\} and b_0 = \sum_{k=1}^K b_k = \alpha + n.

The density of \mathrm{Dir}(\mathbf{b}) with respect to Lebesgue measure on the simplex \Delta_{K-1} is:

Similarly, the prior density is p_0(\mathbf{p}) = \frac{1}{\mathrm{B}(\mathbf{a})} \prod_{k=1}^K p_k^{a_k - 1}. The relative entropy between two Dirichlet distributions on \Delta_{K-1} is:

The ratio of multivariate Beta normalizing constants expands as:

For any Dirichlet distribution \mathbf{p} \sim \mathrm{Dir}(\mathbf{b}), the marginal distribution of p_k is a Beta distribution \mathrm{Beta}(b_k, b_0 - b_k). The expected logarithm of a Beta-distributed random variable has the exact analytical form:

Substituting b_k - a_k = n_k:

Summing both terms gives:

which is the exact analytical identity of Theorem 3.