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: State of the art kernel based two-sample tests operator as two-stage procedures. The available data is split, then first a kernel is optimized, and then an MMD test is performed on the other split of the data. With small available data, however, it is hard to reliably learn a good kernel in the first stage. This problem is particularly relevant if the class of kernels is very large, for example deep kernels.  Assuming abundant access to data from related tasks, the paper proposes to pretrain the kernel selection on the related tasks. This meta learning is argued to meaningfully constrain the space of available kernels. Thus for the actual task we can hope to reliably learn a good kernel even with little training data.  The proposed method “Meta multi-kernel learning” essentially learns a kernel for each related task. These then serve as candidate kernels. With the data from the actual problem they simply optimize a linear combination of the candidate kernels. The latter is a standard apprach. Usually, however, the candidate kernels are some default kernels, potentially chosen by expert knowledge. The propsed meta learning approach thus provides a principled and data-driven way to find suitable candidate kernels.  The authors provide some theoretical analysis and include synthetic and real-data experiments to support the benefits of Meta two-sample testing.