Building a stable classifier with the inflated argmax

Authors: jake a soloff
rina  foygel barber
rebecca  willett
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:50:48

Abstract: We propose a new framework for algorithmic stability in the context of multiclass classification. In practice, classification algorithms often operate by first assigning a continuous score (for instance, an estimated probability) to each possible label, then taking the maximizer---i.e., selecting the class that has the highest score. A drawback of this type of approach is that it is inherently unstable, meaning that it is very sensitive to slight perturbations of the training data, since taking the maximizer is discontinuous. Motivated by this challenge, we propose a pipeline for constructing stable classifiers from data, using bagging (i.e., resampling and averaging) to produce stable continuous scores, and then using a stable relaxation of argmax, which we call the "inflated argmax", to convert these scores to a set of candidate labels. The resulting stability guarantee places no distributional assumptions on the data, does not depend on the number of classes or dimensionality of the covariates, and holds for any base classifier. Using a common benchmark data set, we demonstrate that the inflated argmax provides necessary protection against unstable classifiers, without loss of accuracy.

Main paper: 
 introduction: 
 An algorithm that learns from data is considered to be stable if small perturbations of the training data do not lead to large changes in its output. Algorithmic stability is an important consideration in many statistical applications. Within the fairness literature, for instance, stability is one aspect of reliable decision-making systems [FSVC19; HV19] . In interpretable machine learning, it similarly serves as a form of reproducibility [MSKA19; YK20; YB24]. [CMX11] relates stability to robustness, where a machine learning algorithm is robust if two samples with similar feature vectors have similar test error. In the context of generalization bounds, [BE02] and subsequent authors study stability of an algorithm's real-valued output-for instance, values of a regression function. In the setting of a multiclass classification problem, where the data consists of features X ∈ X and labels Y ∈ [L] = {1, . . . , L}, the results of this literature can thus be applied to analyzing a classifier's predicted probabilities pℓ (X) ∈ [0, 1]-i.e., our learned estimates of the conditional probabilities Pr{Y = ℓ | X}, for each ℓ ∈ [L]. However, as we will see, stability of these predicted probabilities by no means implies stability of the predicted label itself, ŷ = argmax ℓ∈[L] pℓ (x)-an arbitrarily small perturbation in pℓ (x) might completely change the predicted label ŷ. The distinction matters: for system trustworthiness, we care about the model's final output, on which users base their decisions.In this paper, we propose a new framework for algorithmic stability in the context of multiclass classification, to define-and achieve-a meaningful notion of stability when the output of the algorithm consists of predicted labels, rather than predicted probabilities. Our work connects to other approaches to allowing for ambiguous classification, including set-valued classification [Gry93; DDB09; SLW19; CDHL21] and conformal inference [PPVG02; VGS05; Lei14; AB23] (we will discuss related work further, below).In supervised classification, we take a data set D = (X i , Y i ) i∈[n] consisting of n observations, and train a classifier that maps from the feature space X to the set [L] of possible labels. 1 Typically, this map is constructed in two stages. First we run some form of regression to learn a map p : X → ∆ L-1 from features to predicted probabilities, with pℓ (x) denoting our learned estimate of the conditional label probability, Pr{Y = ℓ | X = x} (here ∆ L-1 = {w ∈ R L : w i ≥ 0, i w i = 1} denotes the probability simplex in L dimensions). We will write p = A(D), where A denotes the learning algorithm mapping a data set D (of any size) to the map p. Then the second step is to convert the predicted probabilities pℓ (x) to a predicted label, which is typically done by taking the argmax, i.e., ŷ = argmax ℓ pℓ (x) (with some mechanism for breaking ties). This two-stage procedure can be represented in the following diagram:Training data D = (X i , Y i ) i∈[n] Predicted probabilities p(x) ∈ ∆ L-1 Test features x ∈ X Predicted label ŷ ∈ [L] Learning algorithm A Selection (via argmax)When predictions are ambiguous, meaning two or more classes nearly achieve the maximum predicted probability, the selected label becomes unstable and can change based on seemingly inconsequential perturbations to the training data. In other words, the above workflow is fundamentally incompatible with the goal of algorithmic stability. Consequently, in this paper we instead work with set-valued classifiers, which return a set of candidate labels-in practice, this typically leads to a singleton set for examples where we have high confidence in the label, but allows for a larger set in the case of ambiguous examples. While the idea of returning a set of candidate labels is not itself new, we will see that the novelty of our work lies in finding a construction that offers provable stability guarantees.In this framework, a feature vector x ∈ X is now mapped to a set of candidate labels Ŝ ⊆ [L] (rather than a single label ŷ), via a selection rule s, as illustrated here: denotes a selection rule, mapping a vector of estimated probabilities to the set of candidate labels, and ℘ ([L]) denotes the set of subsets of [L] . Of course, the earlier setting-where the procedure returns a single label ŷ = argmax ℓ pℓ (x), rather than a subset-can be viewed as a special case by simply taking s to be the argmax operator.Training data D = (X i , Y i ) i∈[n] Predicted probabilities p(x) ∈ ∆ L-1 Test features x ∈ X Candidate labels Ŝ ⊆ [L]If we instead allow for Ŝ to contain multiple candidate labels, a commonly used version of this framework is given by a top-5 procedure (or top-k, for any fixed k). That is, after running a learning algorithm A to estimate probabilities p(x), the selection rule s then returns the top 5 labels-the labels ℓ 1 , . . . , ℓ 5 ∈ [L] corresponding to the highest 5 values of pℓ (x). This approach is more stable than a standard argmax. However, the choice of 5 in this setting is somewhat arbitrary, and the set Ŝ always contains 5 candidate labels, regardless of the difficulty of the test sample-while intuitively, we might expect that it should be possible to return a smaller Ŝ for test points x that are easier to classify (and a larger Ŝ if x is more ambiguous). In contrast, in our work we seek a more principled approach, where we provide provable stability guarantees while also aiming for the set Ŝ to be as small as possible.In this work, we introduce selection stability, a new definition of algorithmic stability in the context of classification, which focuses on the stability of the predicted label. We reduce the problem of stabilizing a classifier to separately stabilizing the learning and selection stages, described above. For the selection rule s of the two-stage procedure, we propose the inflated argmax operation: Definition 1 (Inflated argmax). For any w ∈ R L , define the inflated argmax asEQUATIONwhere dist(w, R ε j ) = inf v∈R ε j ∥w -v∥, and whereR ε j = w ∈ R L : w j ≥ max ℓ̸ =j w ℓ + ε/ √ 2 .Our procedure will then return the set Ŝ = argmax ε (p(x)) of candidate labels-intuitively, j ∈ argmax ε (p(x)) indicates that a small perturbation of the predicted probabilities, p(x), would lead to the jth label's predicted probability being largest by some margin. In particular, by construction, the inflated argmax will always include any maximal index-that is, if p(x) j is the (possibly non-unique) largest estimated probability, then we must have j ∈ Ŝ = argmax ε (p(x)) (this fact, and other properties of the inflated argmax, will be established formally in Proposition 10 below).In this work, we derive the stabilizing properties of the inflated argmax, and give an algorithm to compute it efficiently. We prove that combining this operation with bagging at the learning step will provably stabilize any classifier. In particular, our guarantee holds with no assumptions on the data, and no constraints on the dimensionality of the covariates nor on the number of classes. 
framework: stable classification: 
 In this section, we propose a definition of algorithmic stability in the setting of multiclass classification.To begin, we formally define a classification algorithm as a map2 C : ∪ n≥0 (X × [L]) n × X -→ ℘ ([L]),which maps a training data set D of any size n, together with a test feature vector x ∈ X , to a candidate set of labels Ŝ = C(D, x). To relate this notation to our earlier terminology, the two-stage selection procedure described in Section 1.1 corresponds to the classification algorithmC(D, x) = s(p(x)) where p = A(D).Abusing notation, we will write this as C = s • A, indicating that C is obtained by applying the selection rule s to the output of the learning algorithm A.We now present our definition of algorithmic stability for a classification algorithm C. As is common in the algorithmic stability literature, we focus on the stability of the algorithm's output with respect to randomly dropping a data point: does the output of C on a test point x change substantially if we drop a single data point from the training data D? If the algorithm's output were a real-valued prediction ŷ ∈ R, we could assess this by measuring the real-valued change in ŷ when a single data point is dropped. For classification, however, we will need to take a different approach:Definition 2 (Selection stability). We say a classification algorithm C has selection stability δ at sample size n if, for all datasets D ∈ (X × [L]) n and all test features x ∈ X ,1 n n i=1 1 Ŝ ∩ Ŝ\i = ∅ ≤ δ,where Ŝ = C(D, x) and where Ŝ\i = C(D \i , x),for each i ∈ [n].Here the notation D \i denotes the data set D with ith data point removed-that is, for a data setD = (X j , Y j ) j∈[n], the leave-one-out data set is given byD \i = (X j , Y j ) j∈[n]\i .Intuitively, selection stability controls the probability that our algorithm makes wholly inconsistent claims when dropping a single data point at random from the training set. If we interpret Ŝ as making the claim that the true label Y is equal to one of the labels in the candidate set Ŝ, and similarly for Ŝ\i , then as soon as there is even one single value that lies in both Ŝ and Ŝ\i , this means that the two claims are not contradictory-even if the sets are large and are mostly nonoverlapping.At first glance, this condition appears to be too mild, in the sense that we require the two prediction sets only to have some way of being compatible, and allows for substantial difference between the sets Ŝ and Ŝ\i . However, since standard classification algorithms always output a single label, they often cannot be said to be stable even in this basic sense. Thus, we can view this definition as providing a minimal notion of stability that we should require any interpretable method to satisfy.Most prior work on algorithmic stability concerns algorithms with continuous outputs-for example, in our notation above, the algorithm A that returns estimated probabilities p(x). With such algorithms, there are more standard tools at our disposal for quantifying and ensuring stability. In this section, we connect classical algorithmic stability to our notion of selection stability (Definition 2). We first recall the following definition.Definition 3 (Tail stability). [SBW24b] A learning algorithm A has tail stability (ε, δ) at sample size n if, for all datasets D of size n and all test features x ∈ X ,1 n n i=1 1 ∥p(x) -p\i (x)∥ ≥ ε ≤ δ,where p = A(D) and p\i = A(D \i ), and where ∥ • ∥ denotes the usual Euclidean norm on R L .This intuitive notion of stability is achieved by many well-known algorithms A, such as nearestneighbor type methods or methods based on bagging or ensembling (as established by [SBW24a; SBW24b]-see Section 3.1 below for details).To construct a classification algorithm using the two-stage procedure outlined in Section 1.1, we need to apply a selection rule s to the output of our learning algorithm A. We might expect that choosing a stable A will lead to selection stability in the resulting classification algorithm-but in fact, this is not the case: even if the learning algorithm A is itself extremely stable in the sense of Definition 3, the classification rule obtained by applying argmax to the output of A can still be extremely unstable. The underlying issue is that argmax is extremely discontinuous-the perturbation ∥p(x) -p\i (x)∥ in the predicted probabilities can be arbitrarily small but still lead to different predicted labels, i.e.,argmax ℓ pℓ (x) ̸ = argmax ℓ p\i ℓ (x).Since combining argmax with a stable learning algorithm A will not suffice, we instead seek a different selection rule s-one that will ensure selection stability (when paired with a stable learning algorithm A). To formalize this aim, we introduce another definition:Definition 4 (ε-compatibility). A selection rule s : ∆ L-1 → ℘ ([L]) is ε-compatible if, for any v, w ∈ ∆ L-1 , ∥v -w∥ < ε =⇒ s(v) ∩ s(w) ̸ = ∅.This notion of ε-compatibility allows us to construct classification algorithms with selection stability.Combining the above definitions leads immediately to the following result:Proposition 5. Let A be a learning algorithm with tail stability (ε, δ) at sample size n, and let s be a selection rule satisfying ε-compatibility. Then the classification algorithm C = s • A has selection stability δ at sample size n.Therefore, by pairing a stable learning algorithm A with a compatible selection rule s, we will automatically ensure selection stability of the resulting classification algorithm C = s • A.Of course, ε-compatibility of the selection rule s might be achieved in a trivial way-for instance, s returns the entire set [L] for any input. As is common in statistical settings (e.g., a tradeoff between Type I error and power, in hypothesis testing problems), our goal is to ensure selection stability while returning an output Ŝ that is as informative as possible. In particular, later we will consider the specific aim of constructing Ŝ to be a singleton set as often as possible.Before describing these two steps in detail, we first present our main theorem that characterizes the selection stability guarantee of the resulting procedure. Given sample size n for the training data set D, the notation A m will denote a bagged version of the base algorithm A, obtained by averaging over subsamples of D comprised of m data points sampled either with replacement ("bootstrapping", commonly run with m = n) or without replacement ("subbagging", commonly run with m = n/2)-see below for details.Theorem 6. Fix any sample size n, any bag size m, and any inflation parameter ε > 0. For any base learning algorithm A, the classification algorithm C = argmax ε • A m , obtained by combining the bagged version of A together with the inflated argmax, satisfies selection stability δ whereEQUATIONwhere p n,m = 1 -(1 -1 n ) m for bootstrapping, and p n,m = m n for subbagging. The guarantee in Theorem 6 holds for any base learning algorithm A, and applies regardless of the complexity of the feature space X , and allows the test feature x ∈ X to be chosen adversarially. The dependence on the number of classes L is mild-in fact, the tail stability parameter δ in (2) differs only by a factor of two for L = 2 versus L = ∞. Of course, the guarantee does depend on the choice of the bag size m. In general, a smaller value of m leads to a stronger stability guarantee (since p n,m increases with m), but this comes at the expense of accuracy since we are training the base algorithm A on subsampled data sets D r of size m (rather than n). For the common choices of m = n for bootstrap or m = n/2 for subbagging, we have pn,m 1-pn,m = O(1) for each.) For any w ∈ R L and any ε < ε ′ , argmax ε (w) ⊆ argmax ε ′ (w). 
bagging classifier weights: 
 In this section, we formally define the construction of a bagged algorithm to recap the tail stability result that our framework leverages. We consider the two most common variants of this ensembling method: bagging (based on bootstrapping training points) and subbagging (based on subsampling). For any data set D ∈ (X × [L]) n , we can define a subsampled data set of size m as follows: for a sequence r = (i 1 , . . . , i m ) ∈ [n] m (which is called a bag), we define the corresponding subsampled data set D r = (X i , Y i ) i∈r ∈ (X × [L]) m . Note that if the bag r contains repeated indices (i.e., i k = i ℓ for some k ̸ = ℓ), then the same data point from the original data set D will appear multiple times in D r .Definition 7 (Bootstrapping or subbagging a base learning algorithm A). Given a data set D ∈ (X × [L]) n and some x ∈ X , return the outputA m (D)(x) = E r [A(D r )(x)],where the expected value is taken with respect to a bag r that is sampled as follows:• Bootstrapping (sometimes simply referred to as bagging) [Bre96a; Bre96b] The following result [SBW24b] ensures tail stability for any bootstrapped or subbagged algorithm: SBW24b] ). For any base learning algorithm A returning outputs in ∆ L-1 , the bagged algorithm A m has tail stability (ε, δ) for any ε, δ > 0 satisfying (2).Theorem 8 ([Computational cost of bagging. In this section, we have worked with the ideal, derandomized version of bagging for simplicity-that is, we assume that the expected value E r [A(D r )(x)] is calculated exactly with respect to the distribution of r. In practice, of course, this is computationally infeasible (since for bootstrapping, say, there are n m possible bags r), and so we typically resort to Monte Carlo sampling to approximate this expected value, definingA m (D)(x) = 1 B B b=1 A(D r b )(x), for B i.i.d. bags r 1 , . . . , r B sampled via bootstrapping or subbagging. See Appendix C for extensions of our main result, Theorem 6, to the finite-B version of the method.We now return to the inflated argmax operator, as defined in Definition 1. This operator, in place of the standard argmax, allows for stability in classification.While argmax ε is defined as an operator on R L , in the context of classification algorithms, we only apply it to vectors w lying in the simplex ∆ L-1 (in particular, in the context of a two-stage classification procedure as in Section 1.1, we will apply the inflated argmax to the vector w = p(x) of estimated probabilities). In Figure 1 , we visualize the inflated argmax applied to the simplex in a setting with three possible labels, L = 3. Each different shaded area corresponds to the region of the simplex ∆ L-1 where argmax ε (w) returns a particular subset.Unlike the standard argmax, the inflated argmax allows us to achieve stability in the context of classification. This is due to the following theorem, which shows that the inflated argmax is εcompatible, as introduced in Definition 4. (The proofs of this result and all properties of the inflated argmax are deferred to Appendix A.)Theorem 9. For any ε > 0 and any v, w ∈ R L , if ∥v-w∥ < ε then argmax ε (v)∩argmax ε (w) ̸ = ∅.In particular, viewing argmax ε as a map on the simplex ∆ L-1 , argmax ε is ε-compatible.To gain a more concrete understanding of this theorem, we can look again at Figure 1 to examine the case L = 3. Theorem 9 ensures that any two regions corresponding to disjoint subsets-e.g., the top corner region corresonding to {2}, and the bottom center region corresponding to {1, 3}-are distance at least ε apart. In particular, the region in the center, corresponding to vectors w that map to the full set {1, 2, 3}, is a curved triangle of constant width ε, also known as a Reuleaux triangle[Reu76; Wei05]. {2} {1, 2, 3} {1} {3} {1, 2} {2, 3} (1, 0, 0) (0, 0, 1) (0, 1, 0) {1, 3} (1, 0, 0) (0, 0, 1) (0, 1, 0) {2} {1} {3}Figure 1 : The left plot illustrates the inflated argmax (1) over the simplex ∆ L-1 when L = 3. The numbers in brackets correspond to the output of the inflated argmax, argmax ε (w), for various points w in the simplex. The right plot shows the same but for the standard argmax, which corresponds to the limit of argmax ε (w) as ε → 0.Proof of Theorem 6. With the above results in place, we can now see that the inflated argmax allows us to achieve selection stability, when combined with a bagged algorithm. In particular, combining Theorem 8, Theorem 9, and Proposition 5 immediately implies our main result, Theorem 6.The following result establishes some natural properties obeyed by the inflated argmax, and also compares to the standard argmax.Proposition 10 (Basic properties of the inflated argmax). Fix any ε > 0. The inflated argmax operator satisfies the following properties:1. (Including the argmax.) For any w ∈ R L , argmax(w) ⊆ argmax ε (w).3 Moreover, ∩ ε>0 argmax ε (w) = argmax(w).∈ R L , if w j ≤ w k , then j ∈ argmax ε (w) ⇒ k ∈ argmax ε (w). 4. (Permutation invariance.) For any v, w ∈ R L , if v = (w σ(1) , . . . , w σ(L) ) for some permuta- tion σ on [L], then j ∈ argmax ε (v) ⇔ σ(j) ∈ argmax ε (w).Next, while we have established that inflated argmax offers favorable stability properties, we have not yet asked whether it can be efficiently computed-in particular, it is not immediately clear how to verify the condition dist(w, R ε j ) < ε in (1), in order to determine whether j ∈ argmax ε (w). The following result offers an efficient algorithm for computing the inflated argmax set.∈ R L and ε > 0. Let w [1] ≥ • • • ≥ w [L]denote the order statistics of w, and definek(w) = max k ∈ [L] : k ℓ=1 (w [ℓ] -w [k] ) 2 + k ℓ=1 (w [ℓ] -w [k] ) 2 ≤ ε 2 . Then argmax ε (w) = j ∈ [L] : w j > ε √ 2 + Â1 (w) -k(w) + 1 ε 2 k(w) + ( Â1 (w)) 2 -Â2 (w) ,where Â1 (w) =w [1] +•••+w [ k(w)] k(w), and Â2 (w) = Finally, thus far we have established that argmax ε enables us to achieve stable classification with a computationally efficient selection rule-but we do not yet know whether argmax ε is optimal, or whether some other choice of s might lead to smaller output sets Ŝ (while still offering assumptionfree stability). For instance, we might consider a fixed-margin rule,w 2 [1] +•••+w 2 [ k(w)] k(w) . 0.0 0.5 1.0 δ 10 -4 10 -3 10 -2 10 -1 10 0 1 N j∈[N ] 1{δ j > δ} Stability comparison with B = 1, 000 bags argmax•A argmax ε •A argmax• Am argmax ε • AmEQUATIONfor which ε-compatibility also holds-might this simpler rule be better than the inflated argmax? Our next result establishes that-under some natural constraints-argmax ε is in fact the optimal choice of selection rule, in the sense of returning a singleton set (i.e., | Ŝ| = 1) as often as possible.Proposition 12 (Optimality of the inflated argmax). Let s : ∆ L-1 → ℘ ([L]) be any selection rule.Suppose s is ε-compatible (Definition 4), permutation invariant (in the sense of Proposition 10), and contains the argmax. Then for any w ∈ ∆ L-1 and any j ∈ [L],s(w) = {j} =⇒ argmax ε (w) = {j}.In other words, for any selection rule s satisfying the assumptions of the proposition (which includes the fixed-margin rule, s ε margin ), if s(w) is a singleton set then so is argmax ε (w). (See also Appendix D for a closer comparison between the inflated argmax and the fixed-margin selection rule given in (3).) 
experiments: 
 In this section, we evaluate our proposed pipeline, combining subbagging with the inflated argmax, with deep learning models and on a common benchmark data set. 4 Data and models. We use Fashion-MNIST [XRV17] , which consists of n = 60, 000 training pairs (X i , Y i ), N = 10, 000 test pairs ( Xj , Ỹj ), and L = 10 classes. For each data point (X, Y ), X is a 28 × 28 grayscale image that pictures a clothing item, and Y ∈ [L] indicates the type of item, e.g., a dress, a coat, etc. The base model we use is a variant of LeNet-5, implemented in PyTorch [PGML19] tutorials as GarmentClassifier(). The base algorithm A trains this classifier using 5 epochs of stochastic gradient descent.Methods and evaluation. We compare four methods:1. The argmax of the base learning algorithm A. 4 Code to fully reproduce the experiment is available at https://github.com/jake-soloff/ stable-argmax-experiments. Training all of the models for this experiment took a total of four hours on 10 CPUs running in parallel on a single computing cluster. 1 : Results on the Fashion MNIST data set. We display the frequency of returning the correct label as a singleton, β correct-single and the average size, β set-size , both of which are defined in (6). Standard errors for these averages are shown in parentheses. The final column shows the worst-case instability over the test set, β max-instability , defined in (5). For each metric, the symbol ↗ indicates that higher values are desirable, while ↘ indicates that lower values are desirable.2. The ε-inflated argmax of the base learning algorithm A with tolerance ε = .05.3. The argmax of the subbagged algorithm A m , with B = 1, 000 bags of size m = n/2.4. The ε-inflated argmax of the subbagged algorithm A m , with B = 1, 000 bags of size m = n/2 and tolerance ε = .05.We evaluate each method based on several metrics. First, to assess selection stability, for each test point j = 1, . . . , N we compute its instabilityEQUATIONwhere i 1 , . . . , i 500 are sampled uniformly without replacement from [n] (i.e., we are sampling 500 leave-one-out models p\i to compare to the model p trained on the full training set). Since our theory concerns worst-case instability over all possible test points, we evaluate the maximum instabilityEQUATIONSecond, to evaluate how accurate each method is, we compute how often the method returns the correct label as a singleton β correct-single and set size β set-size (the number of labels in the candidate set):EQUATIONIdeally we would want a method to return the correct singleton as frequently as possible (a large value of β correct-single ∈ [0, 1] that is close to 1) and small set size (a value of β set-size ≥ 1 that is close to 1).In Figure 2 , we present results for the instability of each method, plotting the survival function of the instability for all test points (δ j ) j∈ [N ] . The standard argmax applied to the base algorithm, argmax •A, has the longest tail, meaning δ j is large for many test points j. The inflated argmax applied to the base algorithm, argmax ε •A, offers only a very small improvement on the stability. By contrast, applying the standard argmax to the subbagged algorithm, argmax • A m , offers a much more substantial improvement, since the red dashed curve is much smaller than both of the solid curves. Still, some of the largest δ j for this method are near 1, meaning for these test points the returned set is sensitive to dropping a single training instance. Combining the inflated argmax with subbagging, argmax ε • A m , offers a dramatic improvement: in this case δ j = 0 for every j = 1, . . . , N .In Table 1 , we present the average measures, β correct-single and β set-size and the worst-case measure β max-instability . For both the base algorithm and the subbagged algorithm, applying the inflated argmax incurs a small cost both in terms of the frequency of returning the correct label as a singleton and the average size. The inflated argmax increases set size only very slightly, but when combined with subbagging, we improve upon the β correct-single of the original method and achieve an extremely high level of stability (β max-instability = 0), while returning a singleton set in the vast majority of cases.We extend our experiment in Appendix E. In particular, we compare the inflated argmax to some existing methods for set-valued classification, including simple thresholding and top-K rules as well as more involved methods, all of which are Bayes rules for various utility functions [MWDH21] . We also evaluate each method using some additional metrics from set-valued classification, including utility-discounted predictive accuracy [ZCM12] . 
discussion: 
 Related work. There is an enormous literature on set-valued classification, so we only reference some of the most relevant works. Much prior work has considered the possibility of a more relaxed argmax to allow for uncertainty [Bri90; BV16; MA16; BMN20; ZLYW23]. The more recent papers in this line of work have focused on producing a sparse set of weights, but none of these works offer a formal stability guarantee for the support of the weights. Our work is the first to propose a version of the argmax that can provably control a notion of stability for the classification setting.Our definition of selection stability relies on a notion of consistency between sets-two sets of candidate labels are consistent if they have at least one common element. This is similar in flavor to conformal classification [Lei14; SLW19], where a set of candidate values is valid if it contains the correct label; this method does not rely on any distributional assumptions, and consequently has been applied to complex and high-dimensional applications such as generative language models [QFSY23] . These frameworks share the motivation of avoiding 'overconfidence in the assignment of a definite label' [HPW18] .Overview, limitations and open questions. We prove a guarantee on the selection stability based on our methodology combining bagging with the inflated argmax. Theorem 6 does not place any assumptions on the learning algorithm nor on the data, including any distributional assumptions. In fact, the training data and test point may be chosen adversarially based on the algorithm, and the output will still be stable. Moreover, we do not assume that the sample size is large: the guarantee holds for any fixed training set size n and improves as n increases. Furthermore, the inflated argmax selection rule ensures that the returned sets of candidate labels are as small as possible (i.e., that there is as little ambiguity as possible about the predicted label for any given test point x).While our theorem does not require assumptions, our method does require bagging the base learning algorithm. The main limitation of our work is that bagging is computationally intensive. However, training different bags is easily parallelizable, which is what allowed us to easily train a convolutional neural network on B = 1, 000 total subsets of the training data. Moreover, while the original definition of bagging uses the conventional bootstrap, where each bag contains as many samples as the original data set, i.e. m = n, in our framework we allow for arbitrary bag size m, which could be much smaller than the sample size n. Massively subsampling the data (m ≪ n) can actually help scale learning algorithms to large data sets [KTSJ14] . Moreover, bagging with m = n can be expensive, but there are still many areas of machine learning where it is used, notably in Random Forests. Finally, our experiments also show that a modest number of bags (B = 1, 000) is all that we really need to start seeing major gains in selection stability.We leave open several avenues for future work. Practitioners may be interested in other forms of discrete stability, which are more strict than requiring one common element between Ŝ and Ŝ\i . For instance, one popular way to measure set similarity is the Jaccard index, given by| Ŝ∩ Ŝ\i | | Ŝ∪ Ŝ\i | .Another open problem is to extend the stability framework and methods studied here to more structured discrete outputs, such as in clustering and variable selection, where nontrivial metrics on the output are more common.In this section, we prove all the results of Section 3.2, establishing the various properties of the inflated argmax. In fact, all the proofs will rely on the following theorem, which gives an alternative characterization of the set argmax ε (w).Theorem 13. Fix ε > 0. Define the map c ε : R L → R as c ε (w) is the unique solution toj∈[L] (w j -c) + 2 + j∈[L] (w j -c) 2 + = ε 2 ,and define the map t ε : R L → R ast ε (w) = c ε (w) + ε/ √ 2 - j∈[L] (w j -c ε (w)) + .Then for any w ∈ R L , it holds thatargmax ε (w) = {j ∈ [L] : w j > t ε (w)}.Moreover, t ε (w) ≤ c ε (w) for all w and all ε, and t ε (w) and c ε (w) are both nondecreasing functions of ε, and both nondecreasing functions of w (i.e., if w ≤ w ′ holds coordinatewise, thenc ε (w) ≤ c ε (w ′ ) and t ε (w) ≤ t ε (w ′ )).With this key result in place, we are now ready to turn to the proofs of the results stated in Section 3.2. Throughout all the proofs below, the functions c ε (w) and t ε (w) are defined as in the statement of Theorem 13.Let w, w ′ ∈ R L be any vectors withargmax ε (w) ∩ argmax ε (w ′ ) = ∅. Define B = {j ∈ [L] : w j > c ε (w)}, B ′ = {j ∈ [L] : w ′ j > c ε (w ′ )}, and note that we must have B ⊆ argmax ε (w) since for j ̸ ∈ argmax ε (w), w j ≤ t ε (w) ≤ c ε (w) by Theorem 13. Similarly, B ′ ⊆ argmax ε (w ′ ). We also must have B (and similarly B ′ ) nonempty by definition of c ε (w).EQUATIONwhere the first step holds since B ∩ B ′ = ∅ due to argmax ε (w) ∩ argmax ε (w ′ ) = ∅. Next we will need a technical lemma.Lemma 14. Let k, k ′ ≥ 1 be integers and let ε > 0. Then for anyEQUATIONthen ∥x + y + r1 k ∥ 2 + ∥x ′ + y ′ + r ′ 1 k ′ ∥ 2 ≥ ε 2 .To apply this lemma, first we let k = |B| ≥ 1 and k ′ = |B ′ | ≥ 1, and definex = w B -c1 B , y = t ′ 1 B -w ′ B , x ′ = w ′ B ′ -c ′ 1 B ′ , y ′ = t1 B ′ -w B ′, and note that all these vectors have nonnegative entries (specifically, x and x ′ are nonnegative by definition of B and B ′ , while y and y ′ are nonnegative by Theorem 13, using the fact that B ∩ s(w ′ ) = B ′ ∩ s(w) = ∅). Then the condition (8) is satisfied by definition of c = c ε (w) and c ′ = c ε (w ′ ). Define also r = ct ′ , r ′ = c ′t. and note thatr + r ′ = (c -t) + (c ′ -t ′ ) = j∈[L] (w j -c) + -ε/ √ 2 + j∈[L] (w ′ j -c ′ ) + -ε/ √ 2 = x ⊤ 1 k + x ′⊤ 1 k ′ -ε √ 2,where the second step holds by definition of t = t ε (w) and t ′ = t ε (w ′ ). Therefore, (9) is also satisfied. Returning to (7), Lemma 14 then implies that ∥ww ′ ∥ 2 ≥ ε 2 , which completes the proof of the theorem.First we verify that argmax(w) ⊆ argmax ε (w). Let j ∈ argmax(w), i.e.,w j = max ℓ w ℓ . Let v = w + e j • ε/ √ 2, where e j = (0, . . . , 0, 1, 0, . . . , 0) is the jth canonical basis vector. Then v ∈ R ε j , and ∥w -v∥ = ε/ √ 2 < ε, so we have dist(w, R ε j ) < ε and therefore j ∈ argmax ε (w). We also need to verify that ∩ ε>0 argmax ε (w) = argmax(w), i.e., for j ̸ ∈ argmax(w), there is some ε > 0 with j ̸ ∈ argmax ε (w). Let k ∈ [L] be an index with w k > w j , and let ε = √ 2(w kw j ). Then for any v ∈ R ε j , we have∥w -v∥ 2 ≥ (w j -v j ) 2 + (w k -v k ) 2 = (w k -ε/ √ 2 -v j ) 2 + (w k -v k ) 2 ≥ inf t∈R t -(ε/ √ 2 + (v j -v k )) 2 + t 2 = 1 2 (ε/ √ 2 + (v j -v k )) 2 ≥ ε 2 ,where the last step holds since v jv k ≥ ε/ √ 2. This means that dist(w, R ε j ) ≥ ε and so j ̸ ∈ argmax ε (w) at this value of ε > 0.Next we check monotonicity in ε. Fix ε < ε ′ . Then for any j ∈ [L],j ∈ argmax ε (ε) ⇐⇒ w j > t ε (w) =⇒ w j > t ε ′ (w) ⇐⇒ j ∈ argmax ε ′ (w),where each step holds by Theorem 13.Next we verify monotonicity in w. If w j ≤ w k , then by Theorem 13j ∈ argmax ε (w) ⇐⇒ w j > t ε (w) =⇒ w k > t ε (w) ⇐⇒ k ∈ argmax ε (w).Finally we check permutation invariance. Suppose v = (w σ(1) , . . . , w σ(L) ) is a permutation of w. Then by construction, we have c ε (v) = c ε (w) and t ε (v) = t ε (w). In particular,j ∈ argmax ε (v) ⇐⇒ v j > t ε (v) ⇐⇒ w σ(j) > t ε (w) ⇐⇒ σ(j) ∈ argmax ε (w).Define a function f w asEQUATIONThen k(w) can equivalently be defined ask(w) = max{k ∈ [L] : f w (w [k] ) ≤ ε 2 }.(Note that this maximum is well defined, since f w (w [1] ) = 0 and so the set is nonempty.) Since c → f w (c) is strictly decreasing over c ≤ w [1] , we see that the solution c ε (w) to the equationf w (c) = ε 2 must satisfy c ε (w) ≤ w [ k(w)] , and (if k(w) < L) also c ε (w) > w [ k(w)+1] .In particular, this impliesε 2 =   j∈[L] (w j -c ε (w)) +   2 + j∈[L] (w j -c ε (w)) 2 + =   k(w) j=1 (w j -c ε (w))   2 + k(w) j=1 (w j -c ε (w)) 2 = k(w) • ( Â1 (w) -c ε (w)) 2 + k(w) Â2 (w) -2 k(w) Â1 (w)c ε (w) + k(w)c ε (w) 2 = k(w)( k(w) + 1) c ε (w) 2 -2c ε (w) Â1 (w) + k(w)( Â1 (w)) 2 + Â2 (w) k(w) + 1 .This is a quadratic function of c ε (w), and is solved byc ε (w) = Â1 - ( Â1 (w)) 2 -Â2 (w) + ε 2 k(w)k(w) + 1 .We also havet ε (w) = c ε (w) + ε/ √ 2 - j∈[L] (w j -c ε (w)) + = c ε (w) + ε/ √ 2 - k(w) j=1 (w [j] -c ε (w)) = ( k(w) + 1)c ε (w) + ε √ 2 -k(w) Â1 (w).Plugging in our above expression for c ε (w), then,t ε (w) = ε √ 2 + Â1 (w) -k(w) + 1 ( Â1 (w)) 2 -Â2 (w) + ε 2 k(w).Since argmax ε (w) = {j : w j > t ε (w)} by Theorem 13, this completes the proof.First we will need a lemma:Lemma 15. Fix w ∈ R L . Then w ∈ R ε j if and only if argmax ε (w) = {j}. For intuition on this result, we can look back at Figure 1 -for instance, the blue region marked by {1} illustrates the set of vectors w ∈ ∆ L-1 such that argmax ε (w) = {1}, and according to this lemma, this is equal to the region R ε 1 . Using this lemma, we now need to show that, for any ε-compatible and permutation invariant selection rule s, if s(w) = {j} for some w ∈ ∆ L-1 and j ∈ [L], then w ∈ R ε j for some j. First, we must have j = argmax(w), since s is assumed to contain the argmax. Fix any k ̸ = j, and define a vector v ∈ ∆ L-1 by permuting w j and w k :v ℓ =    w k , ℓ = j, w j , ℓ = k, w ℓ , ℓ ̸ = j, k.Then by permutation invariance we have s(v) = {k}. Since s is ε-compatible, we therefore haveε ≤ ∥w -v∥ = (w j -v j ) 2 + (w k -v k ) 2 + ℓ̸ =j,k (w ℓ -v ℓ ) 2 = 2(w j -w k ) 2 .Therefore we have w jw k ≥ ε/ √ 2 for all k ̸ = j, which proves w ∈ R ε j and therefore argmax ε (w) = {j}.Step 1: verifying that c ε (w) is well defined. First we check that the solution c ε (w) exists and is unique. Fix w and define the function f w as in (10). Note that f w (c) ≡ 0 for c ≥ max i w i , and c → f w (c) is strictly decreasing over c ≤ max i w i , with lim c→-∞ f w (c) = ∞; therefore, a unique solution c ε (w) to the equation f w (c) = ε 2 must exist.Step 2: verifying t ε (w) ≤ c ε (w). Next we check that the claim t ε (w) ≤ c ε (w) must hold. First, by definition of c ε (w), we haveε 2 =   j∈[L] (w j -c ε (w)) +   2 + j∈[L] (w j -c ε (w)) 2 + ≤ 2   j∈[L] (w j -c ε (w)) +   2 ,where the inequality holds by the properties of the ℓ 1 and ℓ 2 norms. Therefore,j∈[L] (w j -c) + ≥ ε/ √ 2,which verifies thatt ε (w) = c ε (w) + ε/ √ 2 - j∈[L] (w j -c ε (w)) + ≤ c ε (w).Step 3: checking monotonicity. Now we turn to verifying the monotonicity properties of t ε (w) and c ε (w). First we check that these functions are nonincreasing in ε. First, since f w (c) is nonincreasing in c, and c ε (w) is the solution to f w (c) = ε 2 , this immediately implies that c ε (w) is nonincreasing in ε. Next we turn to t ε (w). Definet ′ (c) = c + 1 √ 2   j∈[L] (w j -c) +   2 + j∈[L] (w j -c) 2 + - j∈[L](w jc) + , so that we have t ε (w) = t ′ (c ε (w)) by construction. We can verify that c → t ′ (c) is nondecreasing, and therefore, t ε (w) = t ′ (c ε (w)) ≤ t ′ (c ε ′ (w)) = t ε ′ (w), where the inequality holds since c ε (w) ≤ c ε ′ (w).Next we check that t ε (w) and c ε (w) are nondecreasing functions of w. Fix any w ≤ w ′ (where the inequality is coordinatewise, i.e.,w j ≤ w ′ j for all j ∈ [L]). Since w → f w (c) is a nondecreasing function, we have f w (c) ≤ f w ′ (c) for all c. We therefore have f w ′ (c ε (w)) ≥ f w (c ε (w)) = ε 2 = f w ′ (c ε (w ′ )). Since c → f w ′ (c) is nonincreasing, therefore, c ε (w ′ ) ≥ c ε (w). Next we consider t ε . Let t ′′ ε (c) = c + ε/ √ 2 -j∈[L] (w j -c) + , which is a nondecreasing function of c. Then we have t ε (w) = t ′′ ε (c ε (w)) ≤ t ′′ ε (c ε (w ′ )) = t ε (w ′ ).Step 4: returning to the inflated argmax. Finally, fixing any w ∈ R L , we will prove that argmax ε (w) = {j : w j > t ε (w)}. First choose any j ∈ [L] with w j > t ε (w). We need to verify that j ∈ argmax ε (w). Define v ∈ R L with entriesv k = max{c ε (w) + ε/ √ 2, w j }, k = j, min{c ε (w), w k }, k ̸ = j.By construction, we have v ∈ R ε j . We calculatedist(w, R ε j ) 2 ≤ ∥w -v∥ 2 = w j -max{c ε (w) + ε/ √ 2, w j } 2 + k̸ =j (w k -min{c ε (w), w k }) 2 = c ε (w) + ε/ √ 2 -w j 2 + + k̸ =j (w k -c ε (w)) 2 + < c ε (w) + ε/ √ 2 -t ε (w) 2 + k̸ =j (w k -c ε (w)) 2 + ≤   k∈[L] (w k -c ε (w)) +   2 + k∈[L] (w k -c ε (w)) 2 + = ε 2 ,where the last two steps hold by definition of t ε (w) and c ε (w), while the strict inequality holds since t ε (w) ≤ c ε (w) < c ε (w) + ε/ √ 2 as established above, while w j > t ε (w) by assumption. Therefore j ∈ argmax ε (w). Now we check the converse. For this last step, we will need a lemma. The following result characterizes the projection of w to the set R ε j .Lemma 16. Fix any w ∈ R L and any j ∈ [L]. Then there is a unique a ∈ R satisfyingEQUATION)Moreover, defining the vector v ∈ R L asEQUATIONit holds that v = argmin u∈R ε j ∥w -u∥, i.e., v is the projection of w to the set R ε j .Next suppose w j ≤ t ε (w). We need to verify that j ̸ ∈ argmax ε (w). Let a and v be defined as in Lemma 16 above. We can compare the equation (11) to the definition of t ε (w),t ε (w) = c ε (w) + ε/ √ 2 - k∈[L] (w k -c ε (w)) + = c ε (w) + ε/ √ 2 - k̸ =j (w k -c ε (w)) + ,where the last step holds since w j ≤ t ε (w) by assumption, and t ε (w) ≤ c ε (w) as proved above. Since c → c + ε/ √ 2 -k̸ =j (w kc) + is an increasing function, and w j ≤ t ε (w), this implies a ≤ c ε (w). We thendist(w, R ε j ) 2 = ∥w -v∥ 2 = w j -a + ε/ √ 2 2 + k̸ =j (w k -a ∧ w k ) 2 by (12) = w j -a + ε/ √ 2 2 + k̸ =j (w k -a) 2 + =   k̸ =j (w k -a) +   2 + k̸ =j (w k -a) 2 + by (11) ≥   k̸ =j (w k -c ε (w)) +   2 + k̸ =j (w k -c ε (w)) 2 + since a ≤ c ε (w) =   k∈[L] (w k -c ε (w)) +   2 + k∈[L] (w k -c ε (w)) 2 + since w j ≤ t ε (w) ≤ c ε (w)= ε 2 by definition of c ε (w). Therefore, dist(w, R ε j ) ≥ ε, and so j ̸ ∈ argmax ε (w).A.6.1 Proof of Lemma 14First, since y is constrained to have nonnegative entries,∥x + y + r1 k ∥ 2 ≥ ∥(x + r1 k ) + ∥ 2, where for a vector v = (v 1 , . . . , v L ) ∈ R L , we write (v) + to denote the vector with jth entry given by (v j ) + = max{v j , 0} for each j. The analogous bound holds for ∥x′ + y ′ + r ′ 1 k ′ ∥ 2 . ∥x + y + r1 k ∥ 2 + ∥x ′ + y ′ + r ′ 1 k ′ ∥ 2 ≥ ∥(x + r1 k ) + ∥ 2 + ∥(x ′ + r ′ 1 k ′ ) + ∥ 2 .We now need to show that∥(x + r1 k ) + ∥ 2 + ∥(x ′ + r ′ 1 k ′ ) + ∥ 2 ≥ ε 2 . Next let r = r + r ′ 2 = x ⊤ 1 k + x ′⊤ 1 k ′ -ε √ 2 2 , ∆ = -r + r ′ 2 ,so that we can write r = r -∆, r ′ = r + ∆ for some ∆ ∈ R. We we therefore need to show inf ∆∈R f (∆) ≥ ε 2 , wheref (∆) = ∥(x + (r -∆)1 k ) + ∥ 2 + ∥(x ′ + (r + ∆)1 k ′ ) + ∥ 2 .First, we observe that we can restrict the range of ∆. Specifically, for any∆ ≥ max j∈[L] x j + r, we have f (∆) = ∥(x ′ + (r + ∆)1 k ′ ) + ∥ 2 , which is a nondecreasing function; therefore, inf ∆≥max j∈[L] xj +r f (∆) = f max j∈[L]x j + r , meaning that we do not need to consider values of ∆ beyond this upper bound. Applying a similar argument for a lower bound, we see that from this point on we only need to verify thatf (∆) ≥ ε 2 for -max j∈[k ′ ] x ′ j -r ≤ ∆ ≤ max j∈[L]x j + r.Moreover, for any ∆, we have(x j + (r -∆)) 2 = (x j + (r -∆)) 2 + + (x j + (r -∆)) 2 -≤ (x j + (r -∆)) 2 + + (r -∆) 2for all j, by nonnegativity of x. we must have(x j + (r -∆)) 2 = (x j + (r -∆)) 2 + for at least one j ∈ [L] when ∆ ≤ max j∈[L]x j + r as specified above (i.e., because for j maximizing the entry x j , the value (x + (r -∆)1 k ) j is nonnegative). Therefore,∥(x + (r -∆)1 k ) + ∥ 2 ≥ ∥x + (r -∆)1 k ∥ 2 -(k -1)(r -∆) 2 . Similarly, we can show that ∥(x ′ + (r + ∆)1 k ′ ) + ∥ 2 ≥ ∥x ′ + (r + ∆)1 k ′ ∥ 2 -(k ′ -1)(r + ∆) 2 . We then calculate f (∆) = ∥(x + (r -∆)1 k ) + ∥ 2 + ∥(x ′ + (r + ∆)1 k ′ ) + ∥ 2 ≥ ∥x + (r -∆)1 k ∥ 2 -(k -1)(r -∆) 2 + ∥x ′ + (r + ∆)1 k ′ ∥ 2 -(k ′ -1)(r + ∆) 2 = ∥x∥ 2 + 2(r -∆)x ⊤ 1 k + (r -∆) 2 + ∥x ′ ∥ 2 + 2(r + ∆)x ′⊤ 1 k ′ + (r + ∆) 2 = 2ε 2 -(x ⊤ 1 k ) 2 -(x ′⊤ 1 k ′ ) 2 + 2(r -∆)x ⊤ 1 k + 2(r + ∆)x ′⊤ 1 k ′ + 2r 2 + 2∆ 2, where the last step holds by (8).Writing z = x ⊤ 1 k and z ′ = x ′⊤ 1 k ′ for convenience, we have f (∆) ≥ 2ε 2 -z 2 -z ′2 + 2(r -∆)z + 2(r + ∆)z ′ + 2r 2 + 2∆ 2 = 2ε 2 -z 2 -z ′2 + 2 z + z ′ -ε √ 2 2 -∆ z + 2 z + z ′ -ε √ 2 2 + ∆ z ′ + 2 z + z ′ -ε √ 2 2 2 + 2∆ 2 = ε 2 + 4 z -ε/ √ 2 z ′ -ε/ √ 2 + 2 ∆ - z -z ′ 2 2 ≥ ε 2 + 4 z -ε/ √ 2 z ′ -ε/ √ 2 .where the second step holds because r = r+r ′ 2 = z+z ′ -ε √ 2 2 by (9).Finally, we also have x ⊤ 1 k = ∥x∥ 1 ≥ ∥x∥, where the first step holds since x is nonnegative, and so by assumption (8) we have z= x ⊤ 1 k ≥ ε/ √ 2. Similarly, z ′ ≥ ε/ √ 2. Therefore, 4 z -ε/ √ 2 z ′ -ε/ √ 2 ≥ 0, and so we have shown that f (∆) ≥ ε 2 , as desired.First, suppose w ∈ R ε j . Then argmax(w) = {j}, and so j ∈ argmax ε (w) by Proposition 10. Next, let c = w jε/ √ 2. Then w ℓ ≤ c for all ℓ ̸ = j, and so defining f w as in (10), we havef w (c) = (w j -c) 2 + (w j -c) 2 = ε 2 =⇒ c ε (w) = c = w j -ε/ √ 2. Then t ε (w) = c ε (w) + ε/ √ 2 - ℓ∈[L] (w ℓ -c ε (w)) + = w j -ε/ √ 2 + ε/ √ 2 -ε/ √ 2 = w j -ε/ √ 2,and so for all ℓ ̸ = j, we have w ℓ ≤ w jε/ √ 2 = t ε (w) and thus ℓ ̸ ∈ argmax ε (w). This verifies that argmax ε (w) = {j}.To prove the converse, suppose argmax ε (w) = {j}. By Proposition 10, we then have argmax(w) = {j}. Let k ∈ argmax ℓ̸ =j w ℓ , then to show w ∈ R ε j it suffices to show thatw j ≥ w k + ε/ √ 2. Define v ∈ R L as v ℓ =    w k , ℓ = j, w k + ε/ √ 2, ℓ = k, w ℓ , ℓ ̸ = j, k..Then v ∈ R ε k , and so we must have ∥w -v∥ ≥ ε since k ̸ ∈ argmax ε (w). We calculate ∥w -v∥ 2 = (w j -v j ) 2 + (w k -v k ) 2 = (w j -w k ) 2 + (ε/ √ 2) 2. This implies that we must have (w jw k ) 2 ≥ ε 2 /2, which completes the proof.First we verify existence and uniqueness of a. Let(a) = w j -a + k̸ =j (w k -a) + .Then f : R → R is a continuous and strictly decreasing bijection, so f (a) = ε/ √ 2 must have a unique solution.Next let a be this unique solution and let v be defined as in (12). Now we verify that v is the unique solution to the optimization problemv = argmin u∈R L 1 2 ∥u -w∥ 2 : u j ≥ u k + ε/ √ 2 for all k ̸ = j ,which defines the projection of w to R ε j . By the Karush-Kuhn-Tucker (KKT) conditions for first-order optimality, it is sufficient to verify thatEQUATIONwhere λ k ≥ 0 for all k, and λ k = 0 for any inactive constraints, i.e., any k for whichv j > v k + ε/ √ 2. Let λ k = (w k -a) + for each k ̸ = j. Note that if v j > v k + ε/ √ 2 (i.e., an inactive constraint), then we must have w k < a and so λ k = 0, by construction of v. Then, for anyℓ ̸ = j,   v -w - k̸ =j λ k (e j -e k )   ℓ = v ℓ -w ℓ + λ ℓ = a ∧ w ℓ -w ℓ + (w ℓ -a) + = 0,where e k denotes the kth canonical basis vector in R L . This proves that the ℓth coordinate of the system of equations (13) holds for each ℓ ̸ = j. Now we verify the jth coordinate. We have  v -w - k̸ =j λ k (e j -e k )   j = v j -w j - k̸ =j λ k = a + ε/ √ 2 -w j - k̸ =j (w k -a) + = ε/ √ 2 -f (a) = 0,by our choice of a as the solution to f (a) = ε/ √ 2. This verifies the KKT conditions, and thus, v is the projection of w to R ε j as claimed.As mentioned in Section 2, in many applications it is common to use randomization in the construction of a classification procedure, in the learning algorithm A and/or in the selection rule s. In this section we formalize this more general framework, and will see how our results apply.First, in the non-random setting, a learning algorithm A maps a data set D ∈ (X × [L]) n to a fitted probability estimate function p : X → ∆ L-1 -we write the regression procedure as p = A(D). In the randomized setting, we also allow for external randomness, p = A(D; ξ), where ξ ∼ Uniform[0, 1] is a random seed (e.g., we might use ξ to randomly shuffle the training data when running stochastic gradient descent).Next, in the non-random setting, a selection rule s is a map from ∆ L-1 to subsets of [L], resulting in a candidate set of labels Ŝ = s(p(x)). In the randomized setting (for instance, if s is the argmax but with ties broken at random), we instead include a random seed ζ ∼ Uniform[0, 1], and write Ŝ = s(p(x); ζ). Formally, then, we have s: ∆ L-1 × [0, 1] → ℘ ([L]). (Note that the selection rule proposed in this work-the inflated argmax-is itself not random.)(1, 0, 0) (0, 0, 1) (0, 1, 0) {2} {1} {3} {1, 2, {1, 2} {2, 3} {1, 3} {2} {1, 2, 3} {1} {3} {1, 2} {2, 3}(1, 0, 0) (0, 0, 1) (0, 1, 0) Comparing to our main result for the derandomized case (Theorem 6, which can be interpreted as taking B → ∞), we see that this finite-B result is essentially the same as Theorem 6 once B ≫ n.{1, 3}In this section, we compare to the simpler fixed-margin selection rule, which was defined earlier in (3)-for convenience, we define it again here:s ε margin (w) = j ∈ [L] : w j > max ℓ∈[L] w ℓ -ε/ √ 2 ,the set of all indices j that are within some margin of being the maximum.In this section, we will see that this fixed-margin selection rule is ε-compatible. Since this rule clearly has the advantage of being much simpler than the inflated argmax (both in terms of its definition and interpretability, and in terms of its computation), we might ask whether this rule is perhaps preferable to the more complex inflated argmax. However, we will also see that the fixed-margin selection rule can be very inefficient compared to the inflated argmax: the inflated argmax can never return a larger set, and will often return a substantially smaller one.First, we verify that this simple rule satisfies ε-compatibility (Definition 4).Proposition 19. For any ε > 0, the selection rule s ε margin is ε-compatible.In particular, since s ε margin clearly contains the argmax (i.e., argmax(w) ⊆ s ε margin (w)) and satisfies permutation invariance (in the sense of Proposition 10), by Proposition 12 this immediately implies that s ε margin (w) = {j} =⇒ argmax ε (w) = {j} for all w ∈ R L , i.e., argmax ε is at least as good as s ε margin at returning a singleton set. However, for this particular selection rule, we can state an even stronger result: Proposition 20. For any w ∈ R L and any ε > 0,s ε margin (w) ⊇ argmax ε (w).In particular, this theoretical result ensures that the set of candidate labels Ŝ returned by the inflated argmax can never be larger than for the fixed-margin selection rule. In low dimensions, however, the improvement is small. Indeed, for L = 2, we actually have s ε margin (w) = argmax ε (w) for all w. For L = 3, Figure 3 shows that the inflated argmax is strictly better (i.e., the set inclusion result of Proposition 20 can, for certain values of w, be a strict inclusion), but the difference in this low-dimensional setting appears minor. Next, however, we will see empirically that as the dimension L grows, the benefit of the inflated argmax can be substantial. In this section, we compare the fixed-margin rule to the inflated argmax using randomly generated probability weights. For our simulation, we consider the size of the two sets for various values of L.To sample the weights, we draw a standard Gaussian random vector Z ∼ N (0, I L ) and define w as the softmax of Z, i.e.,w j = e Zj ℓ∈[L] e Z ℓ for j ∈ [L].We then compute argmax ε (w) and s ε margin (w), where ε = 0.1. In Figure 4 , we present results comparing the average size of the sets returned by each of the two methods. In this setting, we see that the inflated argmax has substantial gains over the fixed-margin selection rule when the number of classes L is large. Even when L = 25, the ratio of the expected sizes is about 78%, so the inflated argmax has a nontrivial advantage in that setting. When L = 100, the ratio of expected sizes is 48%, meaning the inflated argmax is on average more than twice as small. D.3 Proofs for the fixed-margin selection ruleD.3.1 Proof of Proposition 19 Let w, v ∈ R L with s ε margin (w) ∩ s ε margin (v) = ∅. Let j ∈ argmax ℓ w ℓ and k ∈ argmax ℓ v ℓ . Then clearly j ∈ s ε margin (w), so since s ε margin (w) ∩ s ε margin (v) = ∅, this implies j ̸ ∈ s ε margin (v). By definition, then, v j ≤ max ℓ v ℓ -ε/ √ 2 = v k -ε/ √ 2.By an identical argument, we havew k ≤ w j -ε/ √ 2. Then ∥w -v∥ 2 ≥ (w j -v j ) 2 + (w k -v k ) 2 ≥ (w j -v j ) 2 + + (v k -w k ) 2 + ≥ w j -(v k -ε/ √ 2) 2 + + v k -(w j -ε/ √ 2) 2 + ≥ inf t∈R (ε/ √ 2 -t) 2 + + (ε/ √ 2 + t) 2 + = 2(ε/ √ 2) 2 = ε 2 .This proves ∥w -v∥ ≥ ε, and therefore, we have shown that ε-compatibility is satisfied.Suppose j ∈ argmax ε (w). Then dist(w, R ε j ) < ε, so we can find some v ∈ R ε j with ∥w -v∥ < ε. Then, for any k ̸ = j,w j -w k = v j -v k + (w j -v j ) + (v k -w k ) ≥ ε/ √ 2 + (w j v j ) + (v k -w k ) since v ∈ R ε j ≥ ε/ √ 2 -|w j -v j | + |w k -v k | ≥ ε/ √ 2 - √ 2 (w j -v j ) 2 + (w k -v k ) 2 ≥ ε/ √ 2 - √ 2∥w -v∥ > ε/ √ 2 - √ 2 • ε = -ε/ √ 2.Since this holds for all k ̸ = j, then,w j > max k∈[L] w k -ε/ √ 2,which proves that j ∈ s ε margin (w).In this section, we extend our experiment from Section 4 to consider some additional existing methods as baselines and evaluate each method using some common metrics from set-valued classification.Selection rules. We compare the following selection rules:1. Standard argmax. 2. ε-inflated argmax with tolerance ε = .05. 3. Top-K classification with K = 2. 4. Thresholding: including classes in the output set until the sum of probabilities becomes at least τ = 0.8, i.e.,Γ * τ (w) := ℓ ∈ [L] : w ℓ ≥ w [ k] ,where k = min{k :w [1] + • • • + w [k] ≥ τ }.5. Nondeterministic classification optimized for F 1 -score [DDB09] :NDC F1 (w) := ℓ ∈ [L] : w ℓ ≥ w [ k] ,wherek = min{k : w [1] + • • • + w [k] ≥ (k + 1)w [k+1] },with the convention w [L+1] = 0.6. Set-valued Bayes-optimal prediction [MWDH21]: for any utility u :[L]×2 [L] \{∅} → R + , SVBOP u (w) := argmax S⊆[L] ℓ∈[L]w ℓ • u(ℓ, S).We consider two utility functions based on u 65 and u 80 , defined below.Evaluation metrics. We evaluate each method based on a variety of metrics. In addition to β prec and β set-size , defined in Section 4, we assess each method in terms of utility-discounted predictive accuracy [ZCM12] . For a set S ⊆ [L] and label ℓ ∈ [L], define for some parameters α, β > 0, We use the measures u 65 with (α, β) = (1.6, 0.6) and u 80 with (α, β) = (2.2, 1.2), which respectively give small and large rewards for being cautious [NDMH18] .u(ℓ, S) := 1 {ℓ ∈ S} • α |S| - β |S| 2 . Sel.Finally, we directly assess the extent to which each selection rule resorts to returning a non-singleton set on difficult instances. Specifically, we define the superfluous inflation, which, for a selection rule s, is the ratio of the accuracy of the standard argmax given s returns at least two labels divided by the accuracy of s given s returns at least two labels:β sup. infl. := N j=1 1 Ỹj ∈ argmax(p( Xj )), |s(p( Xj ))| ≥ 2 N j=1 1 Ỹj ∈ s(p( Xj )), |s(p( Xj ))| ≥ 2 .If this ratio is close to 1, it means the standard argmax is correct as often as s (when the latter expresses uncertainty by returning multiple labels), so outputting multiple labels could be seen as overly conservative. Note that the argmax never returns more than a singleton, so its superfluous inflation is left blank in our results.Results. This experiment shows that the inflated argmax combined with the subbagged algorithm A m has accuracy (according to a variety of metrics) comparable with several alternative methods, does not output overly large sets of candidate labels, and at the same time admits rigorous stability guarantees. That is, our algorithmic framework does not unduly harm empirical performance.In Table 2 , we present results for each selection rule applied to the base algorithm A and the subbagged algorithm A m described in Section 4. The inflated argmax has significantly higher average precision and significantly smaller set sizes than all of the alternative set-valued classifiers. These two measures are related, since a selection rule can only have high precision if it frequently returns the correct label as a singleton set. The inflated argmax also has the lowest superfluous inflation, meaning that it tends return multiple labels on difficult instances. The u 65 of the inflated argmax is also within two standard errors of the u 65 of SVBOP 65 , which seeks to optimize this utility. Our method does have a significantly lower u 80 than many of the competing methods, since this utility is more forgiving of returning multiple labels.While this experiment considers many different perspectives on set-valued classification, we reiterate that our chief contribution is distribution-free stability guarantees. This means that, regardless of the dataset or base algorithm used, we can guarantee that our method will be stable. In the context of our experiment, Theorem 6 guarantees that δj ≤ δ * = ε -2 • 1-1/L n-1 • pn,m1-pn,m ≈ 0.006 for every test point j = 1, . . . , N . Furthermore, our optimality result shows that the inflated argmax returns a singleton as often as possible among all ε-compatible selection rules.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We attach our code in our submission to OpenReview, and we will deanonymize the link to the Github repository after the review process. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: The paper poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: PyTorch and Fashion MNIST are properly credited. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.38th Conference on Neural Information Processing Systems (NeurIPS 2024).We remark that all the results of this paper apply also to the case of countably infinitely many labels, L = ∞, in which case we should take N to be the label space instead of [L].We will assume without comment that C is measurable. Furthermore, in practice classification algorithms often incorporate randomization (either in the learning stage, such as via stochastic gradient descent, and/or in the selection stage, such as by using a random tie-breaking rule for argmax). All of our definitions and results in this paper can be naturally extended to randomized algorithms-see Appendix B.Methodology: assumption-free stable classificationIn this section, we formally define our pipeline for building a stable classification procedure using any base learning algorithm A. At a high level, we leverage Proposition 5 and separately address the learning and selection stages described in Section 2.1. In Section 3.1, we construct a bagged (i.e., ensembled) version of the base learning algorithm A. The recent work of[SBW24b] ensures that the resulting bagged algorithm has tail stability (ε, δ), with ε ≍ 1 √ nδ . 2. In Section 3.2, we introduce a new selection rule, the inflated argmax, and establish its ε-compatibility. Combined with the bagged algorithm, then, the resulting classification algorithm will be guaranteed to have selection stability δ (by Proposition 5).In this result, argmax(w) should be interpreted as the set of all maximizing entries of w-i.e., in the case of ties, argmax(w) may not be a singleton set. 
 
 Appendix:  


Decision: The reviewers have essentially converged around the idea that the paper is a worthy addition to NeurIPS; the paper indeed contains an interesting theoretical framework to build stability, and I commend the authors for elaborating further on experiments at rebuttal / discussion time.

At camera-ready time, I strongly suggest that the authors integrate not just their additional experimental results, but also make clear(er) point 2. in their rebuttal.

Reviews: review: paper_summary: paper_summary: The paper considers the stability of classifiers in multi-class classification. In the considered framework, the classifier is allowed to return a set of candidate classes instead of only one class. The stability is defined as the frequency of a set predicted by one classifier having no union with a set predicted with a classifier trained on the same training set with only one sample removed. The authors proposed the construction of a stable classifier in this framework using bagging and inflated argmax. The proposed approach is proven to have stability guarantees. Additionally, the effectiveness of the approach in terms of increased stability is confirmed in the empirical experiment.
strength_weakness: strength_weakness: - The authors argue that the argmax is a hard operator that may make the classifier unstable if even small changes of the outputs occur, but the definition of stability used in the paper also uses the hard operator of the set union being equal to an empty set. 
- The definition of the stable classifier as defined in this paper is new to me, and based on the paper, I struggle to find a good motivation for using this specific definition and focusing on it.
- The method requires bagging, what makes it difficult to apply with heavy architectures.
- The empirical results are limited to a single experiment. 
- The empirical experiment uses baselines that do not work in the framework of set-valued prediction and always predict only one class, which puts them at a clear disadvantage under the considered task. The comparison should include other set-valued classifiers. E.g., classifiers mentioned in related works or classifiers that optimize simple set-utility functions like in [1] and [2], that can be very efficiently optimized as discussed in [3].

[1] Del Coz JJ, Díez J, Bahamonde A. Learning nondeterministic classifiers. 2009
[2] Zaffalon M, Giorgio C, Maua DD. Evaluating credal classifiers by utility-discounted predictive accuracy. 2012
[3] Mortier T, Wydmuch M, Dembczynski K, Hullermeier E, Waegeman W. Efficient set-valued prediction in multi-class classification. 2021


NITs:
- "The more recent literature has focused on producing a sparse set of weights, but none of these works offer a formal stability guarantee for the support of the weights." This statement needs citation.
questions: questions: - What are the specific applications benefiting from this notation of stability?
- How does the introduced inflated argmax compare with other set-valued classifiers?
limitations: limitations: Limitations are discussed. I see no negative social impact of this work.
paper_summary: paper_summary: This submission studies the problem of making set-valued predictions, in which the classifier should strive for an optimal balance between the correctness (the true class is among the candidates) and the precision (the candidates are not too many) of its prediction. 

Yet, this submission mentions one method for this purpose [1], which predicts, for each query instance, a Bayes-optimal set valued prediction optimizing the F-measure. It seems to miss a discussion on (several) related methods, which are constructed based on both the probability theory and imprecise probability, e.g., the ones discussed in [3, 4, 6] and references therein. 

On the abstract level, the inflated argmax presented in Section 3.2  might be seen as a way to distort/imprecise the given singleton conditional probability distribution. It might enlarge the set of existing approaches targeting similar purposes, such as the ones discussed in [2, 5] and references therein, meaningfully. 

The proposed algorithm is compared with a variant of LeNet-5, implemented in PyTorch [PGML19] tutorials as GarmentClassifier(), and its ensemble version on the Fashion-MNIST data set. Empirical evidence suggests that the proposed algorithm provides better $\beta_{\text{prec}}$ than the single classifier, but worse than the ensemble version. The average set size produced by the proposed algorithm seems to be reasonably small.  Results on the selection instability $\delta_j$ defined in (4), i.e., an approximate of the loss version of the stability $\delta$ defined in (1), which essentially measures the robustness of the set-valued predictions under tiny changes in the training data set, seems to be in favor of the proposed algorithm. 

[1] Del Coz, J. J., Díez, J., & Bahamonde, A. (2009). Learning Nondeterministic Classifiers. Journal of Machine Learning Research, 10(10).

[2] Montes, I., Miranda, E., & Destercke, S. (2020). Unifying neighbourhood and distortion models: part I–new results on old models. International Journal of General Systems, 49(6), 602-635.

[3] Mortier, T., Wydmuch, M., Dembczyński, K., Hüllermeier, E., & Waegeman, W. (2021). Efficient set-valued prediction in multi-class classification. Data Mining and Knowledge Discovery, 35(4), 1435-1469.  

[4] Nguyen, V. L., Destercke, S., Masson, M. H., & Hüllermeier, E. (2018, July). Reliable multi-class classification based on pairwise epistemic and aleatoric uncertainty. In 27th International Joint Conference on Artificial Intelligence (IJCAI 2018) (pp. 5089-5095).

[5] Nguyen, V. L., Zhang, H., & Destercke, S. (2023, September). Learning Sets of Probabilities Through Ensemble Methods. In European Conference on Symbolic and Quantitative Approaches with Uncertainty (pp. 270-283). Cham: Springer Nature Switzerland.

[6] Troffaes, M. C. (2007). Decision making under uncertainty using imprecise probabilities. International journal of approximate reasoning, 45(1), 17-29.
strength_weakness: strength_weakness: W1: I think adding a few more competitors would be useful in assessing the potential advantages of the proposed algorithms. For example, threshold-based algorithms, which produce Bayes-optimal predictions of F-measure [1], and $U_{65}$ and $U_{80}$ [3], would be a good choice. They would cost $O(L\log(L))$ time to sort the labels according to the decreasing order of the predicted conditional probability masses and select the optimal threshold.   

W2: Assessing the classifiers with respect to utility-discounted predictive accuracies, such as $U_{65}$ and $U_{80}$ which respectively give small and large rewards for being cautious [7] , would help to further highlight the potential (dis)advantages of the classifiers. 
 
W3: Moreover, an imprecise classifier should abstain (i.e., provide set-valued predictions) on difficult cases, on which the precise classifier is likely to fail. This would be assessed in different ways. For example, one might consider reporting the correctness of the cautious classifiers in the case of abstention versus the accuracy of the precise classifiers [4].

W4: The tolerance $\epsilon$ is set to be $.05$. Please provide arguments supporting this choice. Yet, choosing this hyperparameter using a validation set might be challenging. Providing a sensitivity analysis on the choice of  $\epsilon$ would be helpful. 

[1] Del Coz, J. J., Díez, J., & Bahamonde, A. (2009). Learning Nondeterministic Classifiers. Journal of Machine Learning Research, 10(10).

[3] Mortier, T., Wydmuch, M., Dembczyński, K., Hüllermeier, E., & Waegeman, W. (2021). Efficient set-valued prediction in multi-class classification. Data Mining and Knowledge Discovery, 35(4), 1435-1469.

[4] Nguyen, V. L., Destercke, S., Masson, M. H., & Hüllermeier, E. (2018, July). Reliable multi-class classification based on pairwise epistemic and aleatoric uncertainty. In 27th International Joint Conference on Artificial Intelligence (IJCAI 2018) (pp. 5089-5095).  

[7] Zaffalon, M., Corani, G., & Mauá, D. (2012). Evaluating credal classifiers by utility-discounted predictive accuracy. International Journal of Approximate Reasoning, 53(8), 1282-1301.
questions: questions: Q1: Please refer to "Weaknesses" for detailed comments and suggestions for further revision.  

Q2: The notion of selection stability $\delta$ at sample size $n$ seems to be interesting. I guess the experimental setting provides the necessary information to compute the value of the selection stability $\delta$ given in (1) . To provide an idea of how empirical evidence may differ from the theoretical results, it might be helpful to compute that value and report the proportion of test instances whose empirical values $1 - \delta_j$ are smaller than $\delta$. 

Q3: I guess re-defining (4) as $\delta_j = \frac{1}{500}\sum_{k=1}^{500} 1{s(\hat{p}(\tilde{X}_j)) \cap s(\hat{p}^{\setminus i_k}(\tilde{X}_j)) = \emptyset}$ and changing the description of Figure 2 accordingly (to be consistent with the notion of the selection stability $\delta$ defined in Definition 1) might make things a bit clearer.
limitations: limitations: L1: Discussions on related work might need to be extended.

L2:  the empirical study might need to be enlarged with closely related algorithms.  

Please refer to "Weaknesses" for detailed comments and suggestions for further revision.
paper_summary: paper_summary: This paper focuses on the stability of learning multiclass classfiers. It proposes a notion of selection stability for the learning algorithm as well as a modified argmax that returns a set of labels. Evaluations have been conducted using FashionMNIST and simple models.
strength_weakness: strength_weakness: - The presentation of the proposed contribution can be improved.
  - The core proposal of inflated argmax is introduced in section 3.2, which is too late.
  - Bagging is not part of proposal, and shows low relation to the contribution, which can be demonstrated using fewer space.
  - Additional properties of the inflated argmax is also not directly related to the problem.
  - "learning stability by leave-one-out data" and "prediction stability when using argmax for similar probabilities" seem to be mixed and not clearly addressed.
  - Some sentences such as `workflow is fundamentally incompatible with the goal of algorithmic stability` is too extreme, for example, simple streament using proper regularization terms can somehow mitigate the problem.
- The demonstration of empirical evaluation does not match the last sentence in abstract.
  - The combination of $\epsilon$-inflated argmax on the base learning algorithm should be compared.
  - It is also nice to show results on at least two datasets.
questions: questions: - How tail stability and selection stability relates to $\epsilon$-compatibility differently? To confirm, do bagging classfiers only satisfy tail stablity? Are there other algorithm designs satisfy these stabilities?
- How to better interpreter the empirical evaluation results. With similar accuracy shown in Table 1, the proposed method should show better stability in Figure 2?
limitations: limitations: Practical limitations are presented in section 5.
paper_summary: paper_summary: An algorithm is considered stable if small perturbations in the training data do not result in significant changes to its output. Stability has been previously explored in the context of regression. This paper extends this concept to the multiclass framework, where previous work has focused on stability in terms of output probability rather than label selection processes. Given that argmax is not a continuous function, it can lead to instability in the classifier. To address this issue, the authors first define stability for the multiclass domain and then demonstrate that a stable classifier alone is insufficient; a corresponding label selection strategy is also necessary. The authors propose a framework combining bagging and inflated argmax to derive such a combination, which they validate experimentally, showing that their approach achieves higher stability while maintaining precision at a minimal cost.
strength_weakness: strength_weakness: 1.The authors point out my primary concern in their discussion section, highlighting the practicality of the proposed framework. Notably, bagging is an computationally expensive procedure, and while the authors suggest that parallelization can mitigate this issue, I agree that this does not necessarily translate to a reduction in overall cost. This becomes particularly challenging when dealing with large datasets commonly encountered in modern classifiers (e.g., image classification), where the sheer scale of computation required can be prohibitively expensive even with multiple GPUs.
2. The experimental section relies heavily on a single dataset and classifier combination. To fully demonstrate the benefits of the proposed framework, I believe it would be necessary to conduct a more comprehensive set of experiments. Specifically, I suggest that the authors provide additional evidence through experimentation to illustrate how the slight tradeoff in precision is justified by the gained stability in practice.
questions: questions: 1. From the definition of stability given in 4, it is evident that the authors trained the classifiers multiple times while removing training samples. How does the overall average of average precision look like in such cases? Such a metric might show the advantage the proposed framework.
2. My other main question is related to practicality. I would like to know the author's thought more regarding that limitation.
limitations: limitations: N/A
paper_summary: paper_summary: This paper studies and proposes a new theoretical framework to study algorithmic stability for multiclass classification with a focus on a stable class selection given predicted class probability. Based on the proposed theory, the proposed selection criterion based on bagging called "inflated argmax" was proposed. Theoretical analysis suggests that the proposed inflated argmax is stable, while also give a relatively tight decision for the candidate prediction set. Simple experiments using Fashion-MNIST were also provided to show the stability of inflated argmax as well as how tight the candidate set can be.
strength_weakness: strength_weakness: 1. Experiments were quite weak in my opinion to highlight the effectiveness of inflated argmax. For example, more reasonable baseline that has larger than 1 $\beta_{\mathrm{size}}$ could have been provided to show that inflated argmax is quite favorable if it allows to predict more than one class. Curreent baselines are restrictive in the sense that only one class can be chosen (argmax). Several improvements could also be done (please see the questions section for possible improvements). 
2. The proposed inflated argmax requires bagging, which can be quite computationally expensive, or prohibitive in some scenarios under limited resources or large dataset/model.
questions: questions: 1. Can we also see the performance of inflated argmax *without bagging* in the experiment? In my understanding, this should also be possible with the proposed efficient algorithm. This can be useful to see how important the bagging method is needed for the proposed algorithm. 
2. Visualizing the number of bags needed to perform reasonably in the experiment can also be useful. Because it can be expensive if we use more sophisticated algorithm or using larger dataset.
3. In deep learning, overconfident problem of probability can often be an issue when one wants to use probability score. How does this problem affect the proposed inflated argmax in practice?
4. Can we know that it is likely that the second class to be chosen as a candidate class is likely to be the second-largest probability class?
5. Is it reasonable to say that precision is similar given different average precision and standard error provided (Lines 293--294)?
6. I was wondering if there are other possible baselines in related work that could be used in the experiments? 
7. Suggestions on simple baselines
7.1 Top-2 class prediction, size of beta will be 2.
7.2 Predicting classes until sum of probability becomes a certain constant, e.g., if we set threshold = 0.8 and prediction probability vector is [0.1, 0.1, 0.1, 0.2, 0.6], then we will pick the last two classes (0.2 and 0.6). This simple idea looks intuitive and gives flexible set and this is the most simple baseline I had in mind that could be effective when I first read this paper. How do you think about this baseline?

I think the core contribution of this paper is the theory of stable classification rule. Therefore, although the experiment looks quite weak, my current rating is on the accept side (weak accept). I know it can be too much to ask to improve experiments a lot and I do not expect authors to fix everything according to my suggestions. But I would appreciate if the authors can reply to make sure that my understanding is correct.
limitations: limitations: Discussion of limitations of the proposed method and potential future work were appropriate.

Comments: title: Response to author rebuttal
Thanks for your reply
Thanks for your reply
Thanks for your replies
Thank you for the author feedback.
comment: I thank the authors for taking the time to addressing my concerns. I also see that some of these concerns were also discussed with other reviewers. After going through all the discussions, I think that my concerns have been addressed quite a bit. Thus I have decided to improve my score.
We sincerely appreciate your constructive comments throughout this process. We look forward to integrating your suggestions into our revision.
We sincerely appreciate your constructive comments throughout this process. We look forward to integrating your suggestions into our revision.
We sincerely appreciate your constructive comments throughout this process. We look forward to integrating your suggestions into our revision.
Thank you for the additional empirical evidence and discussions. To accommodate these, I increase the rating to $7$.
Thank you for your further detailed comments. We have expanded our table to include the two methods you suggest, Set-Valued Bayes-Optimal Prediction (SVBOP) for $u_{65}$ and $u_{80}$. We bold the statistical winners in each column.$^\dagger$ These new methods overall perform somewhat similar to NDC for $F_1$, but as you anticipated, the expected set size for SVBOP-65 is the smallest of the three. However, the expected set size for the inflated argmax is much smaller than these other three methods. Note also that although its $u_{65}$ utility *is* slightly smaller than SVBOP-65 on the test set, the inflated argmax qualifies as a statistical winner in this column as well. Our proposal thus performs competitively along three of the four criteria we have considered ($\beta_{\text{prec}}$, $\beta_{\text{size}}$ and $u_{65}$).

We agree that, on its own, we can view $\beta_{\text{prec}}$ as an evaluation metric for classification with rejection, since it requires a singleton prediction. Our motivating goal is to return a singleton as often as possible while achieving a user-specified selection stability level. While there are many different perspectives on set-valued classification, we reiterate that what is most novel about our work is that our proposal has distribution-free theory. This means that, regardless of the dataset or base algorithm used, we can guarantee that our method will be stable. Furthermore, our optimality result shows that the inflated argmax is singleton as often as possible among all $\epsilon$-compatible selection rules. 

| | Results for Base Algorithm $\mathcal{A}$| | | | | Results with Subbagging $\widetilde{\mathcal{A}}_m$| | | |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|          |  $u_{65}$ | $u_{80}$ | $\beta_{\text{prec}}$ | $\beta_{\text{size}}$ | | $u_{65}$ | $u_{80}$ | $\beta_{\text{prec}}$ | $\beta_{\text{size}}$ |
|  $\text{argmax}$  |  0.879 (0.003)  |  0.879 (0.003)  |  **0.879** (0.003) | **1.000** (0.000)  | | **0.893** (0.003)  |  0.893 (0.003)  | **0.893** (0.003)  |  **1.000** (0.000)  | 
|  $\text{argmax}^\varepsilon$  |  **0.881** (0.003)  |  0.883 (0.003)  |  **0.873** (0.003)  |  1.015 (0.001)  | | **0.895** (0.003)  |  0.897 (0.003)  | **0.886** (0.003)  |  1.017 (0.001)  | 
|  $\text{top-}2$  |  0.632 (0.001)  |  0.777 (0.001)  |  0.000 (0.000)  |  2.000 (0.000)  | | 0.631 (0.001)  |  0.777 (0.001)  | 0.000 (0.000)  |  2.000 (0.000)  | 
|  Thresholding $\Gamma^*_{0.8}$  |  0.880 (0.002)  |  **0.909** (0.002)  |  0.756 (0.004)  |  1.248 (0.005)  | | 0.868 (0.002)  |  0.907 (0.002)  | 0.703 (0.005) | 1.357 (0.006)  | 
|  NDC for $F_1$  |  **0.889** (0.003)  |  **0.902** (0.003)  |  0.832 (0.004)  |  1.105 (0.003)  | | **0.898** (0.003)  |  **0.915** (0.002)  | 0.825 (0.004) | 1.137 (0.004)  | 
|  SVBOP for $u_{65}$  |  **0.889** (0.003)  |  0.901 (0.003)  |  0.838 (0.004) | 1.091 (0.003)  | | **0.900** (0.003)  |  **0.915** (0.002)  | 0.835 (0.004) | 1.117 (0.003)  | 
|  SVBOP for $u_{80}$  |  **0.885** (0.003)  |  **0.910** (0.002)  |  0.777 (0.004) | 1.190 (0.004)  | | 0.882 (0.002)  |  **0.915** (0.002)  | 0.744 (0.004) | 1.245 (0.005)  | 

Regarding your second suggestion, we looked at, for each imprecise classifier $\hat{s}$, the ratio of "accuracy of standard argmax given $\hat{s}$ abstains" divided by "accuracy of $\hat{s}$ given $\hat{s}$ abstains". We call this the *superfluous inflation*, since if this ratio is close to 1, it means the argmax is correct as often as the set-valued classifier (when the latter abstains), so outputting a set could be seen as overly conservative. We include the table with these results below. The inflated argmax has the smallest ratio, meaning it only abstains in hard cases where it can improve the accuracy by returning a non-singleton set.

| | Results for Base Algorithm $\mathcal{A}$ | Results with Subbagging $\widetilde{\mathcal{A}}_m$ |
| --- | --- | --- |
| |  $\beta_{\text{superfluous-inflation}}$ | $\beta_{\text{superfluous-inflation}}$ |
|  $\text{argmax}^\varepsilon$  |  **0.496** (0.005)  |  **0.504** (0.005)  | 
|  $\text{top-}2$ | 0.905 (0.003) | 0.919 (0.003) | 
|  Thresholding $\Gamma^*_{0.8}$  |  0.622 (0.005) |  0.702 (0.005)  | 
|  NDC for $F_1$ | 0.531 (0.005)  |  0.591 (0.005)  | 
|  SVBOP for $u_{65}$ | 0.525 (0.005)  |  0.575 (0.005)  | 
|  SVBOP for $u_{80}$ | 0.611 (0.005)  |  0.689 (0.005)  | 



$^\dagger$ To check whether a given method wins along a given column, we compute the two-sample Z-score, subtracting that method’s value from the highest value in that column and normalizing by the pooled standard error. We say method is a winner if the two-tailed Z-test is *not* statistically significant at 0.05.
Thank you very much for the author's feedback. I have read the other reviews and the rebuttal. Thank you for the explanation and additional experiments.

I agree with the Author's rebuttal 2 regarding the core contribution. The core contribution of this paper is laying the foundation for stability classification by introducing several notions on classifier stability mentioned in this paper. This idea is novel to the best of my knowledge. Bagging is indeed more expensive than without Bagging, but the paper did a great job demonstrating that a theoretically justified algorithm archives this notion of stability. Having a more efficient algorithm or more analysis of this problem setting can be further studied in the future. 

I appreciate that other algorithms that are not necessarily stable but somewhat reasonable to try in this problem setting were added in the experiment section. This can motivate the practical usefulness of inflated argmax. 

Classification with abstention/rejection is indeed another framework that has been well-studied to abstain from making a prediction when the classifier is not confident. I treat this research direction as a different direction for making classifiers more reliable by outputting a flexible-size prediction set with some guarantee. This approach can also help machine learning systems to collaborate with humans to have better judgment and therefore worth investigating further.  

For these reasons, I increase the score to 7 (Accept).
I thank authors detailed and patient response and I have confirmed my shared concern over limited experiments and computational costs of bagging.
I thank authors for clarifying the organization of the manuscript, as well as my missing attention on section 3.2.2 and proposition 4. I also thank the improved figure 2.
For my concern 1.e. for definition 1, I missed the notation of feature $x$ in $C$ and thought this considers only the classifiers sets output by an learning algorithm.
Thanks for your detailed response. Additional empirical evidence and discussions indeed suggest more about the potential (dis)advantages of the proposed algorithms.   

Based on the response, I think we might see  $\beta_{\text{pre}}$ as an evaluation metric for classification with rejection, i.e., a predictor will only make a prediction for a query instance if it can make a precise/singleton prediction. This might also imply that whether the set-valued prediction covers the true class or not and whether the set size is large or not are not relevant. In this sense, I think it is fair to say that the proposed algorithms provide advantages, compared to top-2, Thresholding $\Gamma^*_{80}$ and NDC for $F_1$. From their definitions, it would be not hard to see that the set-valued predictions that optimize $u_{65}$, $F_1$, $u_{80}$ and a few other criteria, e.g., ones mentioned in [6], cover the top-ranked class, i.e., the output of $argmax$ (unless there are multiple classes with the highest score). By definition, optimizing $u_{65}$ would provide smaller set-valued predictions, compared to others, and therefore might naturally become a related competitor regarding $\beta_{\text{pre}}$.    

Given the (set-valued) predictions on the test set, the test set can be partitioned into 2 parts, where the classifier/predictor makes precise predictions and set-valued predictions.  As far as I understand from [7] (and [3]), $u_{65}$, $u_{80}$, and other utility-discounted predictive accuracies are designed to assess classifiers/predictors on the entire (test) data set. The additional results suggest that the proposed algorithms provide worse scores compared to NDC for $F_1$, which is also not designed to optimize both $u_{65}$ and $u_{80}$. If one opts for $u_{65}$ and $u_{80}$, I think it is reasonable to optimize them directly. Yet, I think it is difficult to find some specific evaluation criterion/criteria to assess multiple aspects of set-valued predictors. I think it is reasonable to expect a proposed set-valued predictor to demonstrate the advantages regarding at least a few criteria that assess the predictor on the entire (test) data set. 

Regarding "W3: Moreover, an imprecise classifier should abstain (i.e., provide set-valued predictions) on difficult cases, on which the precise classifier is likely to fail. This would be assessed in different ways. For example, one might consider reporting the correctness of the cautious classifiers in the case of abstention versus the accuracy of the precise classifiers [4].", I guess I should make the point more explicit. As far as I understand, the correctness of the cautious classifiers is the proportion of the time the set-valued predictions cover the true classes. I think such a criterion is designed to assess the cost of making set-valued predictions, i.e., the cautious classifiers should (only) provide set-valued predictions on difficult cases, on which the precise classifier is likely to fail. In my opinion, this criterion should complement the $\beta_{\text{pre}}$, utility-discounted predictive accuracies, set size, and other evaluation criteria, which take into account both precise predictions and set-valued predictions. 

Yet, I think the notion of selection stability is interesting (as I wrote in the initial review), and it might make the proposed algorithms differ from other cautious classifiers or set-valued predictors. The motivation to go with the proposed algorithms might need to be further elaborated. After due consideration, I keep my initial rating, but I appreciate any additional results and discussions that may strengthen the motivation to go with the proposed algorithms.