---
title: 'ReSWD: Reservoir Sliced Wasserstein Estimator'
url: https://www.emergentmind.com/topics/reservoir-swd-reswd-estimator
type: topic
---

# ReSWD: Reservoir Sliced Wasserstein Estimator

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 [2510.01061].

## 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 $W_p$ term is a 1-dimensional Wasserstein cost, computable in $\mathcal{O}(n\log n)$ via sorting.

In high dimension $d$, random slices $\theta$ 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 $\theta$ where $D(\theta)=W_p(\pi_\theta X, \pi_\theta Y)$ is large. These directions produce more stable, higher-magnitude gradients. WRS maintains a fixed-size reservoir of $K$ directions across steps. At each optimization iteration, $M$ new candidate directions are drawn, and their associated costs are computed. All $K+M$ candidates are then subject to a probabilistic sampling that retains exactly $K$ 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**: $\widehat{S}_p(X,Y) = \frac{1}{L}\sum_{i=1}^L W_p(\pi_{\theta_i} X, \pi_{\theta_i} Y)$, $\theta_i \sim U(S^{d-1})$.
- **ReSWD Reservoir Update:**
  - At iteration $t$, form pool $\mathcal{P}_t = \mathcal{R}_{t-1} \cup \mathcal{N}_t$ (previous $K$ + $M$ new directions).
  - For each $\theta \in \mathcal{P}_t$, compute 1D cost $D_t(\theta)$, then sample key $k_j = u_j^{1/D_t(\theta_j)}$, $u_j \sim U(0,1)$.
  - Retain the $K$ directions with smallest keys as the new reservoir.
- **Estimator weighting:**
  $$
  q(\theta_i) = \frac{D_t(\theta_i)}{\sum_{j=1}^{K}D_t(\theta_j)}
  $$
  $$
  w_i = \frac{1/q(\theta_i)}{\sum_{j=1}^K 1/q(\theta_j)}
  $$
  - **ReSWD estimate:** $\widehat{S}_p^{\mathrm{ReSWD}}(X,Y) = \sum_{i=1}^K w_i D_t(\theta_i)$

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

### Pseudocode Overview

```
Input: sets X, Y; prior reservoir R_{t-1} of size K; fresh sample count M; p, α (ESS threshold), τ (time-decay)
1. Optionally decay keys and weights in reservoir using τ
2. Draw M new directions N_t ~ U(S^{d-1})
3. Form pool P_t = R_{t-1} ∪ N_t
4. For each θ in P_t:
     compute D(θ)=W_p(π_θ X, π_θ Y)
     draw u~U(0,1); set key k(θ)=u^{1/D(θ)}
5. Sort P_t by key; keep K smallest → R_t
6. For survivors: compute q_i ∝ D(θ_i), normalize w_i, ReSWD loss = ∑_i w_i D(θ_i)
7. Compute ESS, reset if ESS < αK
Output: ReSWD loss, updated reservoir R_t
```

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 $W_1$ distance is consistently high (Pearson $\rho \approx 0.95$ vs. $\rho \approx 0.8$ 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 [2510.01061].

## 5. Computational Complexity and Practical Trade-offs

The per-step computation is dominated by sorting $n$ points in each of $(K+M)$ projections:
$$
\mathcal{O}\big((K+M) n \log n\big)
$$
Additional overheads—key generation, reservoir updates, and effective sample size (ESS) checks—are $\mathcal{O}(K+M+d)$. Typical hyperparameter settings (e.g., $K=56$, $M=8$, 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 ($W_1$/$W_2$) | Runtime          |
|----------------|--------------------------|------------------|
| SWD (64 slice) | $0.733 \times 10^{-3}$   | $1.03$ ms/step   |
| ReSWD (K=56,M=8) | $0.622 \times 10^{-3}$ | $1.92$ ms/step   |
| Diffusion SWD  | $1.94 \times 10^2$       | $124$ s          |
| ReSWD+SD3.5-turbo| $0.675 \times 10^2$    | $4$ 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 $24.6$ dB and transform RMSE comparable to, or better than, baseline methods, with efficient runtimes.
- **Diffusion guidance**: When guiding SD3.5-turbo using ReSWD, mean $W_2$ drops to $0.675 \times 10^2$ (from $1.94 \times 10^2$), CLIP-IQA improves (from $0.671$ to $0.800$), and runtime is dramatically reduced.

In all tested settings, ReSWD outperforms standard SWD and recent variance-reduction methods in both accuracy and efficiency [2510.01061].

## 7. Implementation Practices and Parameter Selection

- **Time decay ($\tau$):** Introducing a decay constant $\tau > 0$ enables forgetting of stale directions; recommended $\tau=50–100$ steps.
- **ESS threshold ($\alpha$):** Resetting the reservoir when effective sample size falls below half the nominal value ($\alpha=0.5$) prevents weight collapse.
- **$(K, M)$ selection:** For fixed slice budget $L=K+M$, moderate $M = 8–16$ and $K = L – M$ 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 [2510.01061].

Source: https://www.emergentmind.com/topics/reservoir-swd-reswd-estimator