Meta Two-Sample Testing: Learning Kernels for Testing with Limited Data 

1 Introduction . Two-sample tests ask , “ given samples from each , are these two populations the same ? ” For instance , one might wish to know whether a treatment and control group differ . With very low-dimensional data and/or strong parametric assumptions , methods such as t-tests or Kolmogorov-Smirnov tests are widespread . Recent work in statistics and machine learning has sought tests that cover situations not well-handled by these classic methods [ 1–18 ] , providing tools useful in machine learning for domain adaptation , causal discovery , generative modeling , fairness , adversarial learning , and more [ 19–34 ] . Perhaps the most powerful known widely-applicable scheme is based on a kernel method known as the maximum mean discrepancy ( MMD ) [ 1 ] – or , equivalently [ 35 ] , the energy distance [ 3 ] – when one learns an appropriate kernel for the task at hand [ 10 , 16 ] . Here , one divides the observed data into “ training ” and “ testing ” splits , identifies a kernel on the training data by maximizing a power criterion Ĵ , then runs an MMD test on the testing data ( as illustrated in Figure 1a ) . This method generally works very well when enough data is available for both training and testing . In real-world scenarios , however , two-sample testing tasks can be challenging if we do not have very many data observations . For example , in medical imaging , we might face two small datasets of lung ∗These authors contributed equally . †Now at Department of Statistics , University of Oxford . ‡Work done partially while at the Toyota Technological Institute at Chicago . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . computed tomography ( CT ) scans of patients with coronavirus diseases , and wish to know if these patients are affected in different ways . If they are from different distributions , the virus causing the disease may have mutated . Here , previous tests are likely to be relatively ineffective ; we can not learn a powerful kernel to distinguish such complex distributions with only a few observations . In this paper , we address this issue by considering a problem setting where related testing tasks are available . We use those related tasks to identify a kernel selection algorithm . Specifically , instead of using a fixed algorithm A to learn a kernel ( “ maximize Ĵ among this class of deep kernels ” ) , we want to learn an algorithm Aθ from auxiliary data ( Figure 1b ) : arg max Aθ E ( P , Q ) ∼τ [ Ĵ ( SteP , S te Q ; Aθ ( S tr P , S tr Q ) ) ] . ( 1 ) Here P , Q are distributions sampled from a meta-distribution of related tasks τ , which asks us to distinguish P from Q . The corresponding observed sample sets SP , SQ are split into training ( tr ) and testing ( te ) components . Aθ is a kernel selection algorithm which , given the two training sets ( also called “ support sets ” in meta-learning parlance ) , returns a kernel function . Ĵ finally estimates the power of that kernel using the test set ( “ query sets ” ) . In analogy with meta-learning [ 36–42 ] , we call this learning procedure meta kernel learning ( Meta-KL ) . We can then apply Aθ to select a kernel on our actual testing task , then finally run an MMD test as before ( Figure 1b ) . The adaptation performed by Aθ , however , might still be very difficult to achieve with few training observations ; even the best Aθ found by a generic adaptation scheme might over-fit to StrP , S tr Q . For more stable procedures and , in our experiments , more powerful final tests , we propose meta multi-kernel learning ( Meta-MKL ) . This algorithm independently finds the most powerful kernel for each related task ; at adaptation time , we select a convex combination of those kernels for testing ( Figure 1c ) , as in standard multiple kernel learning [ 43 ] and similarly to ensemble methods in few-shot classification [ 44 , 45 ] . Because we are only learning a small number of weights rather than all of the parameters of a deep network , this adaptation can quickly find a high-quality kernel . We provide both theoretical and empirical evidence that Meta-MKL is better than generic Meta-KL , and that both outperform approaches that do not use related task data , in low-data regimes . We find that learned algorithms can output kernels with high test power using only a few samples , where “ plain ” kernel learning techniques entirely fail . 2 Preliminaries . We will now review the setting of learning a kernel for a two-sample test , following [ 16 ] . Let X ⊂ Rd and P , Q be ( unknown ) Borel probability measures on X , with SP = { xi } mi=1 ∼ Pm and SQ = { yj } mj=1 ∼ Qm observed i.i.d . samples from these distributions . We operate in a classical hypothesis testing setup , with the null hypothesis that P = Q . Maximum mean discrepancy ( MMD ) . The basic tool we use is a kernel-based distance metric between distributions called the MMD , defined as follows . ( The energy distance [ 3 ] is a special case of the MMD for a particular choice of k [ 35 ] . ) Definition 1 ( MMD [ 1 ] ) . Let k : X × X → R be the bounded1 kernel of an RKHS Hk ( i.e. , supx , y∈X |k ( x , y ) | < ∞ ) . Letting X , X ′ ∼ P and Y , Y ′ ∼ Q be independent random variables , MMD ( P , Q ; k ) = sup f∈Hk , ‖f‖Hk≤1 E [ f ( X ) ] − E [ f ( Y ) ] = √ E [ k ( X , X ′ ) + k ( Y , Y ′ ) − 2k ( X , Y ) ] . If k is characteristic , we have that MMD ( P , Q ; k ) = 0 if and only if P = Q . We can estimate MMD using the following U -statistic estimator , which is unbiased for MMD2 ( denoted by M̂MD 2 u ) and has nearly minimal variance among unbiased estimators [ 1 ] : M̂MD 2 u ( SP , SQ ; k ) : = 1 m ( m− 1 ) ∑ i 6=j H ( k ) ij , ( 2 ) H ( k ) ij : = k ( xi , xj ) + k ( yi , yj ) − k ( xi , yj ) − k ( yi , xj ) . ( 3 ) Testing . Under the null hypothesis H0 , mM̂MD 2 u converges in distribution as m → ∞ to some distribution depending on P and k [ 1 , Theorem 12 ] . We can thus build a test with p-value equal to the quantile of our test statistic mM̂MD 2 u under this distribution . Although there are several methods to estimate this null distribution , it is usually considered best [ 10 ] to use a permutation test [ 46 , 47 ] : under H0 , samples from P and Q are interchangeable , and repeatedly re-computing the statistic with samples randomly shuffled between SP and SQ estimates its null distribution . Test power . We generally want to find tests likely to reject H0 when indeed it holds that P 6= Q ; the probability of doing so ( for a particular P , Q , k and m ) is called power . For reasonably large m , Sutherland et al . [ 10 ] and Liu et al . [ 16 ] argue that the power is an almost-monotonic function of J ( P , Q ; k ) : = MMD2 ( P , Q ; k ) σH1 ( P , Q ; k ) , σ2H1 ( P , Q ; k ) : = 4 ( E [ H ( k ) ij H ( k ) i ` ] − E [ H ( k ) ij ] 2 ) . ( 4 ) Here , σ2H1 is the asymptotic variance of √ m M̂MD 2 u under H1 ; it is defined in terms of an expectation of ( 3 ) with respect to the data samples SP , SQ , for i , j , ` distinct . The criterion ( 4 ) depends on the unknown distributions ; we can estimate it from samples with the regularized estimator [ 16 ] Ĵλ ( SP , SQ ; k ) : = M̂MD 2 u ( SP , SQ ; k ) /√ σ̂2H1 ( SP , SQ ; k ) + λ , ( 5 ) σ̂2H1 ( SP , SQ ; k ) : = 4 m3 m∑ i=1 m∑ j=1 H ( k ) ij 2 − 4 m4 m∑ i=1 m∑ j=1 H ( k ) ij 2 . ( 6 ) Kernel choice . Given two samples SP and SQ , the best kernel is ( essentially ) the one that maximizes J in ( 4 ) . If we pick a kernel to maximize our estimate Ĵ using the same data that we use for testing , though , we will “ overfit , ” and reject H0 far too often . Instead , we use data splitting [ 2 , 5 , 10 ] : we partition the samples into two disjoint sets , SP = StrP ∪ SteP , obtain ktr = A ( StrP , StrQ ) ≈ arg maxk Ĵλ ( S tr P , S tr Q ; k ) , then conduct a permutation test based on MMD ( S te P , S te Q ; k tr ) . This process is summarized in Algorithm 2 and illustrated in Figure 1a . This procedure has been successfully used not only to , e.g. , pick the best bandwidth for a simple Gaussian kernel , but even to learn all the parameters of a kernel like ( 8 ) which incorporates a deep network architecture [ 10 , 16 ] . As argued by Liu et al . [ 16 ] , classifier two-sample tests [ 8 , 48 ] ( which test based on the accuracy of a classifier distinguishing P from Q ) are also essentially a special case of this framework – and more-general deep kernel MMD tests tend to work better . Although presented here specifically for M̂MDu , an analogous procedure has been used for many other problems , including other estimates of the MMD and closely-related quantities [ 2 , 5 , 49 ] . When data splitting , the training split must be big enough to identify a good kernel ; with too few training samples mtr , Ĵλ will be a poor estimator , and the kernel will overfit . The testing split , 1For the given expressions to exist and agree , we in fact only need Bochner integrability ; this is implied by boundedness of either the kernel or the distribution , but can also hold more generally . however , must also be big enough : for a given P , Q , and k , it becomes much easier to be confident that MMD ( P , Q ; k ) > 0 as mte grows and the variance in M̂MDu ( P , Q ; k ) accordingly decreases . When the number of available samples is small , both steps suffer . This work seeks methods where , by using related testing tasks , we can identify a good kernel with an extremely small mtr ; thus we can reserve most of the available samples for testing , and overall achieve a more powerful test . Another class of techniques for kernel selection avoiding the need for data splitting is based on selective inference [ 15 ] . At least as studied by Kübler et al . [ 15 ] , however , it is currently available only for restricted classes of kernels and with far-less-accurate “ streaming ” estimates of the MMD , which for fixed kernels can yield far less powerful tests than M̂MDu [ 50 ] . In Section 5.4 , we will demonstrate that in our settings , the data-splitting approach is empirically much more powerful .

Summary: The paper introduces a new way of doing 2 sample testing in the setting where we are given multiple related 2 sample testing tasks and in the setting where we only have limited data. They propose 2 novel methods that allow them to use related information from the "meta training" tasks in order to learn the "best" kernel for 2 sample testing for a new unseen task. The main idea they are trying to stress is how to do 2 sample testing, in the setting where one is only given little data. They show that their 2 proposed method beats standard 2 sample tests by leveraging related tasks. They also provide a theoretical justification and bounds on their method.