[论文解读] Bridging Discrete and Backpropagation: Straight-Through and Beyond

Liyuan Liu, Chengyu Dong+3 more|arXiv (Cornell University)|Apr 17, 2023

Machine Learning and Data Classification被引用 4

一句话总结

本文提出ReinMax，一种针对离散潜在变量的新型梯度估计方法，通过整合Heun方法实现二阶精度，避免了昂贵的Hessian矩阵计算。该方法形式化地将Straight-Through（ST）视为一阶梯度近似，并在多项任务中证明ReinMax在计算开销可忽略的前提下，性能优于当前最先进方法。

ABSTRACT

Backpropagation, the cornerstone of deep learning, is limited to computing gradients for continuous variables. This limitation poses challenges for problems involving discrete latent variables. To address this issue, we propose a novel approach to approximate the gradient of parameters involved in generating discrete latent variables. First, we examine the widely used Straight-Through (ST) heuristic and demonstrate that it works as a first-order approximation of the gradient. Guided by our findings, we propose ReinMax, which achieves second-order accuracy by integrating Heun's method, a second-order numerical method for solving ODEs. ReinMax does not require Hessian or other second-order derivatives, thus having negligible computation overheads. Extensive experimental results on various tasks demonstrate the superiority of ReinMax over the state of the art. Implementations are released at https://github.com/microsoft/ReinMax.

研究动机与目标

• 正式分析Straight-Through（ST）启发式方法在离散潜在变量背景下的梯度近似性质。
• 解决ST及其变体目前缺乏理论基础和超参数指导的问题，这些方法当前依赖试错调参。
• 开发一种新型梯度估计方法，实现高阶精度，同时无需计算Hessian矩阵或二阶导数。
• 在多项任务中（如多项式编程、变分自编码器和结构化预测）展示一致的性能提升。
• 提供理论洞见，以指导未来在深度学习中对离散或不可微组件的梯度估计器设计。

提出的方法

• ReinMax通过将梯度估计过程建模为常微分方程（ODE），并应用Heun方法——一种二阶数值积分方案——推导得出。
• 该方法采用两阶段预测：首先使用当前参数值计算预测梯度，然后利用预测的中间状态进行修正。
• 通过避免显式计算Hessian矩阵，保持计算效率，仅依赖一阶梯度和模型输出。
• 该方法结合可微采样机制与基于预测梯度与实际梯度差异的梯度校正步骤。
• ReinMax采用基线减法策略，利用期望损失 $\mathbb{E}[f(\mathbf{D})]$ 降低方差，提升稳定性。
• 通过修改前向传播过程实现，结合one-hot采样与可微软近似，实现通过离散变量的端到端反向传播。

实验结果

研究问题

• RQ1 Straight-Through（ST）启发式方法与正式的梯度近似方法之间有何关系？
• RQ2 能否在不依赖Hessian矩阵或二阶导数的前提下，构建一个二阶精度的梯度估计器？
• RQ3 ST及其变体的最优超参数配置（如温度）是什么？理论能否指导这一选择？
• RQ4 像ReinMax这样的高阶梯度估计器是否能在涉及离散变量的多样化学习任务中持续优于一阶方法？
• RQ5 在梯度估计中基线的选择如何影响离散潜在变量模型的收敛性和性能？

主要发现

• ReinMax通过集成Heun方法，在梯度近似中实现二阶精度，其泰勒展开与真实梯度在二阶项内完全匹配。
• 形式化证明ST方法等价于前向欧拉法，即一阶数值ODE求解器，为ST的使用提供了理论依据。
• 在多个基准测试中，ReinMax优于所有基线方法，包括基于REINFORCE的方法和SOTA技术如RODEO，涵盖MNIST、Fashion-MNIST、Omniglot及多项式编程任务。
• 在具有2^48维潜在空间的MNIST-VAE中，ReinMax在所有温度设置下均持续获得高于RODEO及其他基线的ELBO分数。
• 在更具挑战性的设置中（如小批量、更多潜在变量、复杂目标函数），ReinMax表现更优，而REINFORCE类方法因方差过高而难以收敛。
• ReinMax具有可忽略的计算开销，其时间和内存消耗与ST及其他一阶方法相当，而不同于方差缩减基线（如GR-MCK）在蒙特卡洛样本数增加时性能急剧下降。