ReSWD: Reservoir Sliced Wasserstein Estimator

• The paper introduces ReSWD, integrating Weighted Reservoir Sampling with SWD to reduce estimator variance and achieve more stable gradients.
• It applies adaptive sampling to prioritize informative projections, enabling faster convergence in tasks like color correction and diffusion model guidance.
• Empirical results show ReSWD reduces Wasserstein error by about 15% over MC-SWD, with improved gradient stability and convergence efficiency.

The Reservoir Sliced Wasserstein Distance (ReSWD) estimator is a variance-reduced extension of the Sliced Wasserstein Distance (SWD), designed for robust, unbiased distribution matching in high-dimensional settings. ReSWD integrates Weighted Reservoir Sampling (WRS) into the SWD Monte Carlo pipeline, adaptively prioritizing informative projection directions, thereby achieving more stable gradients and faster convergence. It provides a drop-in replacement for MC-SWD in gradient-based optimization routines, and has demonstrated empirically superior performance in synthetic and real-world tasks such as color correction and diffusion model guidance (Boss et al., 1 Oct 2025).

1. Sliced Wasserstein Distance and High-Variance Limitations

Given two empirical point clouds $X=\{x_1,\dots,x_{N_X}\}$, $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$, the true Sliced Wasserstein $p$-distance is

$$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$

where $\pi_\theta(x)=\theta^\top x$ projects $x$ onto direction $\theta$. The unbiased Monte Carlo SWD (MC-SWD) estimator approximates this expectation with $L$ i.i.d. samples $\{\theta_i\}_{i=1}^L$:

$$\widehat{S}_p(X,Y) = \frac{1}{L}\sum_{i=1}^L W_p\left(\pi_{\theta_i} X, \pi_{\theta_i} Y\right)$$

where each $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$0 term is a 1-dimensional Wasserstein cost, computable in $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$1 via sorting.

In high dimension $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$2, random slices $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$3 encode limited distributional information. This leads to high variance in the MC-SWD estimator, producing noisy gradients and slow convergence in optimization-based distribution matching tasks.

2. Weighted Reservoir Sampling and the ReSWD Principle

Weighted Reservoir Sampling (WRS) is employed to focus on "informative" projections—directions $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$4 where $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$5 is large. These directions produce more stable, higher-magnitude gradients. WRS maintains a fixed-size reservoir of $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$6 directions across steps. At each optimization iteration, $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$7 new candidate directions are drawn, and their associated costs are computed. All $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$8 candidates are then subject to a probabilistic sampling that retains exactly $Y=\{y_1,\dots,y_{N_Y}\}\subset\mathbb{R}^d$9 directions according to their costs.

This policy preferentially retains influential directions, enabling persistent gradient information and explicit variance reduction, while the reservoir replaces outdated (stale) slices over time. The marginal inclusion probability ensures unbiased Monte Carlo estimation.

3. Formal Definitions and Estimation Algorithm

Monte Carlo and ReSWD Estimators

• MC-SWD: $p$0, $p$1.
• ReSWD Reservoir Update:
  • At iteration $p$2, form pool $p$3 (previous $p$4 + $p$5 new directions).
  • For each $p$6, compute 1D cost $p$7, then sample key $p$8, $p$9.
  • Retain the $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$0 directions with smallest keys as the new reservoir.
• Estimator weighting:

$$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$1

$$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$2

• ReSWD estimate: $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$3

Unbiasedness follows from the inclusion probability structure and self-normalized importance weights.

Pseudocode Overview

$\theta$8

Reservoir persistence across steps ensures retention of informative, high-variance-reducing directions.

4. Variance Reduction Analysis and Empirical Performance

Empirically, on synthetic 3D Gaussian, uniform, and bimodal distribution benchmarks (1024 samples, 300 steps, 64 projections), ReSWD achieves approximately 15% lower final Wasserstein error than both plain SWD and prior variance reduction baselines such as control variates and quasi-Monte Carlo. The correlation of ReSWD estimates with the true $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$4 distance is consistently high (Pearson $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$5 vs. $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$6 for MC-SWD), reflecting reduced estimator variance.

During optimization, an initial "warm-up" phase (reservoir population) introduces a minor early convergence slowdown but yields superior convergence rates after this phase, typically within ∼140 optimization steps.

In color transfer and diffusion guidance applications, ReSWD reduces gradient standard deviation and produces visually smoother, artifact-free results, compared to baseline methods (Boss et al., 1 Oct 2025).

5. Computational Complexity and Practical Trade-offs

The per-step computation is dominated by sorting $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$7 points in each of $$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$8 projections:

$$S_p(X,Y) = \mathbb{E}_{\theta \sim U(S^{d-1})}\left[ W_p(\pi_\theta X, \pi_\theta Y) \right]$$9

Additional overheads—key generation, reservoir updates, and effective sample size (ESS) checks—are $\pi_\theta(x)=\theta^\top x$0. Typical hyperparameter settings (e.g., $\pi_\theta(x)=\theta^\top x$1, $\pi_\theta(x)=\theta^\top x$2, total slice budget 64) incur an approximately 10–20% runtime increase compared to MC-SWD but reduce estimation error by 10–20%. Memory overhead remains negligible relative to core data operations.

Setting	Mean Error ($\pi_\theta(x)=\theta^\top x$3/$\pi_\theta(x)=\theta^\top x$4)	Runtime
SWD (64 slice)	$\pi_\theta(x)=\theta^\top x$5	$\pi_\theta(x)=\theta^\top x$6 ms/step
ReSWD (K=56,M=8)	$\pi_\theta(x)=\theta^\top x$7	$\pi_\theta(x)=\theta^\top x$8 ms/step
Diffusion SWD	$\pi_\theta(x)=\theta^\top x$9	$x$0 s
ReSWD+SD3.5-turbo	$x$1	$x$2 s

6. Empirical Results and Application Domains

• Synthetic 3D distributions: ReSWD exhibits lower mean Wasserstein error and modest runtime overhead relative to MC-SWD.
• Color correction (dual-illumination scenes): ReSWD achieves peak PSNR $x$3 dB and transform RMSE comparable to, or better than, baseline methods, with efficient runtimes.
• Diffusion guidance: When guiding SD3.5-turbo using ReSWD, mean $x$4 drops to $x$5 (from $x$6), CLIP-IQA improves (from $x$7 to $x$8), and runtime is dramatically reduced.

In all tested settings, ReSWD outperforms standard SWD and recent variance-reduction methods in both accuracy and efficiency (Boss et al., 1 Oct 2025).

7. Implementation Practices and Parameter Selection

• Time decay ($x$9): Introducing a decay constant $\theta$0 enables forgetting of stale directions; recommended $\theta$1 steps.
• ESS threshold ($\theta$2): Resetting the reservoir when effective sample size falls below half the nominal value ($\theta$3) prevents weight collapse.
• $\theta$4 selection: For fixed slice budget $\theta$5, moderate $\theta$6 and $\theta$7 yield best trade-offs between exploration (new slices) and exploitation (informative retention).
• Batch optimizations: Amortize cost by using fused GPU sort across all projections in batch mode.
• Differentiability: Detach weight computations from back-propagation so that gradients are only computed with respect to the projected slice costs.
• Data space: For vision tasks, projections in perceptual color spaces, such as CIELAB, are preferred for perceptually meaningful distribution matching.
• Diffusion models: Stop gradient through the U-Net backbone, backpropagating only through the VAE decoder for memory efficiency.

With these operational choices, ReSWD is directly compatible with existing distribution-matching pipelines and delivers substantial variance reduction and convergence acceleration without sacrificing unbiasedness (Boss et al., 1 Oct 2025).