Controlling Multiple Errors Simultaneously with a PAC-Bayes Bound

Authors: reuben  adams
john  shawe-taylor
benjamin  guedj
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:52:02

Abstract: Current PAC-Bayes generalisation bounds are restricted to scalar metrics of performance, such as the loss or error rate. However, one ideally wants more information-rich certificates that control the entire distribution of possible outcomes, such as the distribution of the test loss in regression, or the probabilities of different mis-classifications. We provide the first PAC-Bayes bound capable of providing such rich information by bounding the Kullback-Leibler divergence between the empirical and true probabilities of a set of $M$ error types, which can either be discretized loss values for regression, or the elements of the confusion matrix (or a partition thereof) for classification. We transform our bound into a differentiable training objective. Our bound is especially useful in cases where the severity of different mis-classifications may change over time; existing PAC-Bayes bounds can only bound a particular pre-decided weighting of the error types. In contrast our bound implicitly controls all uncountably many weightings simultaneously.

Main paper: 
 introduction: 
 Generalisation bounds are a core component of the theoretical understanding of machine learning algorithms. For over two decades now, PAC-Bayesian theory has been at the core of studies on generalisation abilities of machine learning algorithms. PAC-Bayes originated in the seminal work of McAllester [1998 McAllester [ , 1999] ] and was further developed by Catoni [2003 Catoni [ , 2004 Catoni [ , 2007]] , among other authors-we refer to the surveys Guedj [2019] and Alquier [2021] for an introduction to the field. The outstanding empirical success of deep neural networks in the past decade calls for better theoretical understanding of deep learning, and PAC-Bayes has emerged as one of the few frameworks that can be used to derive meaningful (and non-vacuous) generalisation bounds for neural networks: the pioneering work of Dziugaite and Roy [2017] has been followed by a number of contributions, including Neyshabur et al. [2018] , Zhou et al. [2019] , Letarte et al. [2019] , Pérez-Ortiz et al. [2021] , Perez-Ortiz et al. [2021] and Biggs and Guedj [2022a,b] , to name but a few.Much of the PAC-Bayes literature focuses on the case of binary classification, or of multiclass classification where one only distinguishes whether each classification is correct or incorrect. This is in stark contrast to the complexity of contemporary real-world learning problems, such as medical diagnosis where the severity of Type I and Type II errors may be crucial and context-dependent. This work aims to bridge this gap by deriving a generalisation bound that provides information-rich measures of performance at test time by controlling the probabilities of errors of any finite number of user-specified types. More precisely, we bound the KL-divergence between the empirical and true distributions over the different error types. From this single bound one can then derive bounds on arbitrary linear combinations of these error probabilities, which will all hold simultaneously with the same probability as the original bound. In addition, these bounds are guaranteed to be non-vacuous (this follows since the KL-divergence blows up on the boundary of the simplex).As a concrete example, if the severity of Type I and Type II errors of a medical test are contextdependent, one would want to be able to bound arbitrary linear combinations of these error probabilities. Existing bounds could only bound finitely many pre-specified weightings by employing a union bound, which would also degrade the bound. In contrast, by constraining the KL-divergence between the true and empirical error probabilities, our bound constrains all uncountably many weightings of the error probabilities simultaneously.The usual setting of PAC-Bayes bounds is that of binary classification, namely an input set X , output set Y = {-1, 1}, hypothesis space H ⊆ Y X and a sample S ∈ (X × Y) m drawn i.i.d. from a datagenerating distribution D. A number of PAC-Bayes bounds in this setting (e.g. Maurer [2004] ) have been unified by a single general bound found in Bégin et al. [2016] . Briefly, Bégin et al. [2016] prove a bound on the discrepancy d(R S (Q), R D (Q)) between the error probability R D (Q) of a stochastic classifier Q (a distribution over H which classifies by first drawing h ∼ Q and then classifying according to h) and its empirical counterpart R S (Q) (the fraction of the sample Q misclassifies). The bound holds with high probability for all Q simultaneously. The bound in Bégin et al. [2016] is binary in the sense that Y contains two elements, but a more subtle way to look at this is that only two cases are distinguished-correct classification and incorrect classification. While it can be applied to multiclassification provided one maintains the second binary characteristic by only distinguishing correct and incorrect classifications. It is this heavy restriction that our result lifts, by considering the new framework of error types.A new framework of errors types We consider a user-specified partition of the space Y × Y of prediction-truth label-pairs into a finite partition of error types E 1 , . . . , E M . Our bound then simultaneously constrains the probability with which errors of each type occur. In multiclass classification for example, one can choose the error types to be the set of all different possible mis-classifications, in which case our bound will control the entire confusion matrix, bounding how far the true confusion matrix (i.e. expected over the data-generating distribution) can diverge from the empirical one (i.e. on the training set). From this one can then derive bounds on the probabilities with which each mis-classification may be made, and arbitrary linear combinations of these error probabilities, and all of these will hold simultaneously with the same probability as the original bound. Our bound therefore paints a far richer picture of the performance of the final learned model than can be provided by any existing PAC-Bayes bound.Formally, we let M j=1 E j be a user-specified disjoint partition of Y 2 into a finite number of M error types, where we say that a hypothesis h ∈ H makes an error of type j on datapoint (x, y) if (h(x), y) ∈ E j (by convention, every pair (ŷ, y) ∈ Y 2 is interpreted as a predicted value ŷ followed by a true value y, in that order). It should be stressed that not all of the E j need correspond to mislabellings-indeed, some of the E j may distinguish different correct labellings.Relation to previous results. Our framework of a finite number of user-specified "error types" includes multiclass classification as a particular case, and it is in this field that one finds the work most closely related to ours. Little is known of multiclass classification from a theoretical perspective and, to the best of our knowledge, only a handful of relevant strategies or generalisation bounds can be compared to the present paper.The closest is the work of Morvant et al. [2012] , which establishes a PAC-Bayes bound on the spectral norm of the difference between the true and empirical confusion matrices. Our bound differs from theirs in two respects. First, they consider the confusion matrix, whereas ours applies to the more general setting of a finite number of error types, which can be the set of all mis-classifications, or some partition thereof. Second, they deal with the spectral norm, whereas we employ the KL-divergence. Since the KL-divergence follows a simple formula, this means we can much more easily infer bounds on the individual error probabilities, which would be very challenging for the spectral norm. The follow-up work Koço and Capponi [2013] shows how a proxy of the spectral norm bound can be used as a training objective that may deal with imbalanced classes. In the present work, we show how our bound can be used as a differentiable training objective directly (without the need of a proxy) and that it can more sensitively deal with imbalanced classes, or errors of different severity, by assigning each error type a user-specified loss value. Laviolette et al. [2017] extend the celebrated C-bound in PAC-Bayes to ensembles, obtaining a bound on the risk of the majority vote classifier in the case of multiclass classification. In this context, our bound is able to distinguish different mis-classifications and control them, whereas they bound the scalar risk which lumps all mis-classifications together. The C-bound has alternately been generalised by Lacasse et al. [2006] (see also Germain et al. [2015] ) to simultaneously control three metrics, namely the so-called expected disagreement, expected joint success and expected joint error of the posterior. While they restricted themselves to the ternary case, some of their proof techniques share similarities with ours. In cases where one has exactly three error types, for example the {-1, 0, 1}valued excess loss, the work of Wu and Seldin [2022] is applicable; they construct so-called 'split-kl' inequalities (both classical and PAC-Bayesian) which deftly handle this specific scenario. Pires et al. [2013] present a comprehensive analysis of convex surrogate losses in cost-sensitive multiclass classification, providing conditions for consistency, bounding the excess loss of a predictor, and extending the analysis to the "Simplex Coding" scheme. We are considering the generalisation gap rather than the excess loss. Lei et al. [2019] study data-dependent bounds for multiclass classification. Their analysis is restricted to SVMs however, whereas ours applies to arbitrary hypothesis spaces.Outline. We fix notation in Section 2. Theorem 1 in Section 3 is our main result-a PAC-Bayes bound on the KL-divergence between the true and empirical error distributions. For multiclass classification with a fully refined partition this becomes a bound on the KL-divergence between the true and empirical confusion matrices. Proposition 1 then bounds the individual error probabilities. Our second main result, Theorem 2 in Section 4, allows us to use bounds on linear combinations of error probabilities as training objectives. We prove Theorem 1 in Section 5 via Proposition 4, which bounds the distribution of errors via a general convex function d and may be of independent interest. Section 6 outlines positive empirical results1 from using our bound as a training objective for neural networks and Section 7 gives perspectives for follow-up work.. 
notation: 
 For any set A, let M(A) be the set of probability measures on A. Let X and Y be arbitrary input (e.g., feature) and output (e.g., label) sets respectively, and D ∈ M(X × Y) be a data-generating distribution. For any sample S ∼ D m drawn i.i.d. from D, let D(S) ∈ M(X × Y) denote the empirical distribution D(S) := 1 m (x,y)∈S δ (x,y) . We consider the setting where the user has specified a partition {E 1 , . . . , E M } of Y 2 into M error types.We are interested in simple hypotheses h : X → Y and soft hypotheses H : X → M(Y). For example, a neural network outputting scores (logits) in R Y is converted to a simple or soft hypothesis, respectively, by passing the scores through the argmax or softmax function, respectively. For any A ⊆ Y, H(x)(A) can be interpreted as the probability according to H that the label of x is in A. We will see in Section 4 that soft hypotheses permit more flexible training procedures and a more finegrained analysis. Note that while soft hypotheses output distributions, they do so deterministically, always returning the same distribution for the same input x, and so are distinct from the stochastic classifiers introduced shortly.For a simple hypothesis h : X → Y and j ∈ [M ], define the j-risk of h to be R j D (h) := P (x,y)∼D ((h(x) , y) ∈ E j ), namely the probability that h makes an error of type E j for a randomly sampled (x, y) ∼ D. For a soft hypothesis H : X → M(Y) define the j-risk of H to be R j D (H) := P (x,y)∼D,ŷ∼H(x) ((ŷ, y) ∈ E j ), namely the probability that one would make an error of type E j on a randomly sampled (x, y) ∼ D if one predicted by sampling ŷ from the distribution H(x). From now until Section 4 it will not matter whether we are dealing with simple or soft hypotheses. So, unless stated explicitly, we will refer to both simply as hypotheses, denote both by lowercase h, and refer to the hypothesis class H, whether it is a subset of Y X or M(Y) X .Our goal is to control the risk vector R D (h) := (R 1 D (h), . . . , R M D (h)), since controlling this vector controls all linear combinations of j-risks. Since this is unobervable, we will control it by bounding how far it diverges from its empirical counterpart R S (h) := R D(S) (h), which we term the empirical risk vector. Note that E S∼Dm R S (h) = R D (h), and that, for a simple hypothesis h ∈ Y X , R S (h) is the vector of proportions of the sample on which h makes an error of type E j 2 . Since the E j partition Y 2 , R D (h) and R S (h) are elements of the M -dimensional simplex△ M := {u ∈ [0, 1] M : u 1 + • • • + u M = 1}.Thus we can choose our divergence measure to be kl(R S (h)∥R D (Q)), where for q, p ∈ △ M we define kl(q∥p) := M j=1 q j ln qj pj . 3 When M = 2 we abbreviate kl((q, 1 -q)∥(p, 1 -p)) to kl(q∥p), which is then the conventional definition of kl(•∥•) found in the PAC-Bayes literature [as in Seeger, 2002, for example] . We define the risk and empirical risk of Q as R D (Q) := E h∼Q R D (h) and R S (Q) := E h∼Q R S (h), respectively, and seek a bound on kl(R S (Q)∥R D (Q)). Note we still have E S [R S (Q)] = R D (Q), this time using Fubini. Moreover, for a sample S of size m, we have that R S (Q) = K/m where K ∼ Mult(m, M, R D (Q)). Recall that for m, M ∈ N and r ∈ △ M , the multinomial distribution Mult(m, M, r) has probability mass function Mult(k; m, M, r) :=m k1 k2 ••• k M M j=1 r kj j , where m k1 k2 ••• k M := m! M j=1 kj ! for k ∈ S m,M := (k 1 , . . . , k M ) ∈ N M 0 : k 1 + • • • + k M = m, and zero otherwise. As a final piece of notation, we let △ >0 M := △ M ∩ (0, 1) M and S >0 m,M := S m,M ∩ N M denote the vector elements of △ M and S m,M , respectively, that have no zero components. 
main result: 
 We now state our main result, which bounds the KL-divergence between the true and empirical risk vectors R D (Q) and R S (Q), interpreted as probability distributions. As is conventional in the PAC-Bayes literature, we refer to sample independent and dependent distributions on M(H) (i.e. stochastic hypotheses) as priors (denoted P ) and posteriors (denoted Q) respectively, even if they are not related by Bayes' theorem.Theorem 1. Let X and Y be arbitrary sets and M j=1 E j be a disjoint partition of Y 2 into M error types. Let D ∈ M(X × Y) be a data-generating distribution and H be a simple (H ⊆ Y X ) or soft (H ⊆ M(Y) X ) hypothesis class. For any prior P ∈ M(H), δ ∈ (0, 1] and sample size m ≥ M , with probability at least 1 -δ over the random draw S ∼ D m , we have that simultaneously for all posteriors Q ∈ M(H), the divergence kl R S (Q)∥R D (Q) is upper bounded byEQUATIONξ(m, M ) := √ πe 1/(12m) m 2 M -1 2 M -1 z=0 M z 2 m z/2 Γ M -z 2 -1 ∈ O (mM ) M .The fact that the logarithmic term is of order O(M ln(mM/δ)) means the bound is linear in M up to logarithmic terms, while this may seem excessive, one should note that the quantity that our theorem bounds also depends on M . Further, the bound has been successfully used in by Biggs and Guedj [2023] to improve on state of the art PAC-Bayes bounds.To see how our bound compares to existing PAC-Bayes bounds for binary classification, takeY = {-1, 1}, M = 2, E 1 = {(-y, y) : y ∈ Y} and E 2 = {(y, y) : y ∈ Y}.The argument of the logarithm then reduces to 1 δ e 1/(12m) 2 + πm 2 ≤ 1.25 √ m when m is large. The corresponding term in Maurer [2004] is 2 √ m, which is only larger because he relaxes the term for aesthetics. Therefore our bound gracefully reduces to Maurer's in the case of binary classification.Suppose after a use of Theorem 1 we have a bound of the form kl(R S (Q)∥R D (Q)) ≤ B. We can then derive bounds on the individual j-risks R j D (Q) or, more generally, on linear combinations thereof. While one could obtain such bounds perhaps more directly with existing PAC-Bayes bounds, the significance of our bound is that all such derived bounds hold with high probability simultaneously. Existing PAC-Bayes bounds would require the use of a union bound in order to bound multiple combinations simultaneously, whereas ours bounds all uncountably many combinations simultaneously, as a package. As for the individual j-risks R j D (Q), the following proposition2 (RS(h))j = R j D (h) = P (x,y)∼ D ((h(x), y) ∈ Ej) = 1 m (x,y)∈S 1[(h(x), y) ∈ Ej].3 We follow the usual convention that 0 ln 0 x = 0 for x ≥ 0 and x ln x 0 = ∞ for x > 0.then yields the boundsL j ≤ R j D (Q) ≤ U j , where L j := inf{p ∈ [0, 1] : kl(R j S (Q)∥p) ≤ B} and U j := sup{p ∈ [0, 1] : kl(R j S (Q)∥p) ≤ B}.Moreover, since in the worst case we have kl(R S (Q)∥R D (Q)) = B, the proposition shows that the lower and upper bounds L j and U j are the tightest possible, since ifR j D (Q) ̸ ∈ [L j , U j ] then kl(R j S (Q)∥R j D (Q)) > B implying kl(R S (Q)∥R D (Q)) > B.For a more precise version of this argument and a proof of Proposition 1, see Appendix C.4. Proposition 1. Let q, p ∈ △ M . Then kl(q j ∥p j ) ≤ kl(q∥p) for all j ∈ [M ], with equality whenp i = 1-pj 1-qj q i . for all i ̸ = j.More generally, suppose we can quantify how costly an error of each type is by means of a loss vector ℓ ∈ [0, ∞) M , where ℓ j is the loss we attribute to an error of type E j . We may then be interested in bounding the total riskEQUATIONwhere we define kl -1 ℓ (u|c) ∈ △ M as follows. To see that it is indeed well-defined (at least when u ∈ △ >0 M ), see the discussion at the beginning of Appendix C.5.EQUATIONThis motivates the following training procedure: search for a posterior Q for which the boundℓ • kl -1 ℓ (R S (Q)|B) on the total risk R T D (Q) is minimised.While this requires a particular choice of loss vector ℓ, we emphasise that at the end of training, Theorem 1 bounds kl(R S (Q)∥R D (Q)) and so can be used to bound any linear combination of the j-risks, not just the one given the loss vector ℓ chosen for training. It is this flexibility which is the main advantage of our bound; changes in the severity of different error types over time do not require union bounds or retraining.In the next section we provide a theorem for calculating kl -1 ℓ (u|c) and its derivatives so that the training procedure can be executed. 
construction of a differentiable training objective: 
 We now state and prove Theorem 2, which provides a speedy method for approximating kl -1 ℓ (u|c) and its derivatives to arbitrary precision, provided c > 0 and ∀j u j > 0. The only approximation step required is that of approximating the unique root of a continuous and strictly increasing scalar function. Thus, provided the u j themselves are differentiable, Theorem 1 combined with Theorem 2 shows that the upper bound on the total risk can be used as a tractable and fully differentiable training objective. See Appendix A for more details, including a pseudocode algorithm and an implementation. Since the proof of Theorem 2 is rather long and technical, we defer it to Appendix C.5. The requirement that the ℓ j are not all equal only rules out trivial cases whereR T D (Q) is independent of R D (Q). Theorem 2. Fix ℓ ∈ [0, ∞) M suchthat not all ℓ j are equal, and definef ℓ : △ M → [0, ∞) by f ℓ (v) := M j=1 ℓ j v j . For all ũ = (u, c) ∈ △ >0 M ×(0, ∞), define v * ( ũ) := kl -1 ℓ (u|c) ∈ △ M and let µ * ( ũ) ∈ (-∞, -max j ℓ j ) be the unique solution to c = ϕ ℓ (µ), where ϕ ℓ : (-∞, -max j ℓ j ) → R is given by ϕ ℓ (µ) := ln(- M j=1 uj µ+ℓj ) + M j=1 u j ln(-(µ + ℓ j )), which is continuous and strictly increasing. Then v * ( ũ) = kl -1 ℓ (u|c) is given by v * ( ũ) j = λ * ( ũ)u j µ * ( ũ) + ℓ j for j ∈ [M ], where λ * ( ũ) =   M j=1 u j µ * ( ũ) + ℓ j   -1.(5)EQUATIONA final wrinkle in evaluating our bound is that while the empirical risk vector R S (Q) = E h∼Q R S (h) does not depend on the data-generating distribution D, the expectation over Q may still be intractable. This would be the default case when Q is a Gaussian over the weights of a multilayer perceptron, for example. In such cases, we can estimateR S (Q) via a Monte Carlo sample R S ( Q) := 1 N N n=1 R S (h n ) (where the h n are drawn i.i.d. from Q) and use the following two results. Proposition 2 shows that the kl(R j S ( Q)∥R j D (Q)) can be simultaneously bounded, whence Proposition 3 can be used to obtain a bound on kl(R S ( Q)∥R D (Q)). Proposition 2. Let X ∼ Multinomial(N, M, p). Then for any δ ∈ (0, 1), with probability at least1 -δ we have that for all j ∈ [M ] simultaneously kl 1 N X j p j ≤ ln 2M δ N .Proof. Each bound holds separately with probability at least 1 -δ/M by Theorem 2.5 in Langford and Caruana [2001] . They then hold simultaneously by application of a union bound. Proposition 3. Suppose q, p, q ∈ △ M are such that kl(q∥p) ≤ B 1 and kl(q j ∥q j ) ≤ B 2 for all j ∈ [M ]. For each j, define q j = inf{r ∈ [0, 1] : kl(q j ∥r) ≤ B 2 }. ThenEQUATIONProof. Deferred to C.1.The fact that the bound on kl( q∥p) → B 1 as B 2 → 0 ensures that as we increase the size of our Monte Carlo sample for estimatingR S (Q) the bound on kl(R S ( Q)∥R D (Q)) approaches that of kl(R S (Q)∥R D (Q)), meaning in the limit we pay an arbitrarily small price in the bound for the approximation. 
proof of the main bound: 
 We split the proof of Theorem 1 into three parts. First, we prove Proposition 4, a bound on d(R S (Q), R D (Q))EQUATIONI d (m, β) := sup r∈△ M k∈S m,M Mult(k; m, M, r) exp βd k m , r .This is a generalisation of the unifying PAC-Bayes bound given in Bégin et al. [2016] where we replace the scalar risk quantities R S (Q) and R D (Q) with their vector counterparts R S (Q) and R D (Q). To see this, note that we can recover it by setting Bégin et al. [2016] . Further, k∈Sm,2 Mult(k; m, 2, r) exp βd ′ k m , r = m k=0 Bin(k; m, r 1 ) exp βd k m , r 1 , so that the supremum over r 1 ∈ [0, 1] of the right hand side equals the supremum over r ∈ △ 2 of the left hand side, which, when substituted into (8), yields the bound given in Bégin et al. [2016] .Y = {-1, 1}, M = 2, E 1 = {(-y, y) : y ∈ Y} and E 2 = {(y, y) : y ∈ Y}. Then, for any convex function d : [0, 1] 2 → R, apply Theorem 4 with the convex function d ′ : △ 2 M → R defined by d ′ ((u 1 , u 2 ), (v 1 , v 2 )) := d(u 1 , v 1 ) so that Theorem 4 bounds d ′ R S (Q), R D (Q) = d R 1 S (Q), R 1 D (Q) which equals d(R S (Q), R D (Q)) in the notation ofTo prove Proposition 4 we require the following two lemmas. The first is the well-known change of measure in equality (Csiszár, 1975, Donsker and Varadhan, 1975) . The second is a generalisation from Binomial to Multinomial distributions of a result found in Maurer [2004] , the proof of which we defer to Appendix C.2. Lemma 1. For any set H, any P, Q ∈ M(H) and any measurable function ϕ :H → R, E h∼Q ϕ(h) ≤ KL(Q∥P ) + ln E h∼P exp(ϕ(h)).Lemma 2. Let X 1 , . . . , X m be i.i.d △ M -valued random vectors with mean µ and suppose that f :△ m M → R is convex. If X ′ 1 , . . . , X ′ m are i.i.d. Mult(1, M, µ) random vectors, then E[f (X 1 , . . . , X m )] ≤ E[f (X ′ 1 , . . . , X ′ m )].The consequence of Lemma 2 is that the worst case (in terms of bounding d(R S (Q), R D (Q))) occurs when R {(x,y)} (h) is a one-hot vector for all (x, y) ∈ S and h ∈ H, namely when H ⊆ M(Y) X only contains hypotheses that, when labelling S, put all their mass on elements ŷ ∈ Y that incur the same error type4 . In particular, this is the case for hypotheses that put all their mass on a single element of Y, equivalent to the simpler case H ⊆ Y X as discussed in Section 2. Thus, Lemma 2 shows that the bound given in Proposition 4 cannot be made tighter only by restricting to such hypotheses.Proof. (of Proposition 4) The case H ⊆ Y X follows directly from the more general case by takingH ′ := {h ′ ∈ M(Y) X : ∃h ∈ H such that ∀x ∈ X h ′ (x) = δ h(x) }, where δ h(x) ∈ M(Y) denotes a point mass on h(x). For the general case H ⊆ M(Y) X , using Jensen's inequality with the convex function d(•, •) and Lemma 1 withϕ(h) = βd(R S (h), R D (h)), we see that for all Q ∈ M(H) βd R S (Q), R D (Q) = βd E h∼Q R S (h), E h∼Q R D (h) ≤ E h∼Q βd R S (h), R D (h) ≤ KL(Q∥P ) + ln E h∼P exp βd R S (h), R D (h) = KL(Q∥P ) + ln(Z P (S)),where Z P (S) := E h∼P exp βd(R S (h), R D (h)) . Note that Z P (S) is a non-negative random variable, so that by Markov's inequality P S∼D m Z P (S)≤ E S ′ ∼D m Z P (S ′ ) δ ≥ 1 -δ. Thus, since ln(•)is strictly increasing, with probability at least 1 -δ over S ∼ D m , we have that simultaneously for allEQUATIONTo boundE S ′ ∼D m Z P (S ′ ), let X i := R {(xi,yi) ′ } (h) ∈ △ M for i ∈ [m], where (x i , y i ) ′ is the i'th element of the dummy sample S ′ . Noting that each X i has mean R D (h), define the random vectorsX ′ i ∼ Mult(1, M, R D (h)) and Y := m i=1 X ′ i ∼ Mult(m, M, R D (h)). Finally let f : △ m M → R be defined by f (x 1 , . . . , x m ) := exp βd 1 m m i=1 x i , R D (h), which is convex since the average is linear, d is convex and the exponential is non-decreasing and convex. Then, by swapping expectations (which is permitted by Fubini's theorem since the argument is non-negative) and applying Lemma 2, we have that E S ′ ∼D m Z P (S ′ ) can be written asE S ′ ∼D m Z P (S ′ ) = E S ′ ∼D m E h∼P exp βd R S ′ (h), R D (h) = E h∼P E S ′ ∼D m exp βd R S ′ (h), R D (h) = E h∼P E X1,...,Xm exp βd 1 m m i=1 X i , R D (h) ≤ E h∼P E X ′ 1 ,...,X ′ m exp βd 1 m m i=1 X ′ i , R D (h) = E h∼P E Y exp βd 1 m Y , R D (h) = E h∼P k∈S m,M Mult k; m, M, R D (h) exp βd k m , R D (h) ≤ sup r∈△ M   k∈S m,M Mult k; m, M, r exp βd k m , r   ,which is the definition of I d (m, β). Inequality (8) then follows by substituting this bound on E S ′ ∼D m Z P (S ′ ) into (9) and dividing by β.We now specialise Proposition 4 to the case d(•, •) = kl(•∥•) to obtain Corollary 1.Corollary 1. In the setting of Theorem 1,EQUATIONEQUATION)Proof. Applying Proposition 4 with d(•, •) = kl(•∥•) and β = m gives that with probability at least1 -δ over S ∼ D m , simultaneously for all posteriors Q ∈ M(H), kl R S (Q)∥R D (Q) ≤ 1 m KL(Q∥P ) + ln I kl (m, m) δ ,whereI kl (m, m) := sup r∈△ M [ k∈S m,M Mult(k; m, M, r) exp mkl( k m , r) ]. Thus it suffices to show that I kl (m, m) ≤ η(m, M ).To prove this, for each fixed r = (r 1 , . . . , r M ) ∈ △ M let J r = {j ∈ [M ] : r j = 0}. Then Mult(k; m, M, r) = 0 for any k ∈ S m,M such that k j ̸ = 0 for some j ∈ J r . For the other k ∈ S m,M , namely those such that k j = 0 for all j ∈ J r , the probability term can be written as Mult(k; m, M, r) = ∀j∈Jr kj =0 M j=1 k kj j k j ! (because 0 0 0! = 1) ≤ m! m m k∈S m,M M j=1 k kj j k j ! ,which is η(m, M ). Since this is independent of r, it also holds after taking the supremum over r ∈ △ M of the left hand side, showing that I kl (m, m) ≤ η(m, M ).The final step in obtaining Theorem 1 is to loosen the bound given in Corollary 1 (which is intractable when m is large) to the tractable form given in Theorem 1. For this we require the following technical lemma, the proof of which we defer to Appendix C.3.Lemma 3. For integers M ≥ 1 and m ≥ M , k∈S >0 m,M 1 M j=1 √ kj ≤ π M 2 m M -2 2 Γ( M 2 ).Proof. (Of Theorem 1) It suffices to show that for all m ≥ M ≥ 1 we have η(m, M ) ≤ ξ(m, M ).We achieve this by applying Stirling's approximation√ 2πn n e n < n! < √ 2πn ne n e 1 12n (valid for n ≥ 1) to the factorials in η(m, M ) and then using Lemma 3. Since Stirling's approximation requires that all the k j are at least one, we partition the sum in η(m, M ) according to the number of coordinates of k at which k j = 0. Let z index the number of such coordinates. Defining f :∞ M =2 S m,M → R by f (k) = |k| j=1 k kj j /k j !and noting that f is symmetric under permutations of its arguments, we then haveEQUATIONStirling's approximation can now be applied to eachk ∈ S >0 m,M f (k) ≤ M j=1 k k j j √ 2πkj k j e k j = M j=1 e k j √ 2πkj = e m (2π) M/2 M j=1 1 √ kj. An application of Lemma 3 now givesk∈S >0 m,M -z f (k) ≤ k∈S >0 m,M -z e m (2π) M -z 2 M -z j=1 1 k j ≤ e m (2π) M -z 2 π M -z 2 m M -z-2 2 Γ M -z 2 = e m m M -z-2 2 2 M -z 2 Γ M -z 2 .Substituting this into equation ( 12) and bounding m! using Stirling's approximation, we haveη(m, M ) ≤ √ 2πme 1/(12m) e m M -1 z=0 M z e m m M -z-2 2 2 M -z 2 Γ( M -z 2 )= ξ(m, M ), which completes the proof of the bound. As for the order of the bound, it is sufficient to bound ln ξ(m, M ) using the crude approximations M z ≤ M M , (2/m) z/2 ≤ 1 and Γ((M -z)/2) ≥ 1. 
numerical experiments: 
 We use binarised versions of MNIST, and HAM10000 Tschandl [2018] . In both cases we partition Y 2 into E 0 = {(0, 0), (1, 1)}, E 1 = {(0, 1)} and E 2 = {(1, 0)}, and take ℓ = (0, 1, 3). Each dataset is split into prior and certification sets. We take H to be two-layer MLPs. As is common in the PAC-Bayes literature, we restrict P and Q to be isotropic and diagonal Gaussian distributions over the parameter space, respectively. The means of P and Q are set to the parameters of an MLP trained on the prior set. The mean and variances of Q and the variance of P are tuned via Theorem 2 to minimize the bound on the total risk R T D (Q). See Appendix A for pseudocode, Appendix B for full experimental details and https://github.com/reubenadams/PAC-Bayes-Control for code. The results for MNIST can be seen in Figure 1 .We estimate R S (Q) with a Monte Carlo and obtain a PAC-Bayes bound on R T D (Q) by combining Proposition 2 (with δ = 0.01 and N = 100000) and Proposition 3. We obtain R T D (Q) ≤ 0.2640 for MNIST and R T D (Q) ≤ 0.8379 for HAM10000, where both bounds hold with probability at least 1 -0.05 -0.01 = 0.94. While these bounds are far from vacuous-the maximum possible value of R T D (Q) is 3 for our choice of ℓ-one might wonder whether one can do better by bounding each error probability individually using Maurer's inequality Maurer [2004] , and then unioning these bounds. As with our Theorem 1, this would also constrain the entire distribution of error types since for any ℓ, one could then calculate the maximimum value of R T D (Q) that satisfies all of these constraints. Both Dataset Volume Our Region Volume Maurer Region MNIST 0.0025 (0.002498, 0.002504) 0.0028 (0.002793, 0.002800) HAM10000 0.0012 (0.001207, 0.001211) 0.0011 (0.001142, 0.001146) Table 1 : Point estimates and 95% confidence intervals for the volumes of the confidence regions for R D (Q) given by Theorem 1 and a union over M individual Maurer bounds, respectively. Our method is superior for MNIST and inferior for HAM100000. methods constrain the region of the simplex in which R D (Q) can lie (with high probability), and a reasonable metric by which to compare them is the volumes of these regions. This can be estimated via a MC sample by uniformly sampling points r from △ M and counting how samples are legal values of R D (Q) according to each method. The 95% confidence intervals for the volumes of the two regions are given in Table 1 . A more comprehensive table for synthetic values of R S (Q) can be found in Appendix B. 
perspectives: 
 We introduce the framework of error types, considering the vectors R S (Q) and R D (Q) of empirical and true probabilities of errors of different types. We prove a PAC-Bayes bound (Theorem 1) on kl(R S (Q)∥R D (Q)) which controls the entire distribution of error probabilities, and hence can be used to derive bounds on arbitrary linear combinations of the error probabilities, all of which hold simultaneously with high probability; this cannot be achieved with any existing PAC-Bayes bound.We construct a differential training objective based on our bound by introducing the the vectorised kl inverse, providing a recipe for quickly computing its value and derivatives (Theorem 2). Our framework is flexible enough to encompass multiclass classification or discretised regression, but also structured output prediction, multi-task learning and learning-to-learn.Another potential application of our work is to the excess risk, since under a misclassification loss there are three different error types, corresponding to excess losses of {-1, 0, 1}. Biggs and Guedj [2023] adapted Theorems 1 and 2 to this setting, leading to an empirically tighter PAC-Bayes bound for certain classification tasks.We require i.i.d. data, which in practice is frequently not the case or is hard to verify. Further, the number of error types M must be finite. In continuous scenarios it would be preferable to be able to control the entire distribution of loss values without having to discretise into finitely many error types.We leave this direction to future work.We here outline more explicitly how Theorem 1 and Theorem 2 may be used to formulate a fully differentiable objective by which a model may be trained.First, if one wishes to make hard labels, namely H ⊆ Y X , it will first be necessary to use a surrogate class of soft hypotheses H ′ ⊆ M(Y) X during training, before reverting to hard labels for example by taking the mean label or the one with highest probability. Using soft hypotheses during training is necessary to ensure that the empirical j-risks R j S (Q) are differentiable with respect to the model parameters. Since how one chooses to do this will depend on the specific use case, we restrict our attention here to the case of soft hypotheses. Specifically, we consider a class of soft hypotheses H = {h θ : θ ∈ R N } ⊆ M(Y) X parameterised by the weights θ ∈ R N of some neural network of a given architecture with N parameters in such a way that the R j S (h θ ) are differentiable in θ. A concrete example would be multiclass classification using a fully connected neural network with output being softmax probabilities on the classes so that the R j S (h θ ) are differentiable. Second, it is necessary to restrict the prior and posterior P, Q ∈ M(H) to a parameterised subset of M(H) in which KL(Q∥P ) has a closed form which is differentiable in the parameterisation. A simple choice for our case of a neural network with N parameters isP, Q ∈ {N (w, diag(s)) : w ∈ R N , s ∈ R N >0 }. For prior a P v,r = N (v, diag(r)) and posterior Q w,s = N (w, diag(s)) we have the closed form KL(Q w,s ∥P v,r ) = 1 2 N n=1 s n r n + (w n -v n ) 2 r n + ln r n s n -N ,which is indeed differentiable in v, r, w and s. While Q w,s and P v,r are technically distributions on R D rather than H, the KL-divergence between the distributions they induce on H will be at most as large as the expression above. Thus, substituting the expression above into the bounds we prove in Section 3 can only increase the value of the bounds, meaning the enlarged bounds certainly still hold with probability at least 1 -δ.Third, in all but the simplest cases R j S (Q w,s ) will not have a closed form, much less one that is differentiable in w and s. A common solution to this is to use the so-called pathwise gradient estimator. In our case, this corresponds to drawing ϵ ∼ N (0, I), where I is the N × N identity matrix, and estimating∇ w,s R j S (Q w,s ) = ∇ w,s E ϵ ′ ∼N (0,I) R j S (h w+ϵ ′ ⊙ √ s ) ≈ ∇ w,s R j S (h w+ϵ⊙ √ s ),where h w denotes the function expressed by the neural network with parameters w. For a proof that this is an unbiased estimator, and for other methods for estimating the gradients of expectations, see the survey Mohamed et al. [2020] .Fourth, one must choose the prior. Designing priors which are optimal in some sense (i.e., minimising the Kullback-Leibler term in the right-hand side of generalisation bounds) has been at the core of an active line of work in the PAC-Bayesian literature. For the sake of simplicity, and since it is out of the scope of our contributions, we assume here that the prior is given beforehand, although we stress that practitioners should pay great attention to its tuning. For our purposes, it suffices to say that if one is using a data-dependent prior then it is necessary to partition the sample into S = S Prior ∪ S Bound , where S Prior is used to train the prior and S Bound is used to evaluate the bound. Since our bound holds uniformly over posteriors Q ∈ M(H), the entire sample S is free to be used to train the posterior Q. For a more in-depth discussion on the choice of prior, we refer to the following body of work: Ambroladze et al. [2006] , Lever et al. [2010 Lever et al. [ , 2013]] Biggs and Guedj [2022a,b] .Finally, given a confidence level δ ∈ (0, 1], one may use Algorithm 1 to obtain a posterior Q w,s with minimal upper bound on the total risk. Note we take the pointwise logarithm of the variances r and s to obtain unbounded parameters on which to perform stochastic gradient descent or some other minimisation algorithm. We use ⊕ to denote vector concatenation. The algorithm can be straightforwardly adapted to permit mini-batches by, for each epoch, sequentially repeating the steps with S equal to each mini-batch. / P v,r , v(S Prior ) ∈ R N , r(S Prior ) ∈ R N ≥0 /* A (data-dependent) prior */ Q w0,s0 , w 0 ∈ R N , s 0 ∈ R N ≥0 /* An initial posterior */ δ ∈ (0, 1] /* A confidence level */ λ > 0 /* Au ← R S h w+ϵ⊙ √ exp(ζ) B ← 1 m KL Q w,exp(ζ) P v,r + ln 1 δ √ πe 1/12m m 2 M -1 2 M -1 z=0 M z 1 (πm) z/2 Γ( M -z 2 ) ũ ← (u 1 , . . . , u M , B) G ← 0 2N ×(M +1) /* Initialise gradient matrix */ F ← 0 M +1 /* Initialise gradient vector */ for j ← 1 to M + 1 do F j ← ∂f * ℓ ∂ ũj ( ũ) /*For MNIST we map labels {0, 1, 2, 3, 4} to 0 and {5, 6, 7, 8, 9} to 1. For HAM10000 we map the cancerous or pre-cancerous labels {Melanoma, Basal Cell Carcinoma, Actinic Keratosis} to 1 and the other labels to 0. In both cases we partition Y 2 into E 0 = {(0, 0), (1, 1)}, E 1 = {(0, 1)} and E 2 = {(1, 0)}, and take ℓ = (0, 1, 3). For HAM10000, E 1 and E 2 then refer to Type I and Type II errors, respectively, and ℓ reflects the greater severity of false negatives.Each dataset is split into prior and certification sets S Prior and S Bound , respectively. For MNIST, we use the conventional training set of size 60000 as the prior set, and the conventional test set of size 10000 as the certification set. For HAM10000 we pool the conventional train, validation and test sets together and then split 50-50 to obtain prior and certification sets each of size 5860. For HAM10000we resize the images to (28, 28) and use just the first channel so that the data dimension is the same for both datasets.We take H to be two-layer MLPs with 784, 100 and 2 units in the input, hidden and output layers, respectively. As is common in the PAC-Bayes literature, we restrict P to be an isotropic Gaussian N (v, λI) and Q to be a diagonal Gaussian N (w, diag(s)). Further, as in Dziugaite and Roy [2017] , we restrict λ to be of the form λ j = c exp(-j/b) for some j ∈ N, taking c = 0.1 and b = 100. Since, at the end of training, we will then have one prior P j for each j ∈ N, we can choose the j that minimizes the PAC-Bayes bound provided we take a union over all of them, taking δ j = 6δ π 2 j 2 so that j δ j = 1 and all the bounds hold simultaneously with probability at least 1 -δ. After applying Algorithm 1 we round λ to a discrete λ j , either up or down depending on which gives the smaller bound.For both datasets we set the prior mean v to be the parameters of an MLP trained on the prior set. In both cases we use SGD with learning rate 0.01 to minimise the cross-entropy loss, using a portion of the prior set as a validation set. For MNIST we train the MLP for 20 epochs to get an error rate of 14%, for HAM10000 we train the MLP for 5 epochs to get an error rate of 22%. We then apply Algorithm 1. By combining Proposition 2 (with δ = 0.01 and N = 100000) and Proposition 3. We obtain R S ( Q) = (0.8879, 0.0919, 0.0203) and R T D (Q) ≤ 0.2640 for MNIST and R S ( Q) = (0.7860, 0.0146, 0.1995) and R T D (Q) ≤ 0.8379 for HAM10000, where both bounds hold with probability at least 1 -0.05 -0.01 = 0.94.The full results are shown in Figure 2 . Figures 2a, 2c and 2e are the same as Figures 1a, 1b and 1c , and are repeated here for easier comparison with the HAM10000 results. Figure 2b shows that Algorithm 1 has failed to reduce the bound on the total risk beyond the initialisation of Q to P , with the small variation being explained by different MC samples being drawn from Q during training rather than Q changing substantially. Indeed, Figure 2h shows that Q does not appreciably move from its initialisation at P -KL(Q∥P ) remains below 0.1 whereas in the MNNIST experiment, which has the same number of parameters, exceeds 30. It is therefore unsurprising that Figures 2d and 2f show negligible change in the empirical error probabilities and the bound on kl(R S (Q)∥R D (Q)), respectively. The divergence in the results is likely due to the difference in sample size; the certification set for the MNIST experiment contains 10000 samples, whereas for the HAM10000 dataset there are only 5000, which, all else equal, makes an increase in KL(Q∥P ) twice as expensive.Recall from Section 6 that while R D (Q) can be effectively constrained to a sub-region of the simple △ M using our Theorem 1, this can also be achieved by unioning M Maurer bounds, one for each error probability. Table 1 gave the 95% confidence intervals for the volumes of the confidence regions in which R D (Q) was likely to lie for experiments on MNIST and HAM10000, but neither region was uniformly smaller, making it unclear which method should be preferred.Table 2 provides additional data by taking synthetic values for R S (Q) and KL(Q∥P ), for different values of m (the size of the certification set) and M (the number of error types). 'Individual' denotes unioning individual Maurer bounds, 'Ours' is our method, 'Intersection' is the intersection of the confidence regions given by the previous two methods (but loosened so that they now both hold simultaneously with probability at least 0.95), and 'Morv.' is the confidence region produced by Morvant's bound Morvant et al. [2012] . The 95% confidence intervals for the volumes of all the regions have been produced by Monte Carlo samples. We see that our confidence region is tighter than the individual one in 4/9 cases (green), worse in 3/9 cases (red) and ties in 2/9 cases (orange). Interestingly, union bounding the naive CR and our CR and intersecting often beats both of these (bold). Morvant's result is either not applicable or their confidence region is much larger than ours and essentially takes up the entire simplex, hence the volume estimate of 1.000. The reason their bound is sometimes inapplicable is because it requires every class to contain at least 8L instances, where L is the number of labels-in the L = 5, M = 25, m = 100 case this would require each class to contain at least 5 × 8 = 40 instances which is impossible with m = 100 samples. Now suppose that equation ( 16) holds for any x (m) ∈ △ m M and let x (m+1) = (x 1 , . . . , x m+1 ) ∈ △ m+1 M . Then the left hand side of equation ( 16) can be written asη (m+1) ∈E m+1 M m+1 i=1 x i • η i = η (m) ∈E m M M j=1 m i=1 x i • η i (x m+1 • e j ) = η (m) ∈E m M m i=1 x i • η i M j=1 (x m+1 • e j ) = 1.We now show that any x (m) = (x 1 , . . . , x m ) ∈ △ m M can be written as a convex combination of the elements of E m M in the following wayEQUATION)We have already shown that the weights sum to one, and they are clearly elements of [0, 1], so the right hand side of equation ( 17) is indeed a convex combination of the elements of E m M . We now show that equation ( 17) holds, again by induction.For m = 1 and any x (1) = (x 1 , ) ∈ △ 1 M , the right hand side of equation ( 17) can be written asM j=1 (x 1 • e j )(e j , ) = (x 1 , ) = x.For the inductive hypothesis, suppose equation ( 17) holds for some arbitrary m ≥ 1, and denote elements of E m+1 M by η (m) ⊕ (e, ) for some η (m) ∈ E m M and e ∈ E M , where ⊕ denotes vector concatenation. Then for any x (m+1) = x (m) ⊕ (x m+1 , ) = (x 1 , . . . , x m+1 ) ∈ △ m+1 M , the right hand side of equation ( 17) can be written as m+1) , where in the penultimate equality we have used the inductive hypothesis and (twice) the result of the previous induction.η (m+1) ∈E m+1 M m+1 i=1 x i • η i η (m+1) = η (m) ∈E m M M j=1 m i=1 x i • η i (x m+1 • e j )η (m) ⊕ (e j , ) = η (m) ∈E m M M j=1 m i=1 x i • η i (x m+1 • e j )η (m) ⊕ η (m) ∈E m M M j=1 m i=1 x i • η i (x m+1 • e j )(e j , ) = M j=1 (x m+1 • e j ) η (m) ∈E m M m i=1 x i • η i η (m) ⊕ η (m) ∈E m M m i=1 x i • η i M j=1 (x m+1 • e j )(e j , ) = 1 • x (m) ⊕ 1 • (x m+1 , ) = x (We can now prove the statement of the Lemma. Applying Jensen's inequality to equation ( 17) with the convex function f , we have thatf (x 1 , . . . , x m ) = f   η (m) ∈E m M m i=1 x i • η i η (m)   ≤ η (m) ∈E m M m i=1 x i • η i f η (m) .Let µ = E[X 1 ] denote the mean of the i.i.d. random vectors X i . Then the above inequality impliesE[f (X 1 , . . . , X m )] ≤ η (m) ∈E m M m i=1 µ • η i f η (m) = η (m) ∈E m M m i=1 P(X ′ i = η i ) f η (m) = E[f (X ′ 1 , . . . , X ′ m )].The proof of Lemma 3 itself requires two technical helping lemmas which we now state and prove. Lemma 4. For any integers n ≥ 2 and p ≥ -1,n-1 k=1 (n -k) p/2 √ k ≤ n p+1 2 1 0 (1 -x) p/2 √ x dx.Proof. The case of p = -1, namelyn-1 k=1 1 k(n -k) ≤ 1 0 1 x(1 -x) dx,has already been demonstrated in Maurer [2004] . For p > -1, letf p (x) := (1 -x) p/2 √ x .We will show that each f p (•) is monotonically decreasing on (0, 1). Indeed,df p dx (x) = - (1 -x) p 2 -1 (px + 1 -x) 2x 3/2 ≤ - (1 -x) p/22x 3/2 < 0, where for the inequalities we have used the fact that p > -1 and x ∈ (0, 1). We therefore see thatn-1 k=1 (n -k) p/2 √ k = n-1 k=1 n p/2 (1 -k n ) p/2 √ n k n = n p+1 2 n-1 k=1 1 n (1 -k n ) p/2 k n = n p+1 2 n-1 k=1 1 n f p k n ≤ n p+1 2 n-1 k=1 k n k-1 n f p (x)dx = n p+1 2 1-1 n 0 f p (x)dx ≤ n p+1 2 1 0 f p (x)dx.Intuitively, the proof of the above lemma works by bounding the integral below by a Riemann sum.In the following lemma we actually calculate this integral, yielding a more explicit bound on the sum in Lemma 4. We found it is easier to calculate a slightly more general integral, where the 1 in the limit and the integrand is replaced by a positive constant a.Lemma 5. For any real number a > 0 and integer n ≥ -1,a 0 (a -x) n/2 √ x dx = √ π Γ( n+2 2 ) Γ( n+3 2 ) a n+1 2 . Proof. Define I n (a) := a 0 (a -x) n/2 √ x dx and f n (a) := √ π Γ( n+2 2 ) Γ( n+3 2 ) a n+1 2 .We proceed by induction, increasing n by 2 each time. This means we need two base cases. First, for n = -1, we haveI -1 (a) = a 0 1 x(a -x) dx = 2 arcsin x a a 0 = π = f -1 (a), since Γ( 1 2 ) = √ π and Γ(1) = 1. Second, for n = 0, I 0 (a) = a 0 1 √ x dx = 2 √ x a 0 = 2 √ a = f 0 (a), since Γ( 3 2 ) = √ π2 . Now, by the Leibniz integral rule, we haved da I n+2 (a) = a 0 ∂ ∂a (a -x) n+2 2 √ x dx = n + 2 2 a 0 (a -x) n 2 √ x dx = n + 2 2 I n (a). Thus I n+2 (a) = n + 2 2 a 0 I n (t)dt + I n (0) = n + 2 2 a 0 I n (t)dt, since I n (0) = 0.Now, for the inductive step, suppose I n (a) = f n (a) for some n ≥ -1. Then, using the previous calculation, we haveI n+2 (a) = n + 2 2 a 0 f n (t)dt = n + 2 2 a 0 √ π Γ( n+2 2 ) Γ( n+3 2 ) t n+1 2 dt = √ π n+2 2 Γ( n+2 2 ) n+3 2 Γ( n+3 2 ) a n+3 2 = √ π Γ( n+2 2 + 1) Γ( n+3 2 + 1) a n+3 2 = √ π Γ (n+2)+2 2 Γ (n+2)+3 2 a (n+2)+1 2 = f n+2 (a).This completes the proof.We are now ready to prove Lemma 3 which, for ease of reference, we restate here. For integersM ≥ 1 and m ≥ M , k∈S >0 m,M 1 M j=1 k j ≤ π M 2 m M -2 2 Γ( M 2 ).Proof. (of Lemma 3) We proceed by induction on M . For M = 1, the set S m,M contains a single element, namely the one-dimensional vector k = (k 1 , ) = (m, ). In this case, the left hand side is 1/ √ m while the right hand side is√ π/( √ mΓ(1/2)) = 1/ √ m, since Γ(1/2) = √ π.Now, as the inductive hypothesis, assume the inequality of Lemma 3 holds for some fixed M ≥ 1 and all m ≥ M . Then for all m ≥ M + 1, we havek∈S >0 m,M +1 1 M +1 j=1 k j = m-M k1=1 1 √ k 1 k ′ ∈S >0 m-k 1 ,M 1 M j=1 k ′ j ≤ m-M k1=1 1 √ k 1 π M 2 (m -k 1 ) M -2 2 Γ( M 2 )(by the inductive hypothesis)= π M 2 Γ( M 2 ) m-M k1=1 (m -k 1 ) M -2 2 √ k 1 ≤ π M 2 Γ( M 2 ) m-1 k1=1 (m -k 1 ) M -2 2 √ k 1 (enlarging the sum domain) ≤ π M 2 Γ( M 2 ) m M -1 2 1 0 (1 -x) M -2 2 √ x dx (by Lemma 4) = π M 2 Γ( M 2 ) m M -1 2 √ π Γ( M 2 ) Γ( M +1 2 ) (by Lemma 5) = π M +1 2 m M -1 2 Γ( M +1 2 ), as required.Proof. The case where q j = 1 or p j = 1 can be dealt with trivially by splitting into the three following subcases• q j = p j = 1 =⇒ kl(q j ∥p j ) = kl(q∥p) = 0• q j = 1, p j ̸ = 1 =⇒ kl(q j ∥p j ) = kl(q∥p) = -log p j• q j ̸ = 1, p j = 1 =⇒ kl(q j ∥p j ) = kl(q∥p) = ∞.For q j ̸ = 1 and p j ̸ = 1 define the distributions q, p ∈ △ M by qj = pj = 0 andqi = q i 1 -q j and pi = p i 1 -p j for i ̸ = j. Then i̸ =j q i log q i p i = i̸ =j (1 -q j )q i log (1 -q j )q i (1 -p j )p i = (1 -q j ) i̸ =j qi log qi pi + qi log 1 -q j 1 -p j = (1 -q j )kl( q∥ p) + (1 -q j ) log 1 -q j 1 -p j ≥ (1 -q j ) log 1 -q j 1 -p j .The final inequality holds since kl( q∥ p) ≥ 0. Further, note that we have equality if and only if q = p, which, by their definitions, translates top i = 1 -p j 1 -q j q ifor all i ̸ = j. If we now add q j log qj pj to both sides, we obtain kl(q∥p) ≥ (1 -q j ) log 1 -q j 1 -p j + q j log q j p j = kl(q j ∥p j ),with the same condition for equality.The following proposition makes more precise the argument found at the beginning of Section 4 for how Proposition 1 can be used to derive the tightest possible lower and upper bounds on each R j D (Q). Proposition 5. Suppose that q, p ∈ △ M are such that kl(q∥p) ≤ B, where q is known and p is unknown. Then, in the absence of any further information, the tightest bound that can be obtained on each p j is p j ≤ kl -1 (q j , B).Proof. Suppose p j > kl -1 (q j , B). Then, by definition of kl -1 , we have that kl(q j ∥p j ) > B.By Proposition 1, this would then imply kl(q∥p) > B, contradicting our assumption. Therefore p j ≤ kl -1 (q j , B). Now, with the information we have, we cannot rule out thatp i = 1 -p j 1 -q j q ifor all i ̸ = j and thus, by Proposition 1, that kl(q j ∥p j ) = kl(q∥p). Further, we cannot rule out that kl(q∥p) = B. Thus, it is possible that kl(q j ∥p j ) = B, in which case p j = kl -1 (q j , B). We therefore see that kl -1 (q j , B) is the tightest possible upper bound on p j , for each j ∈ [M ].Before proving the proposition, we first argue that kl -1 ℓ (u|c) given by Definition 1 is well-defined. First, note that A u := {v ∈ △ M : kl(u∥v) ≤ c} is compact (boundedness is clear and it is closed because it is the preimage of the closed set [0, c] under the continuous map v → kl(u∥v)) and so the continuous function f ℓ achieves its supremum on A u . Further, note that A u is a convex subset of △ M (because the map v → kl(u∥v) is convex) and f ℓ is linear, so the supremum of f ℓ over A u is achieved and is located on the boundary of A u . This means we can replace the inequality constraint kl(u∥v) ≤ c in Definition 1 with the equality constraint kl(u∥v) = c. Finally, if u ∈ △ >0 M then A u is a strictly convex subset of △ M (because the map v → kl(u∥v) is then strictly convex) and so the supremum of f ℓ occurs at a unique point on the boundary of A u . In other words, if u ∈ △ >0 M then kl -1 ℓ (u|c) is defined uniquely. We now prove Theorem 2. While our proof technique is somewhat analogous to the technique used in Clerico et al. [2022] to obtain derivatives of the one-dimensional kl-inverse, our theorem directly yields derivatives on the total risk by (implicitly) employing the envelope theorem (see for example Takayama and Akira, 1985) .Proof Outline: We first derive the expression given for v * ( ũ) = kl -1 ℓ (u|c) given on line (5) of the theorem using the method of Lagrange multipliers. Since we are working on the simplex, we make things easier for ourselves by first making the substitution t j = ln v j to make the v j > 0 constraints unnecessary. The method of Lagrange multipliers yields both the maximum and the minimum (recall that kl -1 ℓ (u|c) is defined as the location of a maximum) for the two values of the Lagrange multiplier µ. We show that exactly one of these values lies in the interval µ ∈ (-∞, -max j ℓ j ) and that this one corresponds to the maximum. This shows that the value µ * Theorem 2 instructs us to find indeed yields v * ( ũ) = kl -1 ℓ (u|c). Finally, we derive the partial derivatives of kl -1 ℓ (u|c) with respect the ũj to obtain the second part of the theorem, namely line (6) by employing the envelope theorem.Proof. (of Theorem 2) We start by deriving the implicit expression for v * ( ũ) = kl -1 ℓ (u|c) given in the proposition by solving a transformed version of the optimisation problem given by Definition 1 using the method of Lagrange multipliers. We obtain two solutions to the Lagrangian equations, which must correspond to the maximum and minimum total risk over the set A u := {v ∈ △ M : kl(u∥v) ≤ c} because, as argued in the main text (see the discussion after Definition 1), A u is compact and so the linear total risk f ℓ (v) attains its maximum and minimum on A u .By definition of v * ( ũ) = kl -1 ℓ (u|c), we know that kl(v * ( ũ)∥u) ≤ c. Since, by assumption, u j > 0 for all j, we see that v * ( ũ) j > 0 for all j, otherwise we would have kl(v * ( ũ)∥u) = ∞, a contradiction. Thus v * ( ũ) ∈ △ >0 M and we are permitted to instead optimise over the unbounded variable t ∈ R M , where t j := ln v j . With this transformation, the constraint v ∈ △ M can be replaced simply by j e tj = 1 and the optimisation problem where e t ∈ R M is defined by (e t ) j := e tj . Note that F (t) = f ℓ (e t ). Following the terminology of mathematical economics, we call the t j the optimisation variables, and the ũj (namely the u j and c) the choice variables. The vector ℓ is considered fixed-we neither want to optimise over it nor differentiate with respect to it-which is why we occasionally suppress it from the notation henceforth.For each ũ, let v * ( ũ) and t * ( ũ) be the solutions to the original and transformed optimisation problems respectively. Since the map v = e t is one-to-one, it is clear that since v * ( ũ) exists uniquely, so does t * ( ũ), and that they are related by v * ( ũ) = e t * ( ũ) . We therefore have the identity f ℓ (v * ( ũ)) ≡ F (t * ( ũ)).Recalling that f * ℓ ( ũ) := f ℓ (v * ( ũ)), we see that ∇ ũf * ℓ ( ũ) ≡ ∇ ũF (t * ( ũ)).(18) the derivatives of f ℓ (kl -1 ℓ (u|c)) with respect to u and c are given by ∇ ũF (t * ( ũ)). Using the method of Lagrange multipliers, there exist real numbers λ * = λ * ( ũ) and µ * = µ * ( ũ) such that (t * , λ * , µ * ) is a stationary point (with respect to t, λ and µ) of the Lagrangian function L(t, λ, µ; ũ) := F (t) + λg(t; ũ) + µh(t).Let F t (•) and h t (•) denote the gradient vectors of F and h respectively, and let g t ( • ; ũ) and g ũ(t; • ) denote the gradient vectors of g with respect to t only and ũ only, respectively. Simple calculation yields Then, taking the partial derivatives of L with respect to λ, µ and the t j , we have that (t, λ, µ) = (t * ( ũ), λ * ( ũ), µ * ( ũ)) solves the simultaneous equations F t (t) + λg t (t; ũ) + µh t (t) = 0, (20) g(t; ũ) = 0, and h(t) = 0, where the last two equations recover the constraints. Substituting the gradients F t , g t and h t , the first equation reduces to ℓ ⊙ e t -λu + µe t = 0, First, note that since, for all ũ, the constraints g = h = 0 are satisfied by t * ( ũ), we have the identities g(t * ( ũ), ũ) ≡ 0 and h(t * ( ũ)) ≡ 0.gDifferentiating these identities with respect to ũj then yields g t (t * ( ũ), ũ) • ∂t * ∂ ũj ( ũ) + g ũj (t * ( ũ), ũ) ≡ 0 and h t (t * ( ũ)) • ∂t * ∂ ũj ( ũ) ≡ 0.As before, we can write these M + 1 pairs of equations as the following pair of matrix equations g t (t * ( ũ), ũ)∇ ũt * ( ũ) + g ũ(t * ( ũ), ũ) ≡ 0 and h t (t * ( ũ))∇ ũt * ( ũ) ≡ 0.Multiplying these identities by λ * ( ũ) and µ * ( ũ) respectively, and combining with equation ( 23), yields∇ ũF (t * ( ũ)) = F t (t * ( ũ)) + λ * ( ũ)g t (t * ( ũ), ũ) + µ * ( ũ)h t (t * ( ũ)) ∇ ũt * ( ũ) + λ * ( ũ)g ũ(t * ( ũ), ũ) = λ * ( ũ)g ũ(t * ( ũ), ũ), where the final equality comes from noting that the terms in the large bracket vanish due to equation (20). Recalling the expression for g ũ(t; ũ) given by Equation 19and that v * ( ũ) = exp(t * ( ũ)) we obtain ∇ ũF (t * ( ũ)) = λ * ( ũ) 1 -t * ( ũ) 1 + log u 1 , . . . , 1 -t * ( ũ) M + log u M , -1= λ * ( ũ) 1 + log u 1 v * ( ũ) 1 , . . . , 1 + log u M v * ( ũ) M , -1Finally, recalling Equivalence (18), namely ∇ ũf * ℓ ( ũ) ≡ ∇ ũF (t * ( ũ)), we see that the above expression gives the derivatives 38th Conference on Neural Information Processing Systems (NeurIPS 2024).Code available here: https://github.com/reubenadams/PAC-Bayes-ControlMore precisely, when ∀h ∈ H ∀(x, y) ∈ S ∃j ∈ [M ] such that h(x)[{ŷ ∈ Y : (ŷ, y) ∈ Ej)}] = 1.Incidentally, this argument also shows that there is at most one solution to the first equation in (22) in the range (-minj ℓj, ∞). There indeed exists a unique solution, which corresponds to the minimum total risk, but we do not prove this. 
 
 Appendix:  


Decision: The paper cleverly extends the PAC-Bayesian theory. Even if it builds on previous works, the theoretical contributions are not straightforward and require several non-trivial steps.

Of note, the reviews have led the authors to commit to further experiments, which will enrich the manuscript. 

The authors should cite Wu and Seldin (2022) as mentioned by Reviewer cfTA.  I also enjoin the authors to mention that the multinomial distribution has been used by Lacasse et al. (2006, Theorem 6) to simultaneously bound the so-called *expected disagreement*, *expected joint success* and *expected joint error* of the C-bound. Even if the latter pursued a different goal and is restricted to a multinomial distribution of solely three events, some proof techniques are similar to the current work. This might be more obvious in the detailed treatment of Germain et al. (2015, Section 5.4).

### References
Lacasse, Laviolette, Marchand, Germain, Usunier. PAC-Bayes Bounds for the Risk of the Majority Vote and the Variance of the Gibbs Classifier. NIPS 2006

Germain, Lacasse, Laviolette, Marchand, Roy. Risk bounds for the majority vote: from a PAC-Bayesian analysis to a learning algorithm. JMLR (2015)

Reviews: review: paper_summary: paper_summary: In the standard statistical learning theory control of generalization
error typically means studying deviations of the empirical risk from
the risk (for instance, with high probability over the sample).  In
the language of statistics such control doesn't make a difference
between type I & II errors.  This paper goes beyond as it looks at a
multi-objective risk, i.e. 'deviation' of the vector or M empirical
risks from the corresponding risks.  In the context of the paper, the
vector of empirical risks lives on a probability simples and the paper
shows how to control a KL divergence between such a vector and a risk
probability vector.  Then, the paper goes on and studies a PAC-Bayes
formulation of this problem (Theorem 1).  In section 4 the paper
proposes how to construct a differentiable objective for this problem.
strength_weakness: strength_weakness: The price for having such a multiobjective setting is M log(M),
whereas the standard union bound would lead to a log(M) cost --- note
that there is a (significant) gap.  However, it should be noted that
this would give a bound on a different objective (not kl() between
points on the simplex).  I think that the paper should argue why such
kl() is interesting or at least design analytical instances and/or
some experiments which would demonstrate that there exist instances
where kl() leads to a better/tigther results (perhaps a performance on
a downstream task?).
questions: questions: Not a question really, but I'd be interested to hear authors' thoughts
on the weakness against a union bound.  E.g. how would one approach
this comparison problem / experiment design.
limitations: limitations: Seems to be adequate.
paper_summary: paper_summary: This paper introduces a novel PAC-Bayes bounds that extend classical kl-based bounds to vector-valued losses that can control several error-types simultaneously. The bound is converted into a differentiable minimization objective and details for the practical implementation of the bound are provided in the Appendix.
strength_weakness: strength_weakness: 1- The proof techniques and the main contribution (Proposition 4) are straightforward extensions of scalar kl-bounds to the multidimensional case. I fear this contribution alone might not be significant enough.

2- The experiment in Section 7 is too simple, further empirical evaluation of the bounds (e.g. simultaneous minimization of Type I and Type II errors in more realistic scenarios) would be very beneficial for the impact of the paper.

3- Wu, Y. S., & Seldin, Y. (2022). Split-kl and PAC-Bayes-split-kl inequalities for ternary random variables. Advances in Neural Information Processing Systems, 35, 11369-11381. is significant related work not discussed in the paper.


In general, I don't think this is a bad paper, but the theoretical contributions are not very original and I feel that the paper needs more experimental backup to be a well-rounded contribution.
questions: questions: See weaknesses.
limitations: limitations: The limitations are properly discussed in Section 6.
paper_summary: paper_summary: The notion of a set of "error-types", introduced by this work, is a user-defined partition of the product space of predictions and responses, generalizing well-known summaries of the erring behavior of predictors. This work presents a PAC-Bayes bound on the divergence between and empirical and true distributions over a set of "error types" for a choice of posterior prediction scheme. The phrasing of the bound in terms of KL divergence naturally generates bounds on arbitrary linear combinations of the true error type probabilities, which hold simultaneously. The authors present a training objective founded on the minimization of the worst-case value of a pre-defined weighting of error types consistent with the empirical error-type profile of a given choice of posterior over predictors.
strength_weakness: strength_weakness: One of the main strengths of this work the simultaneous control of arbitrary linear combinations of the error type probabilities. However, the ability of the previous literature to generate some semblance of these guarantees is not particularly well-exposited. Having not previously read any of this line of work, I think an explicit example of how one can construct a bound on a single linear combination might do the paper some good (and might motivate this work past the simultaneous control of all linear combinations, given that the task seems somewhat tricky given a spectral norm bound).
questions: questions: Is there a quick example of to control a single linear combination of confusion matrix probabilities using a previous bound? 
In line 331, what is an "empirically tighter" bound?
limitations: limitations: Yes
paper_summary: paper_summary: The authors suggest an approach for controlling errors of various types (simultaneously) within a PAC-Bayes framework. They derive a high probability bound on the KL divergence between the empirical and distribution risk vectors. This generalizes earlier work in the binary case by Maurer, 2004 and Begin et al 2016. The authors propose a method for using this bound as a (differentiable) training objective via considering a linear combination of the elements of the risk vector and inverting the resulting binary KL divergence.
strength_weakness: strength_weakness: - The framework considered would be better motivated by providing a specific example (ideally with an experiment) indicating the usefulness of the method (and illustrating why existing approaches with a union bound are insufficient).
- The authors included the NeurIPS checklist from last year (and not from this year). 
- The proof of theorem 2 is very long and either an outline/proof sketch should be provided at the beginning, or the proof should be broken into lemmas/propositions to improve readability.
questions: questions: - Is it really not possible to upper bound arbitrary linear combinations of different types of risk with a union bound? I would think that bounding finitely many would place linear constraints on other linear combinations, allowing at least some bounds to be derived. It isn’t immediately obvious to me how sharp the resulting bounds would be, or if indeed this approach leads to non-vacuous bounds, but I’d appreciate the authors commenting on this and adding some discussion on the topic if the answer to my question is that it is possible. 
- How would $\ell$ be chosen in practice when defining the objective function? It seems the whole purpose of the framework is to provide useful bounds simultaneously for many types of predictive errors. But it also seems the choice of $\ell$ will have a large impact on how well a method trained with proposed approach will perform on tasks that weight certain types of errors heavily.
limitations: limitations: - The authors acknowledge limited experiments as a limitation of the current work. I agree that this is a limitation. In particular, I wonder whether similar results can be obtained using held-out data, and whether or not the bounds derived this way are competitive with this approach. In the one experiment provided, half of MNIST was used to define the prior, while have was used for optimization of the bound. It seems possible that a similar breakdown of data between a training and validation set (trying directly via cross-entropy loss) might lead to similar model performance and tighter bounds on the risk (as in the cases considered in Foong et al). It isn’t immediately obvious whether it is easy to directly invert the multinomial CDF to derive tight bounds on various datasets as in the binary case using holdout data. But I’m curious to hear the authors’ thoughts on this approach.
- The authors point out that the bound is limited to discrete cases. This seems reasonable, and there are many discrete problems where the bound might be useful. I think finding a concrete problem (dataset, with a plausible decision depending on balancing error types) where this bound might be useful would be more convincing to address this limitation than extending the work to continuous problems.

Minor points: 

- Line 147: $S_{m,n}^{>0}$ isn’t the interior of $S_{m,n}$ in a topological sense, assuming $S_{m,n}$ is considered with the discrete topology, which seems the natural choice. Otherwise, I don’t know what is meant by interior.
- Line 198: "Different different"

Comments: comment: We are glad we have addressed the important points you raised, and to read that you believe the paper should be accepted. Thank you again for taking the time to write such a thorough review, and for your suggestions and insights which we believe have allowed us to improve the paper!
We are pleased to read that you found the experiment clarifying! Thank you again for taking the time to write a careful review and for your insightful comments, which we believe have improved the work.
We are glad that you found the additional experiment useful! Thank you for raising your score in light of this, and for your suggestions which we believe have improved the work.
We are glad you found the evaluations interesting and are pleased that you are in favour of accepting the paper! We will indeed expand the evaluations in the final version if it is accepted. Thank you for your detailed review and suggestions, which we believe have improved the paper!
Thanks for your explanations, and empirical evaluations which are indeed interesting. I guess, for the final revision it would be interesting to see a more complete empirical suite (even if it's synthetic) to see where the proposed approach "breaks" and the union bound becomes better. Otherwise it is hard to see a full picture.

Overall, I think this is interesting and proofs have some new ideas, so I'm in favor of accepting this paper.
The authors have addressed the primary concerns raised in my review. I have increased my score by a point, as I think the paper should be accepted.
Thank you for your clarifying comments, with my doubts regarding the theoretical contributions solved, having read the rest of the rebuttals and with more realistic experiments, I am happy to increase my rating to 6. Including these extra experiments in the camera-ready version will make for a solid contribution.
Thanks for further elucidation via experiment. I retain my previous score.