[論文レビュー] Randomized Automatic Differentiation

Deniz Oktay, Nick McGreivy+3 more|arXiv (Cornell University)|Jul 20, 2020

Gaussian Processes and Bayesian Inference参考文献 34被引用数 8

ひとこと要約

この論文は、逆方向自動微分（AD）の過程にランダム化を導入することで、正確な逆方向ADを不偏で低メモリな勾配推定器に置き換えるフレームワーク、ランダム化自動微分（RAD）を紹介する。メモリを犠牲にして分散を増やすことで、確率的最適化における収束速度が向上し、前向き型および再帰型ニューラルネットワーク、および反応拡散型PDEを用いた原子炉制御に応用された。その結果、顕著なメモリ節約が達成された。

ABSTRACT

The successes of deep learning, variational inference, and many other fields have been aided by specialized implementations of reverse-mode automatic differentiation (AD) to compute gradients of mega-dimensional objectives. The AD techniques underlying these tools were designed to compute exact gradients to numerical precision, but modern machine learning models are almost always trained with stochastic gradient descent. Why spend computation and memory on exact (minibatch) gradients only to use them for stochastic optimization? We develop a general framework and approach for randomized automatic differentiation (RAD), which can allow unbiased gradient estimates to be computed with reduced memory in return for variance. We examine limitations of the general approach, and argue that we must leverage problem specific structure to realize benefits. We develop RAD techniques for a variety of simple neural network architectures, and show that for a fixed memory budget, RAD converges in fewer iterations than using a small batch size for feedforward networks, and in a similar number for recurrent networks. We also show that RAD can be applied to scientific computing, and use it to develop a low-memory stochastic gradient method for optimizing the control parameters of a linear reaction-diffusion PDE representing a fission reactor.

研究の動機と目的

• 従来のADが得る正確な勾配と、深層学習における現代の最適化の確率的性質との不一致を解消すること。
• 逆方向ADにおけるメモリ使用量を削減するために、分散を犠牲にして記憶領域を削減するランダム化推定器を導入すること。
• ニューラルネットワークおよび物理シミュレータに適用可能な一般化されたフレームワークとして、不偏で低メモリな勾配推定を構築すること。
• 固定メモリ予算下で、小バッチ学習よりもRADがより速く収束することを実証すること。
• 高コストなメモリを要する大規模な科学計算問題にADの適用範囲を拡大すること。

提案手法

• 計算グラフのパスをランダムにサンプリングすることで、不偏な勾配推定を実現する逆方向ADのランダム化変種を提案する。
• エッジ重み（局所微分）を確率的に部分抽出することで、メモリ削減を図る線形化計算グラフ（LCG）にランダム化を適用する。
• バウアーの公式を用いてLCG内のパス列挙によりヤコビ行列を計算し、中間微分の保存に必要な記憶領域を削減するためのランダムパス選択を実施する。
• LCG内のパスのサブセットを選択するサブサンプリング方式を導入し、不偏性を維持しながらメモリフットプリントを削減する。
• 計算が困難なグラフに対しては、有効な勾配推定を可能にするため、テレスコピング級数推定器を採用する。
• RADを前向き型および再帰型ニューラルネットワーク、および線形反応拡散PDEに適用し、メモリ効率の高さを実証する。

実験結果

リサーチクエスチョン

• RQ1 ランダム化ADは、不偏な勾配推定を維持しつつ、逆方向微分におけるメモリ使用量を削減できるか？
• RQ2 固定メモリ予算下で、RADは小バッチ学習と比較して収束速度が速いか？
• RQ3 反応拡散型PDE最適化のような高メモリ要件を有する科学計算問題に、RADを効果的に適用できるか？
• RQ4 計算グラフのどのような構造的性質が、バイアスを生じさせることなく効果的なランダム化を可能にするか？
• RQ5 RAD推定器の分散は、標準ADと比較してどの程度で、効果的に制御可能か？

主な発見

• 前向き型ネットワークでは、同じメモリ予算下でRADは小バッチ学習よりも少ない反復回数で収束した。
• 再帰型ネットワークでは、わずかなメモリ使用量の削減にもかかわらず、RADは小バッチ学習と同等の反復回数で収束した。
• RADは、原子炉制御のための線形反応拡散PDEの低メモリ最適化を可能にし、科学計算への適用可能性を実証した。
• フレームワークは、線形化計算グラフにおけるパスのサブサンプリングにより、不偏な勾配推定を維持しながら顕著にメモリ要件を削減した。
• 標準ADに比べ、問題固有の構造を活用することでバイアスを生じさせることなく記憶領域を削減できるため、メモリ制約下でRADは優れた性能を示した。
• 本手法は一般化可能であり、非線形性やループを含む他の計算グラフに対しても、ランダム化ラウンドおよびテレスコピング推定器を用いて拡張可能である。