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

会议: ECCV 2026
论文: ECCV 原文
代码: https://ReservoirSWD.github.io
领域: others
关键词: 切片Wasserstein距离, 蓄水池采样, 蒙特卡洛方差缩减, 扩散模型色彩引导, 色彩校正

一句话总结¶

借鉴图形学实时光线追踪中的 ReSTIR 重采样思想，本文将加权蓄水池采样（WRS）引入切片 Wasserstein 距离（SWD），在保持蒙特卡洛无偏性的同时自适应保留高方差投影方向，显著提升优化梯度的信噪比与收敛速度。

研究背景与动机¶

分布匹配是计算机视觉与图形学的基础课题，广泛应用于直方图匹配、电影色彩校正、纹理合成以及生成式扩散模型的引导控制。在度量几何中，Wasserstein 距离因其能够敏锐捕捉连续概率测度间的几何位移而备受青睐。然而，在高维离散样本场景下，直接计算精确 Wasserstein 距离的复杂度严重受制于维度诅咒，离散经验测度计算通常依赖立方级复杂度的线性规划，即使采用 Sinkhorn 熵正则化算法，其复杂度与内存开销在密集迭代优化任务中依然偏高。

切片 Wasserstein 距离（Sliced Wasserstein Distance, SWD）提供了一种极具吸引力的可扩展替代方案。它通过将高维分布沿均匀随机采样的超球面方向投影至一维空间，将原本高复杂度的最优传输问题转化为对一维排序统计量的快速计算，单次投影的复杂度仅为 \(O(n \log n)\)。然而，SWD 实践中高度依赖蒙特卡洛（Monte Carlo, MC）随机积分来近似超球面上所有投影方向的期望。为了避免高维计算量爆炸，实际优化中只能抽取数十个离散随机方向，这导致梯度的随机方差极大，在驱动神经网络权重微调或潜在空间位移优化时极易产生严重震荡，导致收敛缓慢且容易陷入局部次优点。

受到实时路径追踪领域时空重采样重要性采样算法（ReSTIR）的深刻启发，本文发现不同投影方向对分布差异的捕捉能力存在巨大悬殊——真正主导优化梯度的往往是那些两分布投影边缘差异最大、损失贡献最高的少数方向。核心 idea：将加权蓄水池采样（WRS）引入 SWD 优化框架，在迭代步间自适应维护高贡献投影方向池，结合时间指数衰减与有效样本量（ESS）重置机制，以极低开销实现无偏且方差大幅削减的梯度估计。

方法详解¶

整体框架¶

ReSWD 在优化过程中维护一个固定容量为 \(K\) 的持久化投影方向蓄水池（Reservoir）。在每个迭代步 \(t\)，算法接收两组经验样本分布 \(X\) 与 \(Y\)，先对上一步残留的蓄水池权重与随机键实施时间衰减（Time-decay），随后随机抽取 \(M\) 个全新方向候选，合并为容量为 \(K+M\) 的候选池。对候选池中的所有方向计算一维 Wasserstein 距离后，运行加权蓄水池采样挑选出具有代表性的 \(K\) 个幸存方向，并利用自归一化重要性权重构造无偏损失函数驱动梯度反传；若有效样本容量（ESS）过低则触发重置，以防止权重坍缩。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入经验样本分布 X, Y 与历史蓄水池 Rt-1"] --> B["时间指数衰减重加权<br/>老化蓄水池方向权重与随机键"]
    B --> C["新候选方向采样与合并<br/>抽取 M 个新方向构建 K+M 候选池 Pt"]
    C --> D["加权蓄水池采样筛选<br/>按一维传输代价分配采样概率保留 K 个方向"]
    D --> E["自归一化权重与无偏损失计算<br/>梯度截断权重反传驱动优化"]
    E --> F{"有效样本量 ESS 监测<br/>ESS < αK 判定"}
    F -->|触发退化| G["重置蓄水池为空集"]
    F -->|状态健康| H["保留当前蓄水池 Rt 传入下一步"]

关键设计¶

1. 时间指数衰减机制：动态遗忘非平稳优化轨迹中的过时方向

在神经网络参数或潜变量迭代优化场景中，两分布之间的相对几何形态是随时间连续形变的非平稳过程。早期迭代步中产生巨大损失贡献的投影方向，随着优化推进可能迅速变得无关紧要；若直接长期保留历史方向，蓄水池将充斥陈旧过时的非优质方向。为此，ReSWD 在抽取新方向之前，对上一时刻蓄水池 \(R_{t-1}\) 中所有方向的累积权重 \(w_i\) 与随机键值 \(k_i\) 引入基于时间差的指数衰减衰减因子：

\[\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)\]

其中 \(t_i\) 为该方向进入蓄水池的初始步数，\(\tau\) 为时间衰减常数。这种机制既避免了逐轮全盘抛弃历史信息的浪费，又赋予了采样器平滑适应分布演化的遗忘能力。

2. 基于随机键排序的加权蓄水池采样：无偏筛选高贡献投影方向

为了让计算预算集中在对损失梯度最关键的投影方向上，算法对候选池 \(P_t = R_{t-1} \cup N_t\)（共 \(K+M\) 个方向）分别计算一维 \(p\) 阶 Wasserstein 代价 \(D(\theta) = W_p(\pi_\theta \hat{x}, \pi_\theta \hat{y})\)。随后遵循 Efraimidis-Spirakis 单趟加权蓄水池采样算法，为每个方向赋予均匀随机变量 \(u \sim \mathcal{U}(0, 1)\)，并计算采样键值：

\[k(\theta) = u^{1 / D(\theta)}\]

通过在候选池中保留具有最大键值的 \(K\) 个幸存方向，确保每个方向被选入新蓄水池 \(R_t\) 的边际概率严格正比于其传输代价 \(q(\theta) \propto D(\theta)\)。这种方法无须复杂的非均匀超球面参数化采样，即可利用极廉价的标量排序操作，自适应捕获当前两分布差异最剧烈的几何切片。

3. 自归一化重要性加权与 ESS 重置保护：确保无偏性并规避权重坍缩

通过加权蓄水池筛选得到的 \(K\) 个方向天然聚焦于高损失切片，但若直接平均其代价会导致期望估计产生严重正向偏差。为此，算法基于重要性采样的逆包含概率 \(1/q(\theta_i)\) 构造自归一化权重 \(w_i\)，将最终估计器定义为：

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

在反向传播计算梯度时，重要性权重 \(w_i\) 的梯度被显式截断（detach），仅使导数流经一维排序测度 \(D(\theta_i)\)，严格维持了蒙特卡洛积分的无偏特性。同时，为防止个别方向在重要性重归一化后占据绝对主导导致有效采样退化，算法实时监控有效样本大小 \(\text{ESS} = (\sum w_i)^2 / \sum w_i^2\)；一旦 \(\text{ESS} < \alpha K\)（实验中设 \(\alpha = 0.5\)），立即主动清空蓄水池，从源头阻断方差发散与权重崩溃。

损失函数 / 训练策略¶

在具体下游应用中，ReSWD 采用了针对性的轻量化工程策略： - 色彩校正（Color Correction）：采用参数化电影级色彩决策列表（CDL），通过可微优化调节斜率（Slope）、偏移（Offset）、幂次（Power）与饱和度 \(\lambda\)。特征表示转入感知均匀的 CIELAB 空间，图像最大分辨率下采样至 128 像素，在仅 150 步内即可收敛并无损迁移至全分辨率图像或超高清视频片段。 - 扩散模型色彩引导（Diffusion Guidance）：针对 Flow Matching 架构 SD3.5（Medium/Large/Turbo）及 SDXL，在反向去噪的每一步预测 \(x_0\) 并由 VAE 解码至 CIELAB 空间与参考图进行 ReSWD 对齐。为解决大型 Transformer 架构显存爆炸与梯度开销问题，采用类似 SDS 的梯度截断技术（Gradient Stop），阻断骨干网络反传，仅对轻量 VAE 解码器回传梯度，配合 Adam 优化器（学习率 3e-3）在每个扩散步迭代 6 次微调潜在变量偏移。

实验关键数据¶

主实验¶

论文在合成一维/三维分布对齐（1000 组正态、均匀、双峰混合分布）、真实场景电影级色彩对齐以及文本引导扩散生成色彩控制三个基准上进行了系统测试。

在合成通用分布对齐基准（\(d=3, N=1024\)，迭代 300 步，总投影数固定为 64）中，ReSWD 在误差与梯度平稳度上均显著优于传统 SWD、控制变量法（CV）、极大化切片（Max-SW）以及拟蒙特卡洛（QMC）基线：

方法	平均 \(W_1\) 误差 [\(10^{-3}\)] \(\downarrow\)	标准差 \(\sigma\) [\(10^{-3}\)]	95% 置信区间	梯度信噪比 (SNR) \(\uparrow\)	单步耗时 [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 (本文)	0.622	0.075	\(\pm\) 0.005	0.278	1.92
ReSWD + QMC	0.610	0.076	\(\pm\) 0.005	0.274	2.10

在真实图像色彩校正基准（10 个复杂光照场景结合标准色卡）中，ReSWD 与传统色彩迁移方法（Reinhard、Nguyen）及深度神经方法（Yoo、Ho、Larchenko）进行了对比：

方法	色卡 PSNR \(\uparrow\)	变换误差 (RMSE) \(\downarrow\)	CDL 拟合误差 \(\downarrow\)	CTQM 色彩质量 \(\uparrow\)	单次耗时 [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 (本文)	24.64 \(\pm\) 5.70	0.31 \(\pm\) 0.31	0.10 \(\pm\) 0.07	5.17 \(\pm\) 1.80	5

在扩散引导（ContraStyles 提示词与 Unsplash Lite 真实参考图）方面，ReSWD 在 SD3.5-Turbo 上仅耗时 4 秒，且色差 \(W_2\) 降至 0.675，而在 SD3.5-Large 上更进一步将色差降至 0.55，远超基线 Lobashev et al. 的 124 秒与 1.941 \(W_2\) 误差。

消融实验¶

论文在总投影预算固定为 64 的前提下，系统分析了单步新抽取候选方向数量 \(M\)（即蓄水池容量 \(K = 64 - M\)）对分布匹配收敛效果的影响：

单步新候选数量 \(M\)	平均 \(W_1\) 误差 [\(10^{-3}\)] \(\downarrow\)	标准差 \(\sigma\) [\(10^{-3}\)]	95% 置信区间	耗时 [ms] \(\downarrow\)	分析说明
\(M=2\)	0.721	0.257	\(\pm\) 0.016	1.99	候选过少导致蓄水池新陈代谢停滞，方向僵化
\(M=4\)	0.673	0.151	\(\pm\) 0.009	1.96	探索能力略显不足，收敛速度次优
\(M=8\) (最优配置)	0.622	0.075	\(\pm\) 0.005	1.92	历史重用（\(K=56\)）与新方向探索达到最佳平衡
\(M=16\)	0.746	0.122	\(\pm\) 0.008	1.91	历史保留容量不足，方差缩减增益削弱
\(M=32\)	1.192	1.283	\(\pm\) 0.080	1.98	蓄水池退化明显，梯度平稳度显著下降
\(M=48\)	2.122	3.280	\(\pm\) 0.203	1.93	丧失蓄水池累积优势，方差迅速扩大
\(M=56\)	3.811	3.317	\(\pm\) 0.206	1.85	几乎退化为纯蒙特卡洛重采样，震荡加剧

在扩散引导的模块消融中，在 SDXL 骨干上逐项移除组件： - 完整 ReSWD @ SDXL：Mean-\(W_2\) 为 \(1.213 \times 10^{-2}\)，耗时 34s。 - 移除 Adam（退回普通梯度下降）：Mean-\(W_2\) 恶化至 \(1.846 \times 10^{-2}\)，耗时 34s。 - 移除 Gradient Stop 截断（对整个 U-Net 反传）：Mean-\(W_2\) 为 \(1.832 \times 10^{-2}\)，耗时剧增至 126s，证明骨干网络梯度截断在保持对齐精度的同时带来了将近 4 倍的端到端提速。

关键发现¶

新候选方向配比的拐点效应：候选数 \(M=8\) 构成了探索与利用的最优平衡点。\(M\) 过小会导致蓄水池被早期方向“锁死”，缺乏摆脱局部极值的能力；\(M\) 过大则削弱了加权蓄水池降低方差的核心红利。
高维空间的维度诅咒瓶颈：随着问题特征维度上升（如表 5 所示，从 \(d=3\) 扩展至 \(d=64\)），随机击中有效几何切片的概率按超球面表面积呈指数级缩减。当维度 \(d \ge 12\) 时，蓄水池因为难以抽中更高质量的新方向而陷入停滞，导致 ReSWD 表现反而逊于普通 SWD（\(d=12\) 时 ReSWD 误差为 9.768，而 SWD 为 1.190）。这揭示了该方法在低维紧凑特征（如 3 维色彩空间）上的高度适应性与在高维嵌入上的理论边界。

亮点与洞察¶

跨领域技术迁移：将渲染领域解决蒙特卡洛积分方差的经典重采样算法（ReSTIR / WRS）跨界引入最优传输与生成模型优化，思路简洁优雅且极具工程实效。
严格保持无偏性：不同于 Max-SW 这类强行追求极值导致理论性质破坏的启发式方法，ReSWD 通过重要性采样的倒数自归一化权重与梯度截断，在享有极低方差梯度的同时维持了数学上的无偏估计。
轻量即插即用：相较于常规 SWD，仅引入轻微的标量指数衰减与随机键比较，耗时仅增加不到 1 毫秒，却换取了超过 29% 的梯度信噪比提升与更稳健的收敛轨迹。

局限与展望¶

高维任务适应性不足：正如作者坦承，在高维隐空间（如 \(d > 10\)）中，随机超球面采样的稀疏性导致蓄水池极难更新，容易退化或过早过拟合陈旧切片，未来需探索基于学习或主动采样的自适应候选提议分布。
卷积核扩展受阻：将标量投影矩阵推广至结构化可学习卷积核时实验未获成功，由于参数空间过大，随机均匀候选搜索无法稳定捕获优质空域滤波器。
超参数敏感度：衰减常数 \(\tau\) 与重置阈值 \(\alpha\) 需根据具体的非平稳优化速率手工微调，自适应学习率联动机制有待进一步探索。

评分¶

新颖性: ⭐⭐⭐⭐ [成功将图形学路径追踪重采样机制与最优传输切片距离跨界融合，设计紧凑自洽]
实验充分度: ⭐⭐⭐⭐ [覆盖合成分布对齐、真实色彩工程校正与现代扩散模型色彩引导，多维度深入消融]
写作质量: ⭐⭐⭐⭐⭐ [逻辑严密，数学推导清晰，算法描述规范，可视化与消融详实]
价值: ⭐⭐⭐⭐ [对低维色彩传输与扩散引导具有即插即用的工程价值，为解决可微蒙特卡洛方差提供了范例]