[論文レビュー] Neural Networks with Cheap Differential Operators

Ricky T. Q. Chen, David Duvenaud|arXiv (Cornell University)|Dec 8, 2019

Model Reduction and Neural Networks被引用数 12

ひとこと要約

この論文では、入力次元に依存せず、k階微分のための逆方向自動微分（reverse-mode AD）をk回実行するだけで、次元ごとの微分作用素（例：発散やヤコビアン対角成分）を効率的に計算できるニューラルネットワークアーキテクチャ「HollowNet」を提案する。ネットワークをコンditionerとTransformerの2つのコンponentに構造化することで、次元ごとの依存関係を分離し、これらの作用素への正確で低コストなアクセスを可能にし、暗黙のODE、連続正規化流れ、SDEのFokker–Planck一致の計算において顕著な効率向上を実現する。

ABSTRACT

Gradients of neural networks can be computed efficiently for any architecture, but some applications require differential operators with higher time complexity. We describe a family of restricted neural network architectures that allow efficient computation of a family of differential operators involving dimension-wise derivatives, used in cases such as computing the divergence. Our proposed architecture has a Jacobian matrix composed of diagonal and hollow (non-diagonal) components. We can then modify the backward computation graph to extract dimension-wise derivatives efficiently with automatic differentiation. We demonstrate these cheap differential operators for solving root-finding subproblems in implicit ODE solvers, exact density evaluation for continuous normalizing flows, and evaluating the Fokker--Planck equation for training stochastic differential equation models.

研究の動機と目的

• 標準的なニューラルネットワークにおける発散やヤコビアン対角成分といった微分作用素の計算コストが高いため、その問題を解決すること。
• 入力次元に依存せず、計算複雑度が増加しない次元ごとの微分への効率的かつスケーラブルなアクセスを可能にするニューラルネットワークアーキテクチャを設計すること。
• Fokker–Planck方程式の一致を用いて、SDEのパラメータ推定を効率的に行えるようにすること。
• 確率的微分方程式（SDE）のパラメータ推定を、Fokker–Planck方程式の一致を用いて行う。
• 暗黙のODEソルバーを支援するために、低コストな微分作用素を用いて根の探索サブプロブレムを効率的に計算できること。

提案手法

• ネットワークは、入力のうちx_iを除いたすべての入力を用いて隠れ表現h_iを計算するコンditionerネットワークに構造化され、h_iがx_iに依存しないようにする。
• その後、トランスフォーマーネットワークがf_i(x) = τ_i(x_i, h_i)を計算する。ここでτ_iはx_iとh_iの連結を入力とするニューラルネットワークである。
• このアーキテクチャにより、fのヤコビアン行列がx_iからf_iへの対角成分と、h_iからf_iへのハロー成分に分解され、微分の分離が可能になる。
• コンditionerネットワークを経由する勾配を除外するようにバックワード計算グラフを変更することで、次元ごとの微分のみが保持され、1回のバックワードパスで∂f_i/∂x_iを効率的に抽出できる。
• この方法は逆方向自動微分を活用しているが、勾配伝播を対角成分に限定することで、k階微分のための計算複雑度をO(d)からO(k)に削減する。
• このアプローチは標準的なディープラーニング演算と互換性があり、マスク付き重み行列を用いることで、全結合型および畳み込み型アーキテクチャの両方をサポートする。

実験結果

リサーチクエスチョン

• RQ1 発散やヤコビアン対角成分といった次元ごとの微分を効率的に計算できるニューラルネットワークアーキテクチャを設計できるか？
• RQ2 k階微分のための逆方向ADの評価回数をO(d)からO(k)に削減できるか？
• RQ3 このアーキテクチャにより、 stiff な暗黙のODEを解く際の高速な根の探索サブプロブレムの処理が可能になるか？
• RQ4 確率的微分方程式（SDE）のパラメータ推定を、Fokker–Planck方程式の一致を用いて行えるか？
• RQ5 希な観測でも、SDEのパラメータ推定においてFokker–Planck方程式の一致を用いた関数的推定法が有効に機能するか？

主な発見

• HollowNetは、入力次元dに依存せず、k階微分のための逆方向ADをk回実行するだけで、k階微分の次元ごとの微分を計算可能であり、計算複雑度をO(kd)からO(k)に削減する。
• 1回のバックワードパスで正確な発散とヤコビアン対角成分の計算が可能であり、高次元問題において顕著な効率性を発揮する。
• 連続正規化流れでは、Stochasticな尤度推定を用いずに正確な密度評価が可能となり、学習の安定性と精度が向上する。
• SDEのパラメータ推定において、Fokker–Planck一致を用いたHollowNetは、希なデータ環境下で離散化ベースの手法が失敗する状況でも、偽最大尤度法を上回る性能を示す。
• 希な観測下でも、Stochasticな振り子SDEのマルチモーダルなダイナミクスを正確に回復でき、正しい周辺分布と正弦波的挙動を再現する。
• この方法は非定常データに対しても頑健であり、希なサンプリングによる従来の離散化ベースのSDE推定が失敗する状況でも、正確性を維持する。