ReSWD: ReSTIR'd, not shaken. Combining Reservoir Sampling and Sliced Wasserstein Distance for Variance Reduction

TL;DR

The paper addresses high variance in Monte Carlo estimators for Sliced Wasserstein Distance (SWD) when used for distribution matching in vision and graphics. It introduces Reservoir SWD (ReSWD), which employs Weighted Reservoir Sampling to maintain a reservoir of informative projection directions, with time-decayed weights and self-normalized losses to preserve unbiasedness. Experiments on synthetic distributions, color correction, and diffusion guidance show that ReSWD reduces gradient noise, speeds up convergence, and outperforms standard SWD and other variance-reduction baselines. The approach provides a practical, unbiased variance-reduction strategy for high-dimensional distribution matching with broad applicability in graphics and AI-assisted generation tasks.

Abstract

Distribution matching is central to many vision and graphics tasks, where the widely used Wasserstein distance is too costly to compute for high dimensional distributions. The Sliced Wasserstein Distance (SWD) offers a scalable alternative, yet its Monte Carlo estimator suffers from high variance, resulting in noisy gradients and slow convergence. We introduce Reservoir SWD (ReSWD), which integrates Weighted Reservoir Sampling into SWD to adaptively retain informative projection directions in optimization steps, resulting in stable gradients while remaining unbiased. Experiments on synthetic benchmarks and real-world tasks such as color correction and diffusion guidance show that ReSWD consistently outperforms standard SWD and other variance reduction baselines. Project page: https://reservoirswd.github.io/

Figures (8)

• Figure 1: Overview. During the optimization we keep a reservoir of highly influential directions. We use the Sliced Wasserstein Distance (SWD) as a proxy metric for the reservoir update and the final optimization loss. Note that only directions in the reservoir influence the optimization to remain unbiased.
• Figure 1: Comparison on 1D distribution matching. Mean-W$_1$ over 1000 distribution matches for various methods alongside the respective running time. Here, we can see that ours provides the best performance with a comparatively low run-time cost.
• Figure 2: True Wasserstein metric over steps. The effect of reservoir warmup is clear when comparing ReSWD with the true Wasserstein distance during optimization. Initially, ReSWD performs slightly worse due to lower projections in the loss, but can outperform SWD in the end.
• Figure 3: Pearson correlation with the true Wasserstein Distance. Our method achieves a high correlation with the true Wasserstein loss, while improving upon pure SWD (See Fig. \ref{['fig:w1_vs_steps']}). This indicates the unbiased nature of our proposed method.
• Figure 4: Color Matching. Notice the consistent images our method produces with accurate shot matching. ReSWD consistently produces good results which match the final color distribution and without introducing any artifacts. We highlighted challenging areas which our method handled well. Note that the reference implementation or larchenko_color_2025 and nguyenIlluminantAwareGamutBased2014 do not support high resolution images.