[論文レビュー] 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

ひとこと要約

本稿では、ヘンの法を統合することで、コストのかかるヘッセ行列の計算を回避しながら2次精度を達成する、離散潜在変数向けの新たな勾配推定法ReinMaxを提案する。本稿では、ストレートスラッシング（ST）を1次勾配近似として形式的に確立し、最先端手法と比較して顕著な性能向上を示す。計算コストの増加はほとんどない。

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.

研究の動機と目的

• 離散潜在変数の文脈において、ストレートスラッシング（ST）ヒューリスティックを勾配近似法として形式的に分析すること。
• 現在、試行錯誤によるチューニングに依存しているが、STおよびその変種における理論的根拠の欠如とハイパーパrameterの指針の不足を解消すること。
• ヘッセ行列や2次導関数を必要とせず、より高次精度を達成する新しい勾配推定法の開発。
• 多項式プログラミング、変分オートエンコーダ、構造予測など、多様なタスクにおいて一貫した性能向上を示すこと。
• 深層学習における離散的または非微分可能なコンponentsのための勾配推定器の今後の設計を支援する理論的知見を提供すること。

提案手法

• ReinMaxは、勾配推定プロセスを常微分方程式（ODE）としてモデル化し、2次精度の数値積分スキームであるヘンの法を適用することで導出される。
• 2段階の予測を用いる：まず現在のパラメータ値を用いて予測勾配を計算し、次に予測された中間状態を用いてそれを精錬する。
• 1次勾配とモデル出力のみに依存することで、明示的なヘッセ行列の計算を回避し、計算効率を維持する。
• 微分可能サンプリング機構と、予測勾配と実際の勾配の差分を用いた勾配補正ステップを統合する。
• 分散低減のため、期待損失 $\mathbb{E}[f(\mathbf{D})]$ を用いたベースライン差し引き戦略を採用し、安定性を向上させる。
• 1ホットサンプリングと微分可能なソフト近似を組み合わせた修正されたフォワードパスにより、離散変数を介したエンドツーエンドのバックプロパゲーションを可能にする。

実験結果

リサーチクエスチョン

• RQ1 ストレートスラッシング（ST）ヒューリスティックは、形式的な勾配近似法とどのように関係しているか？
• RQ2 ヘッセ行列や2次導関数を必要とせず、2次精度の勾配推定器を構築することは可能か？
• RQ3 STおよびその変種の最適なハイパーパrameter設定（例：温度）は何か？理論がこの選択をガイドできるか？
• RQ4 ReinMaxのような高次勾配推定器は、離散変数を含む多様な学習タスクにおいて、一貫して1次法よりも優れた性能を示すか？
• RQ5 勾配推定におけるベースラインの選択は、離散潜在変数モデルの収束性と性能にどのように影響するか？

主な発見

• ReinMaxはヘンの法を統合することで、真の勾配のテイラー展開を2次まで正確に再現する2次精度の勾配近似を達成する。
• ストレートスラッシング（ST）法が、1次数値ODEソルバである前進オイラー法と形式的に同等であることが示され、その使用に対する理論的裏付けが得られた。
• MNIST、ファッショニストMNIST、オムニグラット、多項式プログラミングタスクを含む複数のベンチマークで、REINFORCEベースの手法や最先端技術（RODEOなど）をすべて上回る性能を示した。
• 2^48次元の潜在空間を持つMNIST-VAEにおいて、ReinMaxはあらゆる温度設定でRODEOや他のベースラインを常に上回るELBOスコアを達成した。
• バッチサイズが小さく、潜在変数が多く、目的関数が複雑な状況でも優れた性能を示した。REINFORCE系手法は分散が高いため、こうした状況で苦戦するが、ReinMaxはその問題を克服した。
• ReinMaxは計算コストの増加がほとんどなく、STや他の1次法と同等の時間的・記憶的消費量を示した。一方、分散低減ベースライン（例：GR-MCK）はモンテカルロサンプル数の増加に伴いスケーリングが著しく悪化するのとは対照的である。