Comparing distributions under severe sample imbalance requires effectively using abundant majority observations while retaining sensitivity to differences involving the smaller sample. This challenge is particularly relevant for complex, non-Euclidean data, such as graphs, whose structural characteristics call for flexible testing procedures. Classifier two-sample tests (C2STs) provide such flexibility, but existing procedures are primarily designed for balanced or comparably sized samples. In this paper, we study C2STs in a highly imbalanced regime, and show that severe sample imbalance may substantially degrade testing performance. Different from conventional imbalanced learning approaches that treat sample imbalance merely as a nuisance for classifier training, we turn it into a source of repeated testing evidence. Building on this insight, we propose an Enhanced Imbalance-Aware Deep Classifier Two-Sample Test (EIAD-C2ST), a classifier-agnostic framework designed for sample imbalance that accommodates diverse data modalities by tailoring classifiers to data structure. For computationally expensive deep classifiers, we further develop a scalable representation-amortized implementation strategy. Theoretically, we establish rigorous guarantees on asymptotic size and characterize conditions for power consistency, allowing the imbalance ratio to diverge. We demonstrate the empirical gains using simulations and an application to blockchain forensics. The results reveal that the majority-sample surplus, when properly leveraged, can be organized into useful statistical evidence, rather than an obstacle, for distributional comparisons with complex data.
Figures & tables
Figure 1: Histograms of the statistic ( 3 ) in a sample-balanced setting (left) and imbalanced settings (middle and right) under H0 from 1,000 replications. Red lines superimposed in the plots are the probability density of N(0,1) . The empirical distribution aligns with the theoretical asymptotic distribution in the balanced setting, whereas there are apparent misalignments in the imbalanced setting.
Figure 2: Flowchart of EIAD-C2ST. The procedure first performs a stratified random split, then constructs two split-specific IAD-C2ST statistics by partitioning the majority sample into approximately balanced slices, training and evaluating classifiers on the resulting blockwise subproblems, and merging the corresponding blockwise p -values. The final EIAD-C2ST p -value is obtained by aggregating the two IAD-C2ST p -values as in ( 7 ).
Figure 3: Representation-amortized implementation of EIAD-C2ST within one train–test orientation. A single nonlinear encoder ϕ is learned from the full training sample and then frozen. For each balanced slice q , a lightweight probe hq is fitted on the frozen embeddings Zq={ϕ(DXtr),ϕ(DY,qtr)} , yielding fqRA=hq∘ϕ . The same encoder is reused across all Q slices; the swapped orientation and p -value aggregation follow the procedure in Algorithm 1 .
n
m
Ratio
ρ
MMD
V-C2ST
W-C2ST
OS-C2ST
IAD-C2ST
EIAD-C2ST
100
19900
199
0.8
0.051
0.000
0.030
0.026
0.078
0.074
100
9900
99
0.8
0.056
0.000
0.026
0.029
0.072
0.070
100
4900
49
0.8
0.046
0.007
0.030
0.047
0.063
0.056
100
1900
19
0.8
0.037
0.017
0.045
0.041
0.057
0.042
100
900
9
0.8
0.043
0.028
0.040
0.048
0.040
0.034
100
400
4
0.8
0.047
0.038
0.057
0.061
0.051
0.038
Table 1: Empirical size for the Euclidean simulation at the significance level α=0.05 . Results are reported based on 1,000 Monte Carlo replications.
Figure 4: Empirical power for the Euclidean simulation at the significance level α=0.05 . Results are based on 1,000 Monte Carlo replications.
Type
n
m
Ratio
Signal( s )
Mixing( h )
Graph-MMD
V-C2ST
W-C2ST
OS-C2ST
IAD-C2ST
EIAD-C2ST
size
100
9900
99
0
0
0.053
0.000
0.081
0.020
0.067
0.066
100
4900
49
0
0
0.047
0.000
0.053
0.025
0.041
0.042
100
1900
19
0
0
0.056
0.000
0.092
0.028
0.040
0.042
100
900
9
0
0
0.043
0.000
0.064
0.016
0.041
0.038
100
400
4
0
0
0.056
0.000
0.056
0.022
0.042
0.038
100
200
2
0
0
0.053
0.006
0.048
0.033
0.039
0.033
Table 2: Empirical size and power for the Graph simulation under varying imbalance and signal settings. Tests are conducted at α=0.05 , and results are based on 1,000 Monte Carlo replications. The highest power in each row is shown in bold .
Figure 5: Schematic illustration of the construction of labeled subgraphs from the Elliptic2 Bitcoin transaction graph.
Figure 6: Illustration of different cluster and path types used in the subgraph construction.
Figure 7: Illustration of licit and suspicious subgraph construction. The left example contains three licit paths: 1→3→6 , 1→3→4→5→6 , and 2→4→5→6 . The right example contains two suspicious paths: 7→8→10→12 and 7→8→9→10→12 , together with one illicit path: 7→8→11→13 . The red dashed boundaries enclose the unknown nodes retained in the final annotations: {3,4,5} for the licit subgraph and {8,9,10} for the suspicious subgraph. Node 8 is retained because it lies on suspicious paths, whereas node 11 is excluded because it lies only on the illicit path. Labeled endpoint nodes are excluded from both annotated subgraphs.
Type
n
m
Ratio
Graph-MMD
V-C2ST
W-C2ST
OS-C2ST
IAD-C2ST
EIAD-C2ST
size
100
9900
99
0.052
0.000
0.001
0.005
0.050
0.045
100
4900
49
0.066
0.000
0.008
0.010
0.038
0.035
100
1900
19
0.061
0.000
0.023
0.019
0.044
0.032
100
900
9
0.044
0.000
0.039
0.028
0.035
0.032
100
400
4
0.051
0.002
0.043
0.031
0.037
0.024
200
9900
49.5
0.044
0.000
0.008
0.015
0.028
0.028
Table 3: Empirical size and power of tests on subgraphs drawn from the Elliptic2 dataset. Tests are conducted at the significance level α=0.05 , and results are based on 1,000 Monte Carlo replications. The highest power in each row is shown in bold .
Data-adaptive two-sample testing assesses if two samples come from the same distribution, using a discrepancy learned from the data (e.g., via kernel-based feature representations). Such methods typically rely on data splitting to decouple learning from testing and control type I error. However, this paradigm is ill-suited to few-shot settings with severe sample-size imbalance: abundant reference samples are available, while only a handful of query samples arrive. In this paper, we show how this imbalance can be leveraged constructively. Using abundant reference data, we learn reference-dependent representations that summarize salient structure of the reference distribution and provide informative signals for detecting departures. We incorporate a collection of representation families that capture both global and local structure, and adaptively weight them using only reference samples via an uncertainty-guided principle. Theoretically, we establish permutation-based type I error control and show consistency of the aggregated test: as the sample sizes grow, the test power converges to one whenever the representation set contains at least one consistent representation. Empirically, our aggregation achieves strong performance across a range of benchmarks while retaining type I error control.
Motivated by the success of modern flow-based generative models in modeling complex data, we study two-sample testing through the lens of flow-based methods. We propose the Zero-Flow Two-Sample Test (ZF2ST), built on the zero-flow criterion, which characterizes distributional equality through a time-reversal antisymmetry of a learnable velocity field. We extend this criterion and further develop the Zero-Flow Discrepancy, an identifying discrepancy that controls the Wasserstein distance, and derive a variational representation in terms of a witness function. This representation naturally leads to a witness-based test whose power is governed by the signal-to-noise ratio (SNR), allowing direct power maximization for witness learning. ZF2ST learns the witness on one data split and performs testing on held-out samples, thereby maintaining Type-I error control and admitting a simple asymptotic null distribution. Experimentally, ZF2ST performs competitively across various synthetic and real-world benchmarks, while showing particularly strong performance in distinguishing image distributions from different sources.
Yakun Wang, Leyang Wang, Song Liu +1
University of Bristol · RIKEN AIP · University of Tokyo & RIKEN AIP
Two-sample testing is a fundamental tool for detecting distributional differences across scientific domains, but classical tests (including kernel-based tests) can be ineffective on high-dimensional structured data such as images. Recent deep two-sample tests improve sensitivity in these settings by learning informative representations, yet they provide limited insight into which data features drive rejection of the null hypothesis H0. To address this issue, we propose a counterfactual explanation framework for deep two-sample testing that generates sample-level edits moving observations from a source group toward a target group while explicitly reducing the discrepancy measured by the test. Our method combines a diffusion autoencoder with a pretrained deep two-sample test model and optimizes a maximum mean discrepancy (MMD) objective in the test model's representation space to produce plausible counterfactuals. We quantify distribution-level effects through changes in the test statistic and the resulting two-sample p-values. We evaluate the method on synthetic 2D shape datasets and two MRI cohorts. Across both settings, the counterfactual transformations consistently increase p-values relative to the original samples, indicating that the edited source set becomes statistically closer to the target distribution under the test. We measure minimality using LPIPS to ensure the counterfactuals remain close to the original samples. The resulting edits provide interpretable evidence of the features associated with the detected group differences. On MRI, the localized changes are consistent with known anatomical differences between cohorts.
Wei-Cheng Lai, Marco Simnacher, Christoph Lippert
Digital Health and Machine Learning Hasso-Plattner-Institute, University of Potsdam Potsdam, Germany · Chair of Statistics Humboldt-Universität zu Berlin Berlin, Germany · Hasso Plattner Institute for Digital Health at Mount Sinai Icahn School of Medicine at Mount Sinai New York, United States of America