[论文解读] ARSM: Augment-REINFORCE-Swap-Merge Estimator for Gradient Backpropagation Through Categorical Variables

Mingzhang Yin, Yuguang Yue+1 more|arXiv (Cornell University)|May 4, 2019

Adversarial Robustness in Machine Learning被引用 17

一句话总结

本文提出 ARSM，一种通过结合变量增强、REINFORCE、Rao-Blackwellization 及一种新颖的变量交换机制，在狄利克雷分布下构建等价梯度期望，从而实现对分类变量反向传播的无偏、低方差梯度估计器。该方法通过在这些期望之间共享随机数，实现显著的方差降低，在变分自编码器中表现优于现有估计器，并为离散强化学习提供了一种无需基线的‘试错自评’策略梯度方法。

ABSTRACT

To address the challenge of backpropagating the gradient through categorical variables, we propose the augment-REINFORCE-swap-merge (ARSM) gradient estimator that is unbiased and has low variance. ARSM first uses variable augmentation, REINFORCE, and Rao-Blackwellization to re-express the gradient as an expectation under the Dirichlet distribution, then uses variable swapping to construct differently expressed but equivalent expectations, and finally shares common random numbers between these expectations to achieve significant variance reduction. Experimental results show ARSM closely resembles the performance of the true gradient for optimization in univariate settings; outperforms existing estimators by a large margin when applied to categorical variational auto-encoders; and provides a "try-and-see self-critic" variance reduction method for discrete-action policy gradient, which removes the need of estimating baselines by generating a random number of pseudo actions and estimating their action-value functions.

研究动机与目标

• 解决深度学习与强化学习中通过分类变量反向传播时高方差梯度估计的挑战。
• 为具有 C ≥ 2 个类别的多变量分类变量开发一种无偏且保持低方差的梯度估计器。
• 通过伪动作滚动和价值函数估计引入自评机制，消除离散动作策略梯度中对基线估计的需求。
• 将原始用于二值变量的 ARM 估计器统一并扩展至一般分类情形（C ≥ 2）。
• 为训练具有离散潜在变量和离散控制策略的模型，提供一种实用且可复现的方法。

提出的方法

• ARSM 估计器使用变量增强，将梯度重表达为狄利克雷分布下的期望。
• 在增强空间中应用 REINFORCE 梯度估计器，获得初始无偏梯度估计。
• 通过条件化充分统计量应用 Rao-Blackwellization 以降低方差。
• 通过变量交换在狄利克雷分布下生成多个等价但表达方式不同的梯度期望。
• 通过在这些交换后的期望之间共享公共随机数，执行合并步骤以进一步降低方差。
• 最终估计器通过计算这些共享的等价期望的加权平均得到，实现显著的方差降低。

实验结果

研究问题

• RQ1 能否通过增强和方差减少技术，为 C ≥ 2 的分类变量构造出低方差无偏梯度估计器？
• RQ2 在狄利克雷分布下，变量交换如何生成等价的梯度期望，从而实现有效的方差降低？
• RQ3 ARSM 能否通过生成伪动作并估计其价值函数，消除离散动作策略梯度中对基线估计的需求？
• RQ4 ARSM 在训练分类变分自编码器和离散强化学习智能体方面，与现有估计器相比表现如何？
• RQ5 ARSM 在单变量设置下是否能接近真实梯度性能，并在多变量和复杂设置下达到最先进水平？

主要发现

• 在单变量设置中，ARSM 非常接近真实梯度，显示出梯度估计的高精度。
• 在分类变分自编码器中，ARSM 相较于现有估计器，实现了最先进水平的训练性能和样本外预测性能。
• 在离散强化学习中，ARSM 提供了一种‘试错自评’方法，通过生成并评估伪动作，消除了对基线估计的需求。
• 该方法通过在交换后的期望之间共享随机数，显著降低了梯度方差，在合成任务和真实世界任务中均优于 REINFORCE 和其他基线方法。
• 当 C ≥ 2 时，仅靠对称采样无法实现 ARSM，因此必须依赖新颖的变量交换机制，该机制对方差降低至关重要。
• 所提出的算法具有可复现性，相关代码已公开发布于 GitHub，涵盖 ARSM 梯度估计器及其在强化学习和 VAE 中的应用。