Bridging Discrete and Backpropagation: Straight-Through and Beyond
This work addresses a fundamental limitation in deep learning for problems involving discrete variables, offering a more efficient solution for researchers and practitioners in fields like natural language processing and reinforcement learning.
The paper tackles the challenge of backpropagation with discrete latent variables by proposing ReinMax, a method that achieves second-order gradient accuracy without significant computational overhead, outperforming state-of-the-art approaches in various tasks.
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.