Wasserstein Filtering: A Sample Selection Method for Robust Distribution Learning

The paper addresses the fundamental question in robust statistics: "Can we design an accurate estimator and a computationally tractable algorithm to recover the true distribution P under the Wasserstein-1 metric, given samples subject to Total Variation (TV) contamination?"*

The authors consider a dataset D = x1,..., xn of i.i.d. samples drawn from a population distribution P*, subject to adversarial corruption where an adversary can modify up to an α-fraction of the samples. The contamination budget α ∈ [0, 1) constrains the adversary such that the number of modified samples is at most αn.

The authors adopt the Wasserstein-1 distance as the metric of distributional discrepancy because it metrizes weak convergence, naturally tolerates support mismatch, and is highly sensitive to the geometry of outliers. They note that existing methods like iterative filtering rely heavily on the unrealistic assumption that P* is isotropic, and they require prior knowledge of unknown spectral parameters associated with the covariance of P*.

The paper's main contributions are:

1. Novel Algorithmic Frameworks: Three distinct algorithms to compute the WF estimator: one marginal scheme (SinkMarg) and two joint optimization schemes (SinkWF and SlicedWF).

2. Rigorous Theoretical Guarantees: Establishment of minimax optimality of the WF estimator under robust distribution estimation, plus elementary optimization convergence rates.

3. Task- and Model-Agnostic Utility: The sample selection formulation is task-independent, serving as a modular preprocessing filter that cleans the dataset prior to downstream modeling.

The core insight is to select a subset of samples whose empirical distribution maximizes its Wasserstein distance to the fully contaminated empirical distribution, thereby preferentially isolating and removing geometrically influential outliers.

Formally, given a contaminated empirical measure P̂n and contamination fraction α, the estimator searches for:

• P̂ WF = argmax W1(P̂n, P) over P ∈ P α

where P α is the family of trimmed empirical measures obtained by removing αn outliers from the contaminated dataset.

The discrete optimization problem is NP-hard, so the authors relax it to its convex hull:

• p̂ WF = argmax W1(P̂n, P p) over p ∈ conv(Δ α)

where conv(Δ α) = v ∈ Rn: vT1 = 1, 0 ≤ vi ≤ 1/((1-α)n), ∀i ∈ [n].

SinkMarg (Marginal Screening): Computes leave-one-out Sinkhorn scores W1,λ(P̂n, P p−i) for each sample and filters out those with the largest scores. This works well for isolated outliers but fails to capture the group structure of clustered outliers.

SinkWF (Sinkhorn Wasserstein Filtering): Solves the entropy-regularized version of the relaxed problem using gradient-based optimization. The gradient is computed via the dual potential from the Sinkhorn algorithm, avoiding memory-intensive backpropagation. Uses Adam optimizer with multiple random initializations.

SlicedWF (Sliced Wasserstein Filtering): Uses the Sliced Wasserstein distance instead of the full Wasserstein distance, which has O(n log n) time and O(n) space complexity, avoiding the O(n2) cost matrix storage. However, the required number of projections S generally scales with the data dimension d.

The authors introduce the Far Exclusion and Local Projection (FELP) contamination model, which characterizes corruptions consisting of:

• Far Outliers (E ρ): Must be sufficiently far from the clean support that they are guaranteed to be excluded by the Wasserstein-maximization estimator.

• Near Outliers (N ρ): Lie close to the clean support; their impact is negligible even if not removed, as they can be modeled as local perturbations.

The model requires three conditions:

1. Support TV Bound: ∥P̂ − P*∥ STV ≤ α

2. Far Exclusion: There exists a subset S+ with P̂(S+) = 1−α and S+ ∩ E ρ = ∅ such that sup W1(P̂ S, P̂) 0

3. Local Projection: There exists a projection map T with d(x, T(x)) ≤ ρ and TP̂ N ρ ≤ P*

Theorem 1 (Population Risk): For any clean distribution P* ∈ G ⊆ W1(ρ, α) and contaminated distribution P̂ ∈ C FELP P*(α, ρ) with 0 ≤ α ≤ ½, the population estimator satisfies: W1(P̂ WF, P*) ≲ ρ.

Theorem 2 (Minimax Optimal Population Risk): For the class of distributions with bounded covariance G cov(σ), if 0 < α0 ≤ α ≤ ½ and dimension d ≲ 1/(α(1−α2)), then the minimax risk satisfies: inf sup R1,P*,T(α, ρ σ) ≍ ρ σ, where ρ σ ≍ σ√(dα) is the sharp resilience factor.

Theorem 3 (Empirical Risk): Under empirical FELP contamination, the empirical estimator satisfies: E[W1(P̂ WF, P*)] ≲ ρ + E[W1(Pn, P*)].

Theorem 4 (Minimax Optimality of Empirical Risk): Under Condition 1 (Far Exclusion with High Probability), the empirical minimax risk satisfies matching upper and lower bounds. For bounded covariance families: inf sup R1,P*,Tn,n(α, ρ σ) ≍ σ√(dα) + σ√(d)n(−1/d).

Theorem 5 (Convergence Rate): Under Conditions 2–4 (unique dual potential, stable active set, smoothness of optimal dual), the gradient ascent scheme achieves a sublinear convergence rate of O(1/√R), matching the standard optimal rate for first-order methods in non-concave smooth optimization.

On two-moons and concentric circles datasets with weak and strong contamination, SinkWF consistently achieves the best performance across most scenarios. The WF algorithms successfully detect outliers even when they are not at extreme edges of the clean distribution. On the concentric circles dataset, DIF, MDE, and COPOD mistakenly flag points on the outer boundary of the clean distribution as outliers, while SinkWF and SinkMarg successfully isolate nearly all hidden outliers.

Tabular Data (ADBench): SinkWF achieves state-of-the-art performance on datasets with higher contamination fractions (yeast, cardiotocography, satellite, shuttle), though performance degrades in low-α regimes. SlicedWF shows advantage on low-to-moderate dimensional datasets. A notable failure mode occurs on the smtp dataset where the total mass of the contaminated samples becomes statistically negligible.

Graph-Structured Data (Tox21): MDE fails completely in the 300-dimensional molecular representation space. SinkWF incurs an O(n2) space complexity but preserves full geometric fidelity, rendering it robust to high dimensionality.

Using DDPM on MNIST with strong and weak contamination, pre-filtering with WF algorithms yields a dramatic 3-to-4-fold reduction in FID score compared to training on the raw contaminated data. The α = 0 baseline experiment demonstrates that applying our filtering algorithms to entirely clean data has a negligible impact on downstream generative modeling.

The paper concludes that the WF framework offers:

• Task-Agnostic Preprocessing: Operates independently of downstream learning tasks

• Minimal Distributional Assumptions: No restrictive parametric assumptions or prior knowledge of parameters

• Computational Tractability: Efficient optimization via modern deep-learning frameworks

• Robustness to Non-Tail Outliers: Successfully identifies structured outliers not on extreme tails

• Flexibility across Dimensions: SlicedWF for large-sample low-dimensional settings, SinkWF for high-dimensional moderate-sample regimes

The authors establish minimax optimality under the FELP contamination model, confirming the framework achieves the optimal statistical rate of robust distribution estimation.

Based on this paper, I can improve AI systems in the following concrete ways:

Improvement: Integrate the Wasserstein Filtering (WF) estimator as a modular preprocessing layer before any downstream task (classification, clustering, generative modeling). The system automatically detects and removes geometrically influential outliers without requiring task-specific labels or parametric assumptions about the data distribution.

Capability: An AI system that can clean contaminated datasets (up to 50% corruption) before training, improving robustness to adversarial data poisoning, sensor failures, or annotation errors—without needing to know the true distribution family or covariance structure in advance.

Improvement: Replace traditional outlier detection (which focuses on tail extremes) with the SinkWF algorithm that maximizes Wasserstein distance to identify clustered, non-tail outliers. The system uses entropy-regularized optimal transport with gradient-based optimization to handle group-structured anomalies that evade conventional methods.

Capability: An anomaly detection system that catches hidden outliers embedded within the bulk of the data distribution (e.g., fraud clusters, coordinated attacks, systematic measurement biases) rather than only flagging extreme values.

Improvement: Implement a hybrid system that automatically switches between SlicedWF (for large-sample, low-to-moderate dimensional data with O(n log n) complexity) and SinkWF (for high-dimensional, moderate-sample data with full geometric fidelity) based on dataset characteristics.

Capability: A scalable robust estimation tool that maintains statistical optimality across diverse data regimes—from tabular datasets with thousands of samples to high-dimensional molecular representations (e.g., 300+ features) where traditional methods fail.

Improvement: Add a WF-based data purification step before training generative models (e.g., diffusion models, GANs). The system removes contaminated samples while preserving clean data structure, as demonstrated by the 3-4x FID improvement on MNIST.

Capability: A generative AI system that produces high-quality outputs even when trained on noisy or adversarially corrupted datasets, with negligible performance loss when applied to already-clean data (validated by the α=0 baseline).

Improvement: Build a learning system with theoretical guarantees matching the minimax lower bound (σ√(dα) + σ√(d)n(−1/d)) under the FELP contamination model. The system automatically determines the contamination fraction α and provides confidence bounds on estimation error.

Capability: An AI system that provides provably optimal worst-case performance guarantees under adversarial contamination, useful for safety-critical applications (medical diagnosis, autonomous systems) where worst-case robustness is essential.

Improvement: Implement the SinkWF algorithm with gradient computation via dual potentials from Sinkhorn iterations (avoiding backpropagation through the cost matrix), enabling O(n2) space usage only for the cost matrix rather than additional memory for gradients.

Capability: A robust learning system that can handle high-dimensional data (e.g., molecular fingerprints, image embeddings) with moderate sample sizes without memory overflow, while maintaining full geometric fidelity for outlier detection.

Given a dataset where a portion of the samples are contaminated, our goal is to recover the underlying clean population distribution. To this end, we propose Wasserstein Filtering (WF), a novel sample selection framework that discards a fraction of suspicious samples and estimates the target distribution using the empirical measure of the remaining data. The core insight is to select a subset of samples whose empirical distribution maximizes its Wasserstein distance to the fully contaminated empirical distribution, thereby preferentially isolating and removing geometrically influential outliers. To render this optimization computationally tractable, we introduce three algorithms: a marginal screening scheme, SinkMarg, and two joint optimization algorithms, SinkWF and SlicedWF, leveraging entropic optimal transport and sliced Wasserstein approximations, respectively. On the theoretical front, we introduce the Far Exclusion and Local Projection (FELP) contamination model, which characterizes corruptions consisting of well-separated outliers and locally indistinguishable perturbations. Under this model, we prove that the WF estimator achieves minimax optimality over distribution families with bounded covariance. Extensive numerical experiments on synthetic datasets, benchmark anomaly detection suites, and robust generative learning with diffusion models demonstrate that WF serves as a highly practical, model-agnostic preprocessing tool. It delivers competitive outlier detection performance and provides substantial downstream benefits for generative modeling under heavy contamination.