Distributionally Robust Weighted k -Nearest Neighbors
Abstract
Learning a robust classiﬁer from a few samples remains a key challenge in machine learning.
A major thrust of research has been focused on developing k -nearest neighbor ( k -NN)
based algorithms combined with metric learning that captures similarities between samples.
When the samples are limited, robustness is especially crucial to ensure the generalization capability of the classiﬁer.
In this paper, we study a minimax distributionally robust formulation of weighted k -nearest neighbors, which aims to ﬁnd the optimal weighted k -NN classiﬁers that hedge against feature uncertainties.
We develop an algorithm, Dr.k-NN , that efﬁciently solves this functional optimization problem and features in assigning minimax optimal weights to training samples when performing classiﬁcation.
These weights are class-dependent, and are determined by the similarities of sample features under the least favorable scenarios.
The proposed framework can be shown to be equivalent to a Lipschitz norm regularization problem.
We also couple our framework with neural-network-based feature embedding.
We demonstrate the competitive performance of our algorithm compared to the state-of-the-art in the few-training-sample setting with various real-data experiments.
Introduction
Machine learning has been proven successful in data-intensive applications but is often hampered when the data set is small.
For example, in breast mammography diagnosis for breast cancer screening [5], the diagnosis of the type of breast cancer requires specialized analysis by pathologists in a highly time- and cost-consuming task and often leads to non-consensual results.
As a result, labeled data in digital pathology are generally very scarce; so do many other applications.
In this paper we aim to tackle the general multi-class classiﬁcation problem when only very few training samples are available for each class [43].
Evidently, k -Nearest Neighbor ( k -NN) algorithm [3] is a natural idea to tackle this problem and shows promising empirical performances.
Notable contributions, including seminal work [19] and the follow-up non-linear version [34], go beyond the vanilla k -NN and propose neighborhood component analysis (NCA).
NCA learns a distance metric that minimizes the expected leave-one-out classiﬁcation error on the training data using a stochastic neighbor selection rule.
Some recent studies in few-shot learning utilize the limited training data using a similar idea, such as matching network [41] and prototypical network [39].
They are primarily based on distance-weighted k -NN, which classiﬁes an unseen sample (aka.
query ) by a weighted [Figure] vote of its neighbors and uses the distance between two data points in the embedding space as their weights.
The classiﬁcation performance of weighted k -NN critically depends on the choice of weighting scheme.
The distance measuring the similarity between samples is typically chosen by metric learning, where a task-speciﬁc distance metric is automatically constructed from supervised data [24, 33, 41].
However, it has been recognized that they may be not robust to the few-training-samples scenario, where an “outlier” may greatly deteriorate the performance.
An example to illustrate this issue is shown in Figure 1.
The training set with only three labeled samples includes two categories we want to classify: mop and Komondor.
As we can see, the query image is visually closer to the third sample and thus more likely to be misclassiﬁed as a Komondor.
Here the third sample in the training set is an “outlier”, since it is a Komondor dressing up as a mop, and it misleads the metric learning model to capture irrelevant details (e.g., the bucket and the mop handle) for the Komondor category.
Such a problem can become even severe when the sample size is small.
The discussion above highlights the importance of choosing a good weighting scheme in weighted k -NN.
To develop algorithms that are more robust in the few-training-samples settings, we propose a new formulation of distributionally robust weighted k -nearest neighbors.
More speciﬁcally, for a given set of features of training samples, we solve a Wasserstein distributionally robust optimization problem that ﬁnds the minimax optimal weight functions for the k -nearest neighbors.
This inﬁnitedimensional functional optimization over weight functions presents a unique challenge for which existing literature on distributionally robust optimization do not consider.
To tackle this challenge, we ﬁrst consider a relaxed problem that optimizes over all randomized classiﬁers, which turns out to admit a ﬁnite-dimensional convex programming reformulation in spite of being inﬁnite-dimensional (Theorem 1).
Next, we show that there is a weighted k -NN classiﬁer achieving the same risk and shares the same least favorable distributions (LFDs) as the robust classiﬁer (Theorem 2).
Thereby we prove the optimality of such weighted k -NN classiﬁer for the original distributionally robust weighted k -nearest neighbors problem.
Furthermore, we prove the equivalence between the proposed robust weighted k -NN classiﬁer with a Lipschitz norm regularization problem (Theorem 3), which provides promising insights on the generalization property of the robust classiﬁer.
Based on these theoretical results, we proposed a novel algorithm called Dr.k-NN .
Unlike the traditional distance-weighted k -NN that uses the same weight for all label classes, our algorithm introduces a vector of weights, one for each class, for each sample in k -NN and performs a weighted majority vote.
These weights are determined from the LFDs and reveal the signiﬁcance of each sample in the worst case, thereby contributing effectively to ﬁnal decision making.
An example is illustrated in Figure 2.
Further, using differentiable optimization [4, 2], we incorporate a neural network into the minimax classiﬁer that jointly learns the feature embedding and the minimax optimal classiﬁer.
Numerical experiments show that our algorithm can effectively improve the multi-class classiﬁcation performance with few training samples on various data sets.
Related work Recently, there has been much interest in multi-class classiﬁcation with few training samples, see [31, 43] for a survey.
The main idea of our work is related to metric learning [19, 22, 33, 34, 32], which essentially translates the hidden information carried by the limited data into a distance metric, and has been widely adopted in few-shot learning and meta learning [15, 24, 39, 41, 30].
However, unlike few-shot and meta learning, where the goal is to acquire meta knowledge from a [Figure] large number of observed classes and then predict examples from unobserved classes, we focus on attacking a speciﬁc general classiﬁcation problem where the number of categories is ﬁxed but labeled data are scarce.
In this paper, we take a different probabilistic approach to exploit information from the data: we construct an uncertainty set for distributions of each class based on the Wasserstein distance.
Wasserstein distributionally robust optimization [10, 7, 13, 1, 8, 17, 38, 6, 36, 16] is an emerging paradigm for statistical learning; see [26] for a recent survey.
Our work is mostly related to [18], a framework for Wasserstein robust hypothesis testing, but is different in three important ways.
First, we focus on multi-class classiﬁcation, while [18] only studied two hypotheses.
Second, we focus on directly minimizing the mis-classiﬁcation error while [18] used a convex relaxation for the 0loss.
Third, we analyze the generalization bound while [18] does not.
Fourth, while [18] requires sample features as an input, we develop a scalable algorithmic framework to simultaneously learn the optimal feature extractor parameterized by neural networks and robust classiﬁer to achieve the best performance.
A recent work [9] studies distributionally robust k -NN regression.
Note that regression and classiﬁcation are fundamentally different as different performance metrics are used.
In [9] the objective is to minimize the mean square error, whereas in our work we minimize classiﬁcation errors.
Last but not least, our minimax formulation is similar to [14], which introduces a generalization of the maximum entropy principle and aims to address a different supervised learning problem.
It is worth noting that the paper considers another type of ambiguity set deﬁned via cross-moment constraints.
In contrast, we adopt Wasserstein sets in our formulation, leading to different tractable optimization reformulations and regularization terms.
Another well-known work on optimal weighted nearest neighbor binary classiﬁer [35] assigns one weight to each sample; the optimal weights minimize asymptotic expansion for the excess risk (regret).
In contrast, we consider minimax robust multi-class classiﬁcation, each training sample is associated with different weights for different classes, and we consider the few sample regime instead of the asymptotic regime in [35].
Distributionally Robust k -NN
In this section, we present our model.
We ﬁrst deﬁne weighted k -NN classiﬁer in Section 2.1, then present the proposed framework of distributionally robust k -NN problem in Section 2.2.
2.1 Weighted k -NN classiﬁer
Let { ( x 1 , y 1 ) , . . .
, ( x n , y n ) } be a set of training samples, where x i denotes the i -th data sample in the observation space X , and
y i ∈ Y := { 1 , . . .
, M } denotes the class (label) of the i -th data sample.
Let φ : X → Ξ be a feature extractor that embeds samples to the feature space Ξ (in Section 4.2 we will train a neural network to learn φ ).
Denote the sample feature vectors and the empirical support as:
[Equation]
Let
[Equation]
Deﬁne empirical distributions:
[Equation]
where δ denotes the Dirac point mass, | · | denotes the cardinality of a set, and I denotes the indicator function.
Let π : Ξ → ∆ M be a randomized classiﬁer that assigns class m ∈ { 1 , . .
, M } with probability π m ( ξ ) to a query feature vector ξ ∈ Ξ , where ∆ M is the probabilistic simplex ∆ M = { π ∈ R M + : {+(cid:80) Mm =
+}
π m = 1 } .
It is worth mentioning that the randomized test is more general than the commonly seen deterministic test.
In particular, the random classiﬁer π reduces to the deterministic test if for any ξ , there exists a m such that π m ( ξ ) = 1 .
Suppose the features in each class m follows a distribution
P m .
We deﬁne the risk of a classiﬁer π as the total error probabilities[-
-]
[Equation]
Recall that the vanilla k -NN is performed as follows.
Let c : Ξ × Ξ → R + be a metric on Ξ that measures distance between features.
For any given query point ξ , let τ S 1 ( ξ ) , . . .
, τ Sn ( ξ ) be a reordering of { 1 , . . . , n } according to their distance to ξ , i.e., c ( ξ, ξ τ Si ( ξ ) ) ≤ c ( ξ, ξ τ Si +1 ( ξ ) ) for all i < n , where the tie is broken arbitrarily.
Here the superscript S indicates the dependence on the sample S .
In vanilla k -NN, we compute the votes as
[Equation]
The vanilla k -NN decides the class for ξ by the majority vote, i.e., accept the class arg max 1 ≤ m ≤ M p m ( ξ ) .
To deﬁne a weighted k -NN, let us replace the equal weights in (2) by an arbitrary weight function w m : Ξ × Ξ → R + for each class m = 1 , . .
, M :
[Equation]
and use a shorthand notation w := ( w 1 , . .
. , w M ) .
In the sequel, we deﬁne a general tie-breaking rule as follows.
For any ξ , denote M 0 ( ξ ) := arg max 1 ≤ m ≤ M p m ( ξ ) .
When |M 0 ( ξ ) | > 1 , there is a tie at ξ .
We denote π m ( ξ ) as the probability of accepting class m for m ∈ M 0 ( ξ ) and we have (cid:80) m ∈M 0 ( ξ ) π m ( ξ ) = 1 .
We deﬁne a weighted k -NN classiﬁer π knn ( ξ ; k, w ) :
Ξ → ∆ M as:
[Equation]
A weighted k -NN classiﬁer involves two parameters: number of nearest neighbors k and weighting scheme w .
Particularly, w m ( ξ, ξ i ) =
I { y i = m } recovers the vanilla k -NN, and w m ( ξ, ξ i ) =
I { y i = m } /c
( ξ, ξ i ) recovers the distance-based weighted k -NN.
Note that our deﬁnition (3) allows different weighting schemes for different classes, which is more general than the standard weighted k -NN.
The goal is to ﬁnd the optimal weighted k -NN classiﬁer π knn ( · ; k, w ) such that the risk Ψ as deﬁned in (1) is minimized.
Since the underlying true distributions are unknown, the commonly used loss function is the empirical loss, i.e., substitute the empirical distributions (cid:98) P m into the risk function (1).
This leads to the following optimization problem:
[Equation]
It is worth mentioning that this minimization problem is an inﬁnite-dimensional functional optimization, since the weighting schemes w is a function on Ξ × Ξ .
2.2 Distributionally robust k -NN
For few-training-sample setting, the empirical distributions might not be good estimates for the true distribution since the sample size is small.
To hedge against distributional uncertainty, we propose a distributionally robust counterpart of the weighted k -NN problem deﬁned in the previous subsection.
Speciﬁcally, suppose each class m is associated with a distributional uncertainty set P m , which will be speciﬁed shortly.
Given P 1 , . . .
, P M , deﬁne the worst-case risk of a classiﬁer π as the worst-case total error probabilities
[Equation]
where Ψ is deﬁned in (1).
We consider the following distributionally robust k -NN problem that ﬁnds the optimal weighted k -NN classiﬁer minimizing the worst-case risk:
[Equation]
Here the optimal solution P ∗ 1 , . .
, P ∗ M to the inner maximization problem is also called least favorable distributions (LFD) in statistics literature [21].
We summarize the architecture of the proposed distributionally robust k -NN framework in Figure 3, more details are provided in Section 4.
Now we describe the uncertainty set P m .
First, since we are going to re-weight the training samples to build the classiﬁer, we restrict the support of every distribution in P m to (cid:98) Ξ , the set of empirical points.
Second, the uncertainty set is data-driven, containing the empirical distribution (cid:98) P m and distributions surrounding its neighborhood.
Third, to measure the closeness between distributions, we choose the Wasserstein metric of order 1 [40], deﬁned as
[Equation]
for any two distributions P and P (cid:48) on Ξ , where the minimization of γ is taken over the set of all probability distributions on Ξ × Ξ with marginals P and P (cid:48) .
The main advantage of using Wasserstein metric is that it takes account of the geometry of the feature space by incorporating the metric c ( · , · ) in its deﬁnition.
Given the empirical distribution (cid:98) P m for m = 1 , . . .
M , we deﬁne
[Equation]
where P ( (cid:98) Ξ) denotes the set of all probability distributions on (cid:98) Ξ ; ϑ m ≥ 0 speciﬁes the size of the uncertainty set for the m -th class that speciﬁes the amount of deviation we would like to control.
Theoretical Properties
In this section, we analyze the computational tractability and statistical properties of the proposed distributionally robust weighted k -NN classiﬁer found in (6).
All proofs are delegated to Appendix B.
3.1 Robust Classiﬁcation
Observe that similar to (5), the formulation (6) is also an inﬁnite-dimensional functional optimization.
Let us ﬁrst relate it to a relaxed robust classiﬁcation problem, which turns out to be more tractable.
Consider the following minimax robust classiﬁcation problem over all randomized classiﬁers π (recalling ∆ M is the probability simplex in R [Equation] ConvexSolver 𝜔 𝑥 Least Favorable Distributions
T r a i n i n g s a m p l e s
T e s t i n g s a m p l e s DR 𝒌 -NN yield (𝑃 !∗ ($𝜔),𝑃 #∗ ($𝜔)) S i a m e s e c a t
Neural Networks (e.g. CNN) End-to-end Learning Framework 𝜋 ∗ (𝜔) Decision project 𝜔 to Ω yield ()𝑃 ! (𝜔), )𝑃 # (𝜔)) Yet still, (8) is an inﬁnite-dimensional functional optimization, since we are optimizing over the set of all randomized classiﬁers.
We establish the following theorem stating a ﬁnite-dimensional convex programming reformulation for the problem (8).
Theorem 1.
For the uncertainty sets deﬁned in (7) , the least favorable distribution of problem (8) can be obtained by solving the following problem:
[Equation]
subject to
[Equation] [Equation]
The decision variable γ m ∈ R n × n + can be viewed as a joint distribution on n empirical points with marginal distributions (cid:98) P m and P m , represented by a vector p m ∈ R n + .
The inequality constraint controls the Wasserstein distance between P m and (cid:98) P m .
Below we give an intuitive explanation for the objective function in (9).
Note that max 1 ≤ m (cid:48) ≤ M p im [-
-](cid:48) − p
im measures the margin between the maximum likelihood of ξ i among all classes and the likelihood of the m -th class.
Thus, the objective in (9) can be equivalently rewritten as minimization of total margin:
[Equation]
When M = 2 , the total margin reduces to the total variation distance.
Also, let y im ∈ { 0 , 1 } be the class indicator variable of sample ξ i , observe that
[Equation]
where the second term on the right side represents the cross-entropy (or negative log-likelihood).
Therefore, problem (9) perturbs ( (cid:98) P 1 , . .
, (cid:98) P M ) to LFDs ( P ∗ 1 , . .
, P ∗ M ) so as to minimize the total margin as well as an upper bound on cross-entropy of LFDs; the smaller the margin (or cross-entropy) is, the more similar between classes and thus the harder to distinguish among them.
3.2 Expressiveness of Weighted k -NN
In this subsection we study the expressive power of the class of weighted k -NN classiﬁers { π
knn ( · ; k, w ): 1 ≤ k ≤ n,w m : Ξ × Ξ → R + , 1 ≤ m ≤ M }
deﬁned in Section 2.1.
The following theorem establishes the equivalence between the original problem (6) and the relaxed robust classiﬁcation problem (8) studied in Section 3.1.
Theorem 2.
For the uncertainty set deﬁned in (7) , formulations (6) and (8) have identical optimal values.
In addition, there exists optimal solutions of (6) and (8) that share common LFDs that are optimal to (9) .
Theorem 2 implies that the set of weighted k -NN classiﬁers is exhaustive , in the sense that it achieves the same optimal robust risk as optimizing over the set of all randomized classiﬁers.
In our proof, we show that the weighted 1-NN classiﬁer , with weights equal to the LFDs of (9), is an optimal solution to (6).
Therefore, instead of solving (6) directly, by Theorem 1, we can solve the convex program (9) for the LFDs, based on which we construct a robust k -NN classiﬁer.
This justiﬁes the Dr.k-NN algorithm to be described in Section 4.
3.3 Lipschitz Regularization
Next, we discuss the connection between the proposed distributional robust k -NN framework in (6), (8) and an equivalent Lipschitz regularization problem.
Using duality for Wasserstein DRO [17], problem (8) is equivalent to
[Equation]
Then by [16], this problem is upper bounded by the following Lipschitz regularized classiﬁcation problem
[Equation]
where
[Equation]
is the Lipschitz norm of the function π m for each m = 1 , . . .
, M .
Perhaps surprisingly, the next result shows that (10) and (11) are actually equivalent (thus by Theorem 2, are both equivalent to (6)), despite that the loss function does not satisfy existing criteria ensuring the equivalence [13, 36, 16].
Theorem 3.
Formulations (10) and (11) are equivalent, and there exists an optimizer π ∗ of (11) satisfying (cid:107)
π ∗ m (cid:107) Lip = λ ∗ m , m = 1 , . .
, M , where ( λ ∗ m ) m is the optimizer of (10) when π = π ∗ .
The theorem is proved by using c -transform [40] to show that any optimizer π ∗ m can be modiﬁed into a λ ∗ m -Lipschitz classiﬁer while maintaining the optimality.
Note that Lipschitz classiﬁers are known to enjoy good generalization property [42, 20], thus the equivalence provides promising insights for the generalization property of the robust classiﬁer.
Proposed Algorithm Dr.k-NN
In this section, we present the Distributional robust k -Nearest Neighbor ( Dr.k-NN ) algorithm, which is a direct consequence of the theoretical justiﬁcations in Section 3.
4.1 Dr.k-NN algorithm
Based on Theorem 2, our algorithm contains two steps.
Step 1.
[ Sample re-weighting ]
For each class m , re-weight n samples using a distribution P
∗
m , where ( P ∗ 1 , . . P ∗ M ) is the p -component of the minimizer of (9).
Step 2.
[ k-NN ] Given a query point ξ , ordering the training samples according to their distance to ξ : [Equation]
Compute the weighted k -NN votes, deﬁne
[Equation]
[Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Table] [Figure] Decide the class for a query feature point ξ as arg max 1 ≤ m ≤ M (cid:101) p m ( ξ ) , where the tie is broken according to the rule (4).
Figure 4 gives an illustration showing the probabilistic weights ( P ∗ 1 , P ∗ 2 , P ∗ 3 ) for three classes and its corresponding decision boundary yielding from the weighted k -NN.
For the sake of completeness, we also extend our algorithm to non-few-training-sample setting, which is referred to as truncated Dr.k-NN .
The key idea is to keep the training samples that are important in deciding the decision boundary based on the maximum entropy principle [11].
This can be particularly useful for the general classiﬁcation problem with an arbitrary size of training set.
An illustration (Figure 6) and more details can be found in Appendix C.
4.2 Joint learning framework
In this section, we propose a framework that jointly learns the feature mapping and the robust classiﬁer.
Let the feature mapping φ (; θ ) be a neural network parameterized by θ whose input is a batch of training samples (Figure 3), and then compose it with an optimization layer that packs the convex problem (9) as an output layer that outputs the LFDs of (8).
The optimization layer is adopted from differentiable optimization [4, 2], in which the optimization problem is integrated as an individual layer in an end-to-end trainable deep networks and the solution of the problem can be backpropagated through neural networks.
To apply the mini-batch stochastic gradient descent, we need to ensure that each batch comprises of multiple “mini-sets”, one for each class, containing at least one training sample from each class fed into the convex optimization layer.
In light of (9), the objective of our joint learning framework is min [Equation] J [Equation]
θ ; P ∗
[Equation]
, P ∗ M ) , where
[Equation]
and { P ∗ m ( φ ( · ; θ )) } 1 ≤ m ≤ M are the LFDs generated by the convex solver deﬁned in (9) given input variables { ξ i = φ ( x i ; θ ) } 1 ≤ i ≤ n .
The algorithm is summarized in Algorithm 1 (Appendix A). [Figure]
Experiments
In this section, we evaluate our method and eight alternative approaches on four commonly-used image data sets: MNIST [28], CIFAR-10 [25], Omniglot [27], and present a set of comprehensive numerical examples.
We also test our method on two medical diagnosis data sets: Lung Cancer [12], and [-COVID-
 CT [38],-]COVIDCT [44], where very few data samples are available for study, due to privacy concerns and high costs associated with harvesting data.
Speciﬁcally, Lung Cancer data record 56 attributes for only 32 patients who have been diagnosed with three types of pathological lung cancers; [-COVID-
 Computed-]COVIDComputed Tomography (CT) data contain 349 COVID-
CT images from 216 patients andnon-COVID-19 CTs.
Here, we present a few examples of COVID-
CT images and [-non-COVID-
 CT-]non-COVIDCT images in Appendix E.
Experiment set-up.
We compare our method including Dr.k-NN and its truncated version with the following baselines: (1) k -NN based methods with different dimension reduction techniques, including Principal Component Analysis ( PCA+ k -NN ),
Singular Value Decomposition ( SVD+ k -NN ),
Neighbourhood Components Analysis ( NCA+ k -NN )
[19], and feature embeddings generated by Dr. kNN (Feature embedding + k -NN) as a sanity check; (2) matching networks [41]; (3) prototypical networks [39]; (4) MetaOptNet [29].
To make these methods comparable, we adopt the same naive neural network with a single CNN layer on matching network, prototypical network, and our model, respectively, where the kernel size is 3 , the stride is 1 and the width of the output layer is d = 400 .
In our experiments, we focus on an M -class K -sample ( K training samples for each class) learning task.
To generate the training data set, we randomly select M classes and for each class we take K random samples.
So our training data set contains MK samples overall.
We then aim to classify a disjoint batch of unseen samples into one of these M classes.
Thus random performance on this task stands at 1 /M .
We test the average performance of different methods using 1 , 000 unseen samples from the same M classes.
To obtain reliable results, we repeat each test 10 times and calculate the average accuracy.
Other experimental conﬁgurations are described as follows:
The Adam optimizer [23] is adopted for all experiments conducted in this paper, where learning rate is 10 − 2 .
The mini-batch size is 32.
The hyper-parameter ϑ m is chosen by cross-validation, which varies from application to application.
The differentiable convex optimization layer we adopt is from [2].
To make all approaches comparable, we use the same network structure in matching network, prototypical network, and MetaOptNet as we described above.
We use the Euclidean distance c ( ξ, ξ (cid:48) )
= (cid:107) ξ − ξ (cid:48) (cid:107) 2 throughout our experiment.
All experiments are performed on Google Colaboratory (Pro version) with 12GB RAM and dual-core
Intel processors, which speed up to 2.3 GHz (without GPU).
Results.
We present the average test accuracy in Table 1 for the unseen samples with different
M = 2 , 5 and K = 5 , 10 on small subsets of MNIST, mini ImageNet, CIFAR-10, Omniglot, and with M = 3 and K = 5 , 8 on lung cancer data.
Note that random performance for two-class and ﬁve-class classiﬁcations are 0 .
5 and 0 . 2 , respectively.
The ﬁgures in Table 1 show that Dr.k-NN ( k = 5 ) outperforms other baselines in terms of the average test accuracy on all data sets that 95% conﬁdence interval of our method’s performance on all the data sets is smaller than 0.08.
The truncated Dr.k-NN also yields competitive results using only 20% training samples ( τ = 0 . 9 ), compared to standard Dr.k-NN .
To conﬁrm that the proposed learning framework will affect the distribution of hidden representation of data points, we show the training and query samples in a 2D feature space and the corresponding decision boundary in Figure 5.
It turns out our framework ﬁnds a better feature representation in the 2D space with a smooth decision boundary and a reasonable decision conﬁdence map (indicated by the color depth in Figure 5 (a)).
Comparison to kernel smoothing.
We also compare with an approach using kernel-smoothing of the LFDs (in contrast to using k -NN) for performing classiﬁcation.
Consider a Gaussian kernel κ ( x ) =
| H h | − 1 / 2 κ ( H − 1 / 2 h x ) , where H h = hI p is the isotropical kernel with bandwidth h .
Then we replace (cid:101) p m ( ξ ) in Step 2 of Dr.k-NN with the following
[Equation]
We evaluate both methods on a subset of MNIST, which contains 1,000 testing samples (small dots) and 20 training samples (large dots) from two categories (indexed by blue and red, respectively).
As shown in Figure 7 and Figure 8 (Appendix D), our experimental results have shown the importance of using k -NN in our proposed algorithm, where the Dr.kNN signiﬁcantly outperforms the parallel version using kernel smoothing (even after the kernel bandwidth being optimized).
We ﬁnd that the performance when using kernel smoothing (13) heavily depends on selecting an appropriate kernel bandwidth h as illustrated by Figure 8.
Moreover, the best kernel bandwidth may vary from one dataset to another.
Therefore, the crossvalidation is required to be carried out to ﬁnd the best kernel bandwidth in practice, which is quite time-consuming.
In contrast, choosing the hyper-parameter k is an easy task, since we only have limited choices of k in few-training-sample scenario and the performance of Dr.k-NN is insensitive to the choices of k (see Figure 7).
Conclusion
We propose a distributionally robust k -NN classiﬁer ( Dr.k-NN ) for tackling the multi-class classiﬁcation problem with few training samples.
To make a decision, each neighboring sample is weighted according to least favorable distributions resulting from a distributionally robust problem.
As shown in the theoretical results and demonstrated by experiments, our methods achieve outstanding performance in classiﬁcation accuracy compared with other baselines using minimal resources.
The robust classiﬁer layer (9) serves an alternative to the usual softmax layer in a neural network for classiﬁcation and we believe it is promising for other machine learning tasks
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Section(c)
Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[N/A] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[Yes] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
B Proofs for SectionB.1 Proof of Theorem The proof of Theorem 1 is based on the following two lemmas. Lemma 1.
Fix probability distributions P 1 , . . . , P M ∈ P ( (cid:98) Ξ) , where (cid:98) Ξ = { ξ 1 , . . . , ξ n } . Then [Equation]
Furthermore, the optimal classiﬁer π ∗ satisﬁes that for any ξ i ∈ (cid:98) Ξ ,
[Equation]
This lemma gives a closed-form expression for the risk of the optimal classiﬁer if P 1 , . .
. , P M are known, and shows that the optimal decision π ∗ accepts the class with the maximum likelihood.
Moreover, when there is a tie (i.e., the set arg max 1 ≤ m ≤ M P m ( ξ ) is not singleton), the optimal decision π ∗ can break the tie arbitrarily.
Proof of Lemma 1.
We here prove a more general result for an arbitrary sample space Ξ .
Note that each P m , 1 ≤ m ≤ M , is absolutely continuous with respect to P 1 + · · · + P M , hence the RadonNikodym derivative dP
[Equation]
d
[Equation]
+ P M ) exists.
Using the interchangeability principle [37] that enables us to exchange the minimization and integration, we have
[Equation]
where the ﬁrst equality is obtained by plugging in the deﬁnition of Ψ in (1); the second equality is due the interchangeability principle; and the last equality holds because for any ξ , the inner minimization attains its minimum at one of the vertices of ∆ M .
More speciﬁcally, note that for each ξ , the objective of the inner minimization problem equals to 1 − (cid:80)
Mm =1 π m dP m d ( P 1 + ··· + P M ) ( ξ ) .
Under the constaint that π ∈ ∆ M , i.e., (cid:80) Mm =1 π m = 1 , we have:
[Equation]
and the equality holds when π are chosen such that
[Equation]
If there is a single maximum in { dP m d
[Equation]
P 1 + ··· + P M ) ( ξ ) , 1 ≤ m ≤ M } , say at index m ∗ , then this simply implies that the optimal π is chosen as π m ∗ = 1 and π m = 0 for m (cid:54)
= m ∗ .
If we substitute Ξ with the empirical support (cid:98) Ξ , the above formulation in Equation (14) translates into
[Equation]
therefore the lemma is proved.
Lemma 2.
For the uncertainty sets deﬁned in (7) , the problem max P m ∈P m , 1 ≤ m ≤ M ψ ( P 1 , . . .
, P M ) is equivalent to (9) .
Proof of Lemma 2.
Recall that the Wasserstein metric of order 1 is deﬁned as
[Equation]
for any two distributions P and P (cid:48) on Ξ , where the minimization of γ is taken over the set of all probability distributions on Ξ × Ξ with marginals P and P (cid:48) , i.e., the set
[Equation]
where P (Ξ × Ξ) denotes the joint probability distributions on Ξ × Ξ .
Therefore, the Wasserstein metric W ( P, P (cid:48) ) can be rewritten as
[Equation]
By the deﬁnition of uncertainty sets in (7) which contains discrete distributions supported on (cid:98) Ξ , we can introduce additional variables γ m ∈ R n × n + which represents the distribution on (cid:98) Ξ × (cid:98) Ξ , with marginals
P m ∈ P m and (cid:98) P m , for 1 ≤ m ≤ M .
For any ξ i , ξ j ∈ (cid:98) Ξ , let γ i,jm denotes γ m ( ξ i , ξ j ) for simplicity.
Thus the objective function in the above reformualtion of W ( P m , (cid:98) P m ) is (cid:80) ni
=1 (cid:80) nj =1 γ
i,jm c ( ξ i , ξ j ) .
The constraints W ( P m , (cid:98) P m ) ≤ ϑ m in (7) can be rewritten using γ m as n
[Equation]
i = n
[Equation]
j
[Equation]
[Equation]
[Equation]
Furthermore, the marginal distribution constraint of γ m reads:
[Equation]
Thereby the problem max P m ∈P m , 1 ≤ m ≤ M ψ ( P 1 , . .
P M ) is equivalent to the convex optimization formulation in (9).
Proof to Theorem 1.
By Lemmas 1 and 2, we have
[Equation]
To prove Theorem 1, it remains to verify the validity of exchanging max and min .
We identify π as ( π 1 , . . .
, π n ) , where π i ∈ R M
+ satisﬁes (cid:80) Mm =1 π im =
Similar to the proof of Lemma 2,
P m , 1 ≤ m ≤ M , can also be identiﬁed as a vector in R n .
Note that the objective function Ψ( π ; P 1 , . .
P M ) is linear in ( π 1 , . .
. , π n ) and concave in ( P 1 , . .
. , P M ) , and the Slater condition holds.
Hence applying convex programming duality we can exchange max and min and thus the result follows.
It is worth mentioning that the optimal solution π ∗ and corresponding LFDs P ∗ 1 , . . .
, P ∗ M always exist since they are solutions to a saddle point problem.
B.2 Proof of Theorem
Proof of Theorem 2.
On the one hand, since π knn ( · ; k, w ) can be regarded as a special case of the general classiﬁer π : Ξ → ∆ M , it holds that
[Equation]
On the other hand, by Lemma 2, there exists an optimal solution to the minimax problem (8), denoted as ( P ∗ 1 , . .
, P ∗ M ) , and the optimal classiﬁer π ∗ as given in Lemma 1.
Note that there exists 1 ≤ k ∗ ≤ n and weight functions w ∗ 1 , . .
, w ∗ M such that
[Equation]
for example, by taking k ∗ = 1 and w ∗ m =
P ∗ m , 1 ≤ m ≤ M .
This implies that
[Equation]
Thereby we have shown that formulations (6) and (8) have identical optimal values.
Moreover, by the strong duality results in Theorem 1, we know ( π ∗ ; P ∗ 1 , . . .
, P ∗ M ) is the saddle point for the formulation (8), and by the above arguments, we see that ( π knn ; P ∗ 1 , . .
, P ∗ M ) leads to the same optimal value for formulation (6) as ( π ∗ ; P ∗ 1 , . .
, P ∗ M ) for formulation (8).
Therefore, we show that π knn is indeed the optimal solution to (6).
B.3 Proof of Theorem
Proof of Theorem 3.
We ﬁrst show the equivalence between the Lipschitz regularized problem (11) and the minimax problem (8).
Denote by v Lip the optimal value of (11) and v dual the optimal value of (10).
Observe that if (cid:107) π m (cid:107) Lip ≤ λ m , then
[Equation]
Therefore, we have
[Equation]
If we can show v dual ≥ v Lip , then we prove the equivalence between (11) and (10), thus the equivalence between (11) and (8).
Let ( π ∗ ; λ ∗ 1 , . . .
, λ ∗ M ) be a dual minimizer of problem (10), whose existence is ensured by [17].
Deﬁne
[Equation]
Then it follows that
[Equation]
Deﬁne
[Equation]
Then by deﬁnition, (cid:107) ˜
π m (cid:107) Lip ≤ λ ∗ m .
Indeed, for any ξ, ˜ ξ ∈ (cid:98) Ξ , there exists a ξ 0 ∈ arg max ˆ ξ ∈ (cid:98) Ξ { 1 − λ ∗ m c (ˆ ξ, ξ ) − φ m (ˆ ξ ) } such that:
[Equation]
Furthermore, since ˜ π m (˜ ξ ) ≥ 1 − λ ∗ m c ( ξ, ˜ ξ ) − φ m ( ξ ) , ∀ ξ, ˜ ξ , we have φ m ( ξ ) ≥ 1 − ˜ π m (˜ ξ ) − λ ∗ m c ( ξ, ˜ ξ ) , ∀ ˜ ξ ∈ (cid:98) Ξ .
Hence, we have φ m ( ξ ) ≥ max ˜ ξ ∈ (cid:98) Ξ (cid:110) 1 − ˜ π m (˜ ξ ) − λ ∗ m c ( ξ, ˜ ξ ) (cid:111) .
Recall that ( π
[Equation]
; λ ∗ 1 , . .
. , λ ∗ M ) is a dual minimizer of problem (10):
[Equation]
thus v dual ≥
[Equation]
Since v dual is the minimum value, this means that if ˜ π is a feasible solution, then it is also an optimal solution to (10).
Next we verify ˜ π is a feasible classiﬁer, i.e., it satisﬁes 0 ≤ ˜ π ( ξ ) ≤ 1 and (cid:80) Mm =1 ˜ π m ( ξ ) = 1 , ∀ ξ .
First, by deﬁnition, π ∗ m (˜ ξ ) ≥ ˜ π m (˜ ξ ) , ∀ ˜ ξ ∈ (cid:98) Ξ , therefore ˜ π m (˜ ξ ) ≤ 1 and (cid:80) Mm =1 ˜ π m (˜ ξ ) ≤ 1 .
If we are able to show that
[Equation]
then we can show ˜ π is indeed a feasible classiﬁer.
To show (16), ﬁrst note that if π ∗ m (˜ ξ ) = 0 , then we have by deﬁnition ˜
π m (˜ ξ ) = 0 .
{+Moreover, for any ˜ ξ ∈ (cid:98) Ξ , there is a set M 
+}
⊂ { 1 , . . .
, M } such that for all m ∈ M 0 , π ∗ m (˜ ξ ) > 0 , and the worst-case distribution P ∗ m transports probability mass from supp (cid:98) P m to ˜ ξ , which suggests that there exits ˆ ξ m ∈ supp (cid:98) P m such that
[Equation]
It follows from the deﬁnition of φ m that
[Equation]
Meanwhile, by deﬁnition of ˜ π m ,
[Equation]
Thereby we have shown (16).
The proof is completed by noting that the optimal solution ˜ π satisﬁes (cid:107) ˜ π m (cid:107) Lip ≤ λ ∗ m and thus v dual ≥ v Lip .
Combine with the previous result that v dual ≤ v Lip , we have shown v dual = v Lip and the proof is completed.
C Memory-efﬁcient implementation of Dr.k-NN in data-intensive scenario
For the sake of completeness, we extend our algorithm to non-few-training-sample setting.
This can be particularly useful for the general classiﬁcation problem with an arbitrary size of training set.
In fact, k -NN methods notoriously suffer from computational inefﬁciency if the number of labeled samples n is large, since it has to store and search through the entire training set [19].
[Figure]
The main idea is to only keep the training samples that are important in deciding the decision boundary based on the maximum entropy principle [11].
As a measure of importance, we choose the samples with the largest entropy across all categories, based on the intuition that the samples with higher entropy has larger uncertainty and will be more useful for classiﬁcation purposes since they tend to lie on the decision boundary.
The entropy of a sample is deﬁned as follows.
Consider a random variable which takes value m with probability π m , (cid:80) Mi
=1 π m = 1 ; then the entropy of this random variable is deﬁne as
[Equation]
As a simple example, for Bernoulli random variable (which can represent, e.g., the outcome for ﬂipping a coin with bias p ), the entropy function is H ( p ) = − p log p − (1 − p ) log(1 − p ) , and it is a concave function achieving the maximum at p ∗ = 1 / 2 , which means that the fair-coin has the maximum entropy; this is intuitive as indeed the outcome of a fair coin toss is the most difﬁcult to predict.
Now we use this entropy to deﬁne the “uncertainty” associated with each training points.
With a little abuse of notation, deﬁne
[Equation]
Denote the minimal and maximal entropy of all the training points as
[Equation]
Deﬁne the τ -truncated training set as
[Equation]
The truncated Dr.k-NN is obtained similarly as Step 2 of Dr.k-NN by restricting the training set (cid:98) Ξ only to the samples in (cid:101) Ξ (samples with larger entropy).
Figure 6 reveals that the most informative samples usually lie in between categories.
We can see that a truncated Dr.k-NN classiﬁer with τ = 0 .
9 only uses 20% samples with little performance loss.
More experimental details is presented in Section 5.
D Comparison to kernel smoothing
Figure 7 and Figure 8 present a comparison of the results using Dr.k-NN and the kernel smoothing deﬁned in (13).
The results suggest that the performance of Dr.k-NN is insensitive to the choice of k , while the performance of the kernel smoothing is heavily depended on the choice of h
[Figure]
[Figure]
E Real data examples for COVID-19 CT
Figure 9 and Figure 10 show 16 real CT images collected from patients who have been diagnosed with COVID-19 and other diseases (non-COVID-19), respectively.