Source-linked AI summary
ReSWD: ReSTIR'd, not shaken. Combining Reservoir Sampling and Sliced Wasserstein Distance for Variance Reduction
Mark Boss, Andreas Engelhardt, Simon Donné, Varun Jampani
TL;DR
Distribution matching in vision and graphics needs scalable alternatives to Wasserstein distances, while SWD’s Monte Carlo estimator can produce noisy gradients. ReSWD integrates weighted reservoir sampling to retain informative directions during optimization. It improves on standard SWD and variance-reduction baselines across synthetic and real-world tasks, with current limitations for learned convolution-kernel projections.
Problem
Wasserstein distance is costly for high-dimensional distributions, while Monte Carlo estimation of SWD can have high variance and produce noisy optimization gradients.
Source: §1 · p. 1, §3.1 · p. 4Method
ReSWD integrates Weighted Reservoir Sampling into SWD to reuse and reweight informative projection directions during optimization while preserving unbiasedness.
Source: §1 · p. 2Results
ReSWD improves over standard SWD and existing variance-reduction baselines on synthetic distribution matching and real-world color correction and color diffusion guidance.
Source: §1 · p. 2, §5 · p. 10Takeaways & Limitations
The authors report that ReSWD reduces stochastic variance while supporting faster and more robust optimization.
Source: §1 · p. 2The technique is limited to simple matrix projections; learned convolution kernels did not yield satisfying results because their search space was too large.
Source: §4 · p. 9
Abstract
from arXiv · showhide
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/
1 INTRODUCTION
ReSWD applies weighted reservoir sampling to SWD to retain informative projection directions, addressing Monte Carlo variance that can make optimization gradients noisy and slow convergence. The paper reports improvements over standard SWD and variance-reduction baselines on synthetic and real-world distribution-matching tasks.
1 INTRODUCTION
- The Wasserstein distance captures shifts between distributions, but its curse of dimensionality limits its direct use in many iterative optimization settings.Distribution matching appears in vision and graphics tasks including histogram matching and color grading.Source: §1 · pp. 1–2
- SWD reduces Wasserstein computation to averaged one-dimensional transport costs, but Monte Carlo projection sampling can have high variance and scale with the number of directions.In optimization, this variance produces noisy gradients and slower convergence.Source: §1 · p. 2
- ReSWD adapts Weighted Reservoir Sampling to SWD, reusing and reweighting informative projection directions during optimization while preserving unbiasedness.It preferentially retains directions where the distributions are most dissimilar.Source: §1 · p. 2
- ReSWD improves over standard SWD and existing variance-reduction baselines on synthetic distribution matching and real-world color correction and color diffusion guidance.The introduction points to Fig. 1 for the real-world examples.Source: §1 · p. 2
2 RELATED WORK
The related work situates ReSWD among optimal-transport and sliced-Wasserstein approaches to distribution matching, including prior efforts to improve efficiency and reduce variance. It distinguishes ReSWD through a reservoir of high-contributing directions built during training and points to color matching and diffusion as application contexts.
Distribution Matching and Optimal Transport
- Optimal transport aligns distributions by minimizing the cost of moving probability mass, and Wasserstein distances are a common means of doing so.Entropic regularization enables scalable Sinkhorn solvers and divergences that interpolate between optimal transport and kernel MMD.Source: §2 · p. 2
- Sliced Wasserstein addresses high-dimensional costs by averaging one-dimensional transports; variants include Max-SW, which selects contributing directions, and methods using quasi-Monte Carlo or control variates to reduce variance.The literature also studies sliced-Wasserstein geometry and its use in graphics, vision, and generative modeling.Source: §2 · p. 2
- ReSWD differs from direction-selection work by building a reservoir of high-contributing directions during training, while drawing on variance-reduction methods from real-time path tracing.The authors describe this as keeping more directions in optimization.Source: §2 · p. 3
Color Transfer and Correction
- Color-transfer research spans classical histogram and optimal-transport methods, neural style-transfer approaches, and diffusion controls such as LoRA, IP-Adapter, and Control-Nets.A cited recent work uses optimal-transport constraints during generation to enforce a color distribution.Source: §2 · p. 3
- The authors position ReSWD as a way to improve efficiency in traditional color matching and support newer color-guided diffusion tasks..Source: §2 · p. 3
3 METHOD
SWD estimates Wasserstein distance using random one-dimensional projections, but Monte Carlo sampling can produce noisy optimization gradients. ReSWD reuses weighted reservoir directions while retaining an unbiased estimate, and is applied to color correction and diffusion guidance.
3.1 PRELIMINARIES
- Randomly projected one-dimensional Wasserstein costs make SWD efficient, but Monte Carlo direction sampling can yield noisy gradients during optimization.The one-dimensional Wasserstein calculation uses sorting, with complexity O(n log n).Source: §3.1 · pp. 3–4
3.2 RESERVOIR SLICED WASSERSTEIN DISTANCE (RESWD)
- Exponential time decay downweights stale reservoir directions, while an ESS threshold of 0.5K resets the reservoir to prevent weight collapse.The experiments use α = 0.5 for the ESS reset threshold.Source: §3.2 · pp. 4–5
- ReSWD combines stored and newly sampled directions, retaining high-contribution candidates through weighted reservoir sampling for an unbiased Monte Carlo loss estimate.The candidate pool contains K stored and M new directions.Source: §3.2 · pp. 4–5
3.3 APPLICATIONS
- For color correction, ReSWD optimizes CDL parameters in CIELAB space at a maximum image dimension of 128, then applies them to full-resolution images or video.The parameters are optimized within 150 steps.Source: §3.3 · pp. 5–6
- For SD3.5 diffusion guidance, matching uses predicted images decoded through the VAE, stops gradients at the backbone, and applies six Adam steps for 95% of denoising steps.The reservoir is reset twice during generation rather than after every denoising step.Source: §3.3 · p. 6
4 RESULTS
Across synthetic distribution matching and real-world color tasks, ReSWD improves on prior SWD techniques, with diffusion-guidance modifications also outperforming the prior method. Results include a warmup trade-off, a best-performing fresh-candidate setting, and a limitation to simple matrix projections.
General Distribution Matching
- ReSWD provides the best performance on 1D distribution matching with comparatively low runtime cost.The evaluation averages Mean-W1 over 1,000 distribution matches and reports running time.Source: §4 · p. 6, Table 1 · p. 7
- After roughly 140 steps, ReSWD's initially slower trajectory produces better results than pure SWD, while its loss has high correlation with true Wasserstein distance.The initial lag is attributed to reservoir building; the high correlation is presented as indicating an unbiased method.Source: §4 · p. 6, Figure 3 · p. 7, Figure 4 · p. 7
Color Matching
- ReSWD achieves the lowest transform error at competitive runtime in color matching, and its visual results match the final color distribution without artifacts.The comparison reports color-chart metrics and runtimes; the figure highlights challenging areas.Source: §4 · p. 6, Table 3 · p. 7, Figure 5 · p. 8
Diffusion Guidance
- Diffusion-guidance modifications outperform the prior technique; even on SDXL, they are faster and provide better results.The main speed difference comes from gradient stopping, which the paper says does not influence guidance.Source: §4 · pp. 7–8, Table 4 · p. 8
- With the upgraded SD3.5 base model and necessary changes, example generations are faithful to the input image's color distribution.The figure caption also reports that ReSWD handles challenging guidance cases where Lobashev et al. (2025) sometimes degenerates.Source: §4 · p. 9, Figure 6 · p. 9
Ablation
- With 64 total projections, 8 fresh candidates achieve the best overall results across applications.Too few candidates starve the reservoir of new directions, whereas too many leave too few directions supporting optimization.Source: Table 2 · p. 7, §4 · p. 9
Limitations
- The current technique is limited to simple matrix projections; learned convolution kernels did not provide satisfying results because randomly finding substantially better kernels is difficult.The reported degradation resembled limiting projection directions and redrawing them only every few optimization steps.Source: §4 · p. 9
5 CONCLUSION
The paper combines rendering-influenced variance reduction techniques for SWD optimization into an efficient, unbiased approach that reports reduced overhead and state-of-the-art results across applications.
5 CONCLUSION
- The combined approach offers a more efficient and unbiased SWD optimization solution than recent variance reduction techniques.The paper reports reduced overhead while maintaining optimal performance.Source: §5 · p. 10
- The technique demonstrates superior performance across real-world and synthetic applications.The paper reports state-of-the-art results.Source: §5 · p. 10
SUPPLEMENTS
The supplements present additional application results.
SUPPLEMENTS
- The supplements present more results from the paper’s applications.Source: SUPPLEMENTS · p. 12
A DIFFUSION GUIDANCE
Further diffusion-guidance results show that the earlier method often produces blur, collapse, or artifacts in some color-guidance settings, while the updated pipeline avoids these issues.
A DIFFUSION GUIDANCE
- The earlier diffusion-guidance approach often blurs, collapses, or produces artifacts when SWD guidance enforces colors unlikely under the prompt.Source: A DIFFUSION GUIDANCE · p. 12
- The updated pipeline does not produce those issues in the further diffusion-guidance results.Source: A DIFFUSION GUIDANCE · p. 12
B COLOR MATCHING
Additional color-matching results report consistent, artifact-free performance and support arbitrary resolutions.
B COLOR MATCHING
- The method achieves consistent, artifact-free color matching and supports arbitrary resolutions.Source: B COLOR MATCHING · p. 12
- The additional color-matching results highlight subtle changes and artifact-free, high-resolution matching.Source: Figure 8 · p. 14