Neural Wasserstein Two-Sample Tests

TL;DR

This paper introduces a neural Wasserstein two-sample test that leverages deep learning to identify low-dimensional projections where high-dimensional distributions differ, improving testing power and calibration without resampling.

Contribution

It develops a novel learning-based two-sample test using neural networks and Wasserstein distance, with theoretical guarantees and adaptive aggregation for unknown projection features.

Findings

The test is valid and consistent under the null hypothesis.

It achieves strong finite-sample performance in simulations.

The method provides asymptotically distribution-free calibration.

Abstract

The two-sample homogeneity testing problem is fundamental in statistics and becomes particularly challenging in high dimensions, where classical tests can suffer substantial power loss. We develop a learning-assisted procedure based on the projection 1-Wasserstein distance, which we call the neural Wasserstein test. The method is motivated by the observation that there often exists a low-dimensional projection under which the two high-dimensional distributions differ. In practice, we learn the projection directions via manifold optimization and a witness function using deep neural networks. To adapt to unknown projection dimensions and sparsity levels, we aggregate a collection of candidate statistics through a max-type construction, avoiding explicit tuning while potentially improving power. We establish the validity and consistency of the proposed test and prove a Berry--Esseen type bound for the Gaussian approximation. In particular, under the null hypothesis, the aggregated statistic converges to the absolute maximum of a standard Gaussian vector, yielding an asymptotically pivotal (distribution-free) calibration that bypasses resampling. Simulation studies and a real-data example demonstrate the strong finite-sample performance of the proposed method.