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

Conference: ECCV 2026
Paper: ECCV Official
Code: https://ReservoirSWD.github.io
Area: others
Keywords: Sliced Wasserstein Distance, Weighted Reservoir Sampling, Monte Carlo Variance Reduction, Diffusion Color Guidance, Color Correction

TL;DR

Inspired by the ReSTIR resampling algorithm in real-time computer graphics path tracing, this paper integrates Weighted Reservoir Sampling (WRS) into the Sliced Wasserstein Distance (SWD) framework, adaptively preserving highly informative projection directions across optimization steps to boost gradient signal-to-noise ratio and accelerate convergence while strictly maintaining Monte Carlo unbiasedness.

Background & Motivation

Distribution matching represents a foundational pillar across computer vision and computer graphics, underpinning diverse workflows from classical histogram matching and cinema-grade color correction to texture alignment and controlled guidance for modern generative models. Within metric geometry and optimal transport, the Wasserstein distance provides an exceptionally robust framework due to its capacity to faithfully measure underlying geometric shifts between continuous or discrete probability measures. However, computing the exact Wasserstein distance in high-dimensional discrete empirical settings suffers severely from the curse of dimensionality, incurring cubic complexity for exact solvers and substantial memory/computational bottlenecks even when utilizing entropy-regularized Sinkhorn divergences in iterative gradient descent regimes.

The Sliced Wasserstein Distance (SWD) presents a computationally efficient alternative by projecting high-dimensional empirical measures onto randomly sampled 1D unit directions on the hypersphere. This reduces the intractable multidimensional transport problem to independent 1D order statistics sorting with \(O(n \log n)\) complexity per slice. Nonetheless, practical SWD implementations approximate the continuous expectation over projection directions via Monte Carlo (MC) integration. Due to computational budgets, only a modest number of random directions (typically dozens) are drawn at each optimization step. This introduces severe stochastic variance, leading to noisy, unstable gradient estimates that impede convergence, induce oscillations, and cause optimization trajectories to settle into suboptimal configurations.

Drawing insights from the ReSTIR spatiotemporal reservoir resampling framework widely utilized in real-time ray tracing, the authors observe that distinct projection directions exhibit vastly different degrees of discriminative power: directions capturing the largest geometric deviations and transport costs dominate the optimization gradient. Core Idea: integrate Weighted Reservoir Sampling (WRS) into the SWD optimization pipeline to adaptively maintain a persistent pool of high-contribution projection directions across iterations, combining exponential time-decay with effective sample size (ESS) flushing to achieve an unbiased Monte Carlo gradient estimator with dramatically reduced variance.

Method

Overall Architecture

ReSWD maintains a persistent reservoir of fixed capacity \(K\) containing influential projection directions across optimization steps. At iteration step \(t\), given empirical distributions \(X\) and \(Y\), the algorithm first applies an exponential time-decay to age the weights and keys of historical directions residing in reservoir \(R_{t-1}\). It then draws \(M\) freshly sampled candidate directions \(N_t\) from the uniform hypersphere distribution, forming a candidate pool \(P_t = R_{t-1} \cup N_t\) of size \(K+M\). After evaluating the 1D Wasserstein transport costs for all candidates, Weighted Reservoir Sampling selects the top \(K\) survivors proportional to their transport loss contribution. Finally, self-normalized importance weights with detached gradients form an unbiased Monte Carlo loss, while an effective sample size (ESS) monitoring threshold resets degenerate reservoirs to prevent weight collapse.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input empirical distributions X, Y and reservoir Rt-1"] --> B["Exponential Time-Decay Reweighting<br/>Age reservoir weights and random keys"]
    B --> C["Candidate Sampling and Joining<br/>Sample M new directions to construct candidate pool Pt"]
    C --> D["Weighted Reservoir Sampling Filtering<br/>Select K survivors proportional to 1D transport costs"]
    D --> E["Self-Normalized Weights and Unbiased Loss<br/>Gradient-detached importance loss evaluation"]
    E --> F{"Effective Sample Size ESS Monitor<br/>Check if ESS < αK"}
    F -->|Degenerate| G["Flush reservoir to empty set"]
    F -->|Healthy| H["Pass updated reservoir Rt to next step"]

Key Designs

1. Exponential Time-Decay Reweighting: Dynamically Forgetting Outdated Slices in Non-Stationary Optimization

During iterative optimization of neural parameters or continuous latent representations, the relative geometry between source and target distributions evolves continuously, rendering the loss landscape non-stationary. Projection directions that yielded massive loss contributions in early iterations frequently become irrelevant or redundant as distributions align; blindly accumulating historical directions causes the reservoir to saturate with stale, uninformative slices. To address this non-stationarity, ReSWD ages stored reservoir weights \(w_i\) and selection keys \(k_i\) before introducing fresh candidates:

\[\tilde{w}_i \leftarrow w_i \exp\left(-\frac{t - t_i}{\tau}\right), \quad \tilde{k}_i \leftarrow k_i \exp\left(-\frac{t - t_i}{\tau}\right)\]

where \(t_i\) denotes the step index when direction \(\theta_i\) was admitted into the reservoir, and \(\tau\) denotes the decay time constant. This exponential forgetting ensures continuous adaptation along the optimization trajectory without forfeiting the variance reduction gained from recent historical samples.

2. Key-Based Weighted Reservoir Sampling: Unbiased Selection of Dominant Projection Slices

To direct computational effort toward projection directions that yield the strongest optimization signal, the algorithm evaluates the 1D \(p\)-power Wasserstein distance \(D(\theta) = W_p(\pi_\theta \hat{x}, \pi_\theta \hat{y})\) for all \(K+M\) candidates in pool \(P_t\). Following the one-pass Efraimidis-Spirakis weighted sampling algorithm, each candidate is assigned a random selection key:

\[k(\theta) = u^{1 / D(\theta)}, \quad u \sim \mathcal{U}(0, 1)\]

By retaining the \(K\) candidates possessing the largest keys, the marginal inclusion probability \(q(\theta)\) of each direction entering the updated reservoir \(R_t\) is strictly proportional to its loss magnitude, \(q(\theta) \propto D(\theta)\). This eliminates the need for expensive non-uniform parametric sampling on the hypersphere, leveraging simple scalar sorting to isolate directions with the largest distributional discrepancy.

3. Self-Normalized Importance Weighting with ESS Reset Protection: Preserving Unbiasedness and Preventing Collapse

Because the \(K\) retained reservoir directions are intentionally biased toward high-cost projections, computing an unweighted arithmetic average would severely overestimate the true Wasserstein expectation. ReSWD constructs self-normalized importance weights using inverse inclusion probabilities \(1/q(\theta_i)\), formulating the estimator as:

\[\widehat{S}_p(X, Y) = \sum_{i=1}^K w_i D(\theta_i), \quad \text{where } w_i = \frac{1/q(\theta_i)}{\sum_{j=1}^K 1/q(\theta_j)}\]

Crucially, gradients through the importance weights \(w_i\) are detached during backpropagation, allowing derivative flows solely through the 1D sorting distances \(D(\theta_i)\). This preserves the unbiased Monte Carlo nature of the SWD metric. Furthermore, to guard against weight degeneration where a handful of directions dominate the normalized weights, the system continuously evaluates the effective sample size \(\text{ESS} = (\sum w_i)^2 / \sum w_i^2\). Whenever \(\text{ESS} < \alpha K\) (with \(\alpha = 0.5\)), the reservoir is completely flushed, eliminating catastrophic variance escalation.

Loss & Training

The framework is deployed across two domain-specific downstream tasks: - Color Correction: Differentiable Color Decision Lists (CDL) parameterize slope (\(s\)), offset (\(o\)), power (\(p\)), and saturation (\(\lambda\)). Pixels are converted into the perceptually uniform CIELAB space and evaluated at a downscaled resolution of 128 pixels. Optimization converges within 150 steps and applies directly to original high-resolution footage or video sequences without artifacts. - Diffusion Guidance: Applied to flow-matching models (SD3.5 Medium, Large, Turbo) and SDXL. At each diffusion step, the predicted clean image \(x_0\) is decoded via the VAE into CIELAB space and matched against reference images using ReSWD. To mitigate prohibitive memory and computational overhead in large Transformer backbones, a gradient stop is applied after the backbone—backpropagating solely through the VAE decoder. An Adam optimizer (\(\text{lr} = 3\times 10^{-3}\)) performs 6 gradient steps for 95% of the denoising trajectory, with the reservoir reset only twice throughout image generation.

Key Experimental Results

Main Results

Evaluation spans synthetic 1D/3D distribution matching (1,000 distribution pairs across normal, uniform, and bimodal Gaussian mixtures), cinematic color matching on real illumination scenes with ground-truth color charts, and diffusion color guidance using ContraStyles prompts and Unsplash Lite references.

On the 3D distribution matching benchmark (\(d=3, N=1024\), 300 steps, total budget of 64 projections), ReSWD outperforms standard SWD, control variate formulations, Max-SW, and Quasi-Monte Carlo (QMC):

Method	Mean-\(W_1\) [\(10^{-3}\)] \(\downarrow\)	Std Dev \(\sigma\) [\(10^{-3}\)]	95% CI	Gradient SNR \(\uparrow\)	Time per step [ms] \(\downarrow\)
SWD	0.733	0.050	\(\pm\) 0.003	0.215	1.03
LCV (Linear Control Variates)	0.735	0.044	\(\pm\) 0.003	0.218	1.81
UCV (Unbiased Control Variates)	0.726	0.045	\(\pm\) 0.003	0.200	2.19
Max-SW	29.152	38.278	\(\pm\) 2.373	0.0497	1.13
IS-EBSW-e	0.698	0.075	\(\pm\) 0.005	0.2291	1.42
QMC (Quasi-Monte Carlo)	0.670	0.043	\(\pm\) 0.003	0.208	1.38
ReSWD (Ours)	0.622	0.075	\(\pm\) 0.005	0.278	1.92
ReSWD + QMC	0.610	0.076	\(\pm\) 0.005	0.274	2.10

On real-world color matching benchmarks across 10 diverse lighting setups with GretagMacbeth color charts:

Method	Color PSNR \(\uparrow\)	Transform Error (RMSE) \(\downarrow\)	CDL Error (RMSE) \(\downarrow\)	CTQM Metric \(\uparrow\)	Time per match [s] \(\downarrow\)
Reinhard et al.	21.94 \(\pm\) 5.10	0.31 \(\pm\) 0.28	0.14 \(\pm\) 0.09	5.12 \(\pm\) 1.82	1
Nguyen et al.	18.76 \(\pm\) 3.69	1.27 \(\pm\) 1.32	0.33 \(\pm\) 0.23	5.16 \(\pm\) 1.78	3
Yoo et al.	20.97 \(\pm\) 5.44	0.44 \(\pm\) 0.38	0.21 \(\pm\) 0.17	5.09 \(\pm\) 1.80	20
Ho et al.	10.43 \(\pm\) 2.53	0.54 \(\pm\) 0.29	0.45 \(\pm\) 0.25	5.11 \(\pm\) 1.77	2
Larchenko et al.	14.80 \(\pm\) 3.47	0.47 \(\pm\) 0.34	0.20 \(\pm\) 0.14	5.05 \(\pm\) 1.82	24
Ours with SWD	24.30 \(\pm\) 6.12	0.34 \(\pm\) 0.30	0.11 \(\pm\) 0.08	5.15 \(\pm\) 1.81	5
ReSWD (Ours)	24.64 \(\pm\) 5.70	0.31 \(\pm\) 0.31	0.10 \(\pm\) 0.07	5.17 \(\pm\) 1.80	5

In diffusion guidance experiments, ReSWD on SD3.5-Turbo completes generation in only 4 seconds per image while achieving a Mean-\(W_2\) color distance of 0.675, and reaching 0.55 on SD3.5-Large, vastly outperforming the prior art Lobashev et al. (124s runtime and 1.941 Mean-\(W_2\)).

Ablation Study

Fixing the total projection budget at 64, the authors examine the ratio between freshly sampled candidate directions \(M\) and reservoir survivors \(K = 64 - M\):

Fresh Candidates \(M\)	Mean-\(W_1\) [\(10^{-3}\)] \(\downarrow\)	Std Dev \(\sigma\) [\(10^{-3}\)]	95% CI	Time [ms] \(\downarrow\)	Analysis
\(M=2\)	0.721	0.257	\(\pm\) 0.016	1.99	Severe candidate starvation; direction pool stagnates
\(M=4\)	0.673	0.151	\(\pm\) 0.009	1.96	Suboptimal exploration of changing geometry
\(M=8\) (Optimal)	0.622	0.075	\(\pm\) 0.005	1.92	Ideal balance between reservoir history and fresh exploration
\(M=16\)	0.746	0.122	\(\pm\) 0.008	1.91	Diminished reservoir capacity weakens variance reduction
\(M=32\)	1.192	1.283	\(\pm\) 0.080	1.98	Degraded reservoir retention leads to gradient jitter
\(M=48\)	2.122	3.280	\(\pm\) 0.203	1.93	Near loss of temporal accumulation benefits
\(M=56\)	3.811	3.317	\(\pm\) 0.206	1.85	Degenerates to noisy unbuffered Monte Carlo sampling

Ablating optimization components on the SDXL diffusion backbone reveals: - Full ReSWD @ SDXL: Mean-\(W_2 = 1.213 \times 10^{-2}\), runtime 34s. - Removing Adam (reverting to SGD): Mean-\(W_2\) degrades to \(1.846 \times 10^{-2}\), runtime 34s. - Removing Gradient Stop (backpropagating through the entire U-Net): Mean-\(W_2 = 1.832 \times 10^{-2}\) while runtime escalates to 126s, demonstrating that truncating backbone backpropagation yields a ~4x speedup with no loss in alignment accuracy.

Key Findings

• Exploration vs. Exploitation Sweet Spot: Setting \(M=8\) fresh candidates alongside \(K=56\) reservoir survivors achieves peak convergence performance. Undersized \(M\) starves the sampler and traps optimization in stale directions, whereas oversized \(M\) destroys the variance-reduction benefits of reservoir buffering.
• Dimensionality Breakdown: As feature dimension \(d\) scales up (shown in Table 5 across \(d=3\) to \(d=64\)), the probability of randomly drawing a slice that improves upon the existing reservoir decays exponentially. At \(d \ge 12\), ReSWD experiences direction staleness and degrades relative to standard SWD (at \(d=12\), ReSWD yields 9.768 vs. 1.190 for SWD). This delineates the operational sweet spot of ReSWD as low-to-medium dimensional spaces (such as 3D color channels).

Highlights & Insights

• Cross-Disciplinary Methodology: Adapting spatiotemporal reservoir sampling from graphics path tracing into optimal transport distribution matching represents an elegant and practical conceptual crossover.
• Strict Monte Carlo Unbiasedness: Unlike greedy extrema methods such as Max-SW that sacrifice theoretical fidelity, ReSWD leverages self-normalized importance weighting and detached gradients to ensure mathematical unbiasedness.
• Negligible Overhead: Adding reservoir maintenance costs under 1 millisecond per step while delivering over 29% improvement in gradient SNR and markedly smoother convergence trajectories.

Limitations & Future Work

• High-Dimensional Degradation: In high-dimensional spaces (\(d > 10\)), random hypersphere sampling fails to supply informative replacement slices, leading to reservoir stagnation and staleness. Future work must incorporate learned or active proposal distributions.
• Incompatibility with Convolutional Kernels: Attempts to extend the reservoir mechanism from matrix projections to learned spatial convolution kernels proved ineffective due to excessive search space complexity.
• Hyperparameter Sensitivity: The time-decay parameter \(\tau\) and ESS reset threshold \(\alpha\) currently require empirical tuning according to task-specific non-stationarity rates.

Related Work & Insights

• vs Sliced Wasserstein Distance (SWD): Standard SWD relies on memoryless uniform sampling at each step, suffering from high gradient variance. ReSWD reuses high-impact directions across steps to suppress variance without sacrificing unbiasedness.
• vs Max-SW: Max-SW selects solely the single direction maximizing projection loss, destroying the global distribution topology and causing severe mode collapse (Table 1 Mean-\(W_1\) error of 29.152). ReSWD preserves full distribution statistics via soft importance weighting.
• vs ReSTIR: Inherits the fundamental insight of reusing spatio-temporal reservoir samples from real-time ray tracing, but generalizes it to iterative hyperplane slicing in optimal transport metric spaces.

Rating

• Novelty: ⭐⭐⭐⭐ [Creative and mathematically sound integration of graphics reservoir resampling into optimal transport slicing]
• Experimental Thoroughness: ⭐⭐⭐⭐ [Evaluated across synthetic benchmarks, production-grade color calibration, and multi-model diffusion guidance with thorough ablations]
• Writing Quality: ⭐⭐⭐⭐⭐ [Clear structural organization, crisp mathematical formulations, and rigorous algorithmic definitions]
• Value: ⭐⭐⭐⭐ [Immediate utility for color-grading pipelines and generative diffusion control; establishes a valuable precedent for variance reduction in differentiable rendering]