[论文解读] Rao-Blackwellized Stochastic Gradients for Discrete Distributions

Runjing Liu, Jeffrey Regier+3 more|arXiv (Cornell University)|Oct 10, 2018

Domain Adaptation and Few-Shot Learning参考文献 25被引用 5

一句话总结

本文提出了一种基于Rao-Blackwellization的元过程，以在不引入偏差的情况下降低离散分布的随机梯度方差，尤其在分布稀疏（即集中于少数类别）时效果显著。该方法对高概率类别进行解析求和，显著提升了半监督分类和像素注意力等任务中的收敛速度与最终性能，优于REINFORCE、NVIL、RELAX和Gumbel-softmax，在速度与准确率上均表现更优。

ABSTRACT

We wish to compute the gradient of an expectation over a finite or countably infinite sample space having $K \leq \infty$ categories. When $K$ is indeed infinite, or finite but very large, the relevant summation is intractable. Accordingly, various stochastic gradient estimators have been proposed. In this paper, we describe a technique that can be applied to reduce the variance of any such estimator, without changing its bias---in particular, unbiasedness is retained. We show that our technique is an instance of Rao-Blackwellization, and we demonstrate the improvement it yields on a semi-supervised classification problem and a pixel attention task.

研究动机与目标

• 在不引入偏差的情况下，降低对离散、高维或无限维样本空间期望的随机梯度估计方差。
• 开发一种通用的方差减少技术，适用于任何针对离散分布的无偏随机梯度估计器，尤其适用于概率质量函数稀疏的场景。
• 提升具有离散隐变量模型（如极端分类或注意力机制）中基于梯度优化的实用性，因为标准方法在这些场景中常因方差过高而表现不佳。
• 证明对高概率类别的解析求和可带来显著的方差减少，尤其在仅有少数类别具有显著质量时。

提出的方法

• 该方法应用Rao-Blackwellization将期望分解为两部分：高概率原子（精确求和）和剩余低概率原子（随机估计）。
• 采用一种元过程，可与任何现有无偏随机梯度估计器（如REINFORCE或带控制变量的REINFORCE）结合使用，以在保持无偏性的同时减少方差。
• 该算法识别在分布 $ q_{\eta}(z) $ 下概率最高的K个类别，精确计算其贡献，其余部分通过蒙特卡洛采样估计。
• 该方法具有即时性（anytime）特性，即随着计算资源的增加，方差减少程度也随之提升——资源越多，可精确求和的类别越多。
• 该方法基于统计理论，特别是Rao-Blackwellization，其在全方差定律下可保证方差减少。
• 在稀疏场景（如极端分类或注意力机制）中尤为有效，此时仅有少数类别是合理或活跃的。

实验结果

研究问题

• RQ1 能否在不引入偏差的情况下降低离散分布随机梯度估计的方差？
• RQ2 当应用于具有大量类别的单变量离散随机变量时，Rao-Blackwellization的效果如何？
• RQ3 在实际深度学习任务中，对高概率类别进行解析求和是否能带来显著的方差减少？
• RQ4 与REINFORCE、NVIL、RELAX和Gumbel-softmax等基线方法相比，所提方法在收敛速度和最终性能方面表现如何？
• RQ5 该方法能否高效应用于结构化或多元离散设置（如隐马尔可夫链或注意力机制）？

主要发现

• 在半监督分类任务中，Rao-Blackwellized REINFORCE估计器收敛更快，且达到的负ELBO低于REINFORCE、NVIL和RELAX。
• 在像素注意力任务中，该方法在约44秒内达到负ELBO为500，而RELAX耗时110秒，表明收敛速度显著更快。
• Gumbel-softmax因直通估计器引入的高偏差而未能收敛至合理ELBO，而所提方法保持稳定且准确。
• NVIL虽收敛迅速，但达到的ELBO劣于所提方法，表明更快收敛并不意味着性能更优。
• 在移动MNIST任务中，该方法在所有基线中实现了最小的负ELBO，证明其优化质量更优。
• 每轮训练的计算成本高于其他方法（15.4秒 vs. 8.7–11.1秒），但收敛速度与最终性能的提升足以证明该开销的合理性。