[論文レビュー] Efficient and Modular Implicit Differentiation

Mathieu Blondel, Quentin Berthet+6 more|arXiv (Cornell University)|May 31, 2021

Advanced Optimization Algorithms Research参考文献 70被引用数 37

ひとこと要約

本論文は自動的な暗黙微分を導入します。Python/JAXフレームワークで、最適性条件 F を指定することにより最適化問題の解を微分可能にし、既存ソルバーの上にモジュラーでソルバー非依存の微分を実現します。ヤコビ行列の誤差保証と複数の応用を提供します。

ABSTRACT

Automatic differentiation (autodiff) has revolutionized machine learning. It allows to express complex computations by composing elementary ones in creative ways and removes the burden of computing their derivatives by hand. More recently, differentiation of optimization problem solutions has attracted widespread attention with applications such as optimization layers, and in bi-level problems such as hyper-parameter optimization and meta-learning. However, so far, implicit differentiation remained difficult to use for practitioners, as it often required case-by-case tedious mathematical derivations and implementations. In this paper, we propose automatic implicit differentiation, an efficient and modular approach for implicit differentiation of optimization problems. In our approach, the user defines directly in Python a function $F$ capturing the optimality conditions of the problem to be differentiated. Once this is done, we leverage autodiff of $F$ and the implicit function theorem to automatically differentiate the optimization problem. Our approach thus combines the benefits of implicit differentiation and autodiff. It is efficient as it can be added on top of any state-of-the-art solver and modular as the optimality condition specification is decoupled from the implicit differentiation mechanism. We show that seemingly simple principles allow to recover many existing implicit differentiation methods and create new ones easily. We demonstrate the ease of formulating and solving bi-level optimization problems using our framework. We also showcase an application to the sensitivity analysis of molecular dynamics.

研究の動機と目的

• ユーザーが最適性条件を直接Pythonで指定できるようにして、暗黙微分の利用障壁を下げる。
• 暗黙微分と自動微分を組み合わせ、ソルバーを再実装することなく最適化解を微分する。
• 最先端のソルバーとさまざまな最適性条件に対応するモジュラーなフレームワークを提供する。
• 近似解に対する理論的なヤコビ行列の精度保証を提供する。
• 多様な応用を通じて実用的な二階最適化と感度分析を実証する。

提案手法

• ユーザーは、アルゴリズムが解く問題の最適性条件を捉えるマッピング F を定義する。
• Pythonデコレータ (@custom_root) を用いて、F に基づく暗黙微分をソルバーの上につける。
• 暗黙の関数定理を適用し、線形系 -∂1F(x*(theta),theta) ∂x*(theta) = ∂2F(x*(theta),theta) を介して x*(theta) を theta に関連付ける。
• 行列表なしの線形ソルバー (CG, GMRES, BiCGSTAB) を用いて、ヤコビ行列-ベクトル積(JVP)とベクトル-ヤコビ積(VJP)を効率的に計算する。
• 微分を他のニューラルネットワークや損失演算と組み合わせられるよう、前処理/後処理のマッピングをサポートする。
• さまざまな最適性条件マッピング（ stationary、KKT、 proximal gradient fixed point、 projected gradient fixed point など）の実用的実装を提供する。

実験結果

リサーチクエスチョン

• RQ1 ユーザー定義の F を用いることで、自動的な暗黙微分は広範な最適性条件をカバーできるか？
• RQ2 内部最適化が近似的に解かれた場合、ヤコビ行列の精度保証はどうなるか？
• RQ3 提案されたフレームワークは、異なるソルバーと不動点表現に対して、展開法（unrolling）と比べて効率性と柔軟性でどうなるか？
• RQ4 ハイパーパラメータ最適化、データセット蒸留、辞書学習などの二階問題を容易に扱えるか？
• RQ5 このフレームワーク内で、一般的な最適化スキーム（proximal、projection、mirror descent）を実装し微分するための実践的ガイドラインは何か？

主な発見

• 本フレームワークは、残差マップ F を自動微分で微分し、暗黙の関数定理を適用することで最適化問題の解を微分可能にする。
• 近似解からのヤコビ行列誤差は有界で、内部解の残差に比例してスケールする（Theorem 1）。
• このアプローチは既存の暗黙微分法を回収し、単一のモジュラーなシステム内で新たな方法を可能にする。
• 実験は、多クラスSVMのハイパーパラメータ最適化、データセット蒸留、タスク駆動辞書学習、分子動力学の感度分析の効率的な微分を示す。
• 固定点定式化（mirror descent、proximal gradient、projected gradient）と暗黙微分を組み合わせた実用的で拡張性のある二階最適化ワークフローを生む。
• いくつかの二階タスクで展開法と比較して実行時間が有利で、速度と単純さの実用的な利得を示す。