[논문 리뷰] Efficient and Modular Implicit Differentiation

Mathieu Blondel, Quentin Berthet외 6명|arXiv (Cornell University)|2021. 05. 31.

Advanced Optimization Algorithms Research참고 문헌 70인용 수 37

한 줄 요약

이 논문은 자동 암시적 미분(auto implicit differentiation)을 소개합니다. 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를 기반으로 솔버 위에 암시적 미분을 부착한다.
• 내부 최적화가 θ에 따라 x*(θ)로 연결되도록 선형 시스템 -∂1F(x*(θ),θ) ∂x*(θ) = ∂2F(x*(θ),θ) 를 적용한다.
• 행렬-자유 선형 솔버(CG, GMRES, BiCGSTAB)로 자이코비-벡터곱(JVP) / 벡터-자이코비곱(VJP)을 효율적으로 계산한다.
• 미분을 다른 신경망이나 손실 연산과 합성할 수 있도록 전처리/후처리 매핑을 지원한다.
• 다양한 최적조건 매핑(정지점, KKT, proximal gradient 고정점, projected gradient 고정점 등)의 실용 구현을 제공한다.

실험 결과

연구 질문

• RQ1 사용자 정의 F를 통해 자동 암시적 미분이 광범위한 최적 조건의 카탈로그를 다룰 수 있는가?
• RQ2 내부 최적화가 근사적으로 해결되었을 때 자이코비 정밀도 보장은 무엇인가?
• RQ3 다른 솔버와 고정점 표현에 대해 unrolling 방법과 비교해 효율성과 유연성은 어떻게 되는가?
• RQ4 하이퍼파라미터 최적화, 데이터셋 증류, 딕셔너리 학습과 같은 이중 문제를 이 방법으로 쉽게 다룰 수 있는가?
• RQ5 일반적인 최적화 스킴(프로ximal, projection, mirror descent)에서의 구현 및 차분을 위한 실용 가이드라인은 무엇인가?

주요 결과

• 이 프레임워크는 잔여 맵 F를 자동 미분으로 미분하고 암시적 함수 정리를 적용함으로써 최적화 문제 해를 차분화할 수 있게 한다.
• 내부 해결의 근사로 인한 자이코비 잔차 오차는 경계에 있으며 내부 해의 잔차와 비례해 스케일링된다(정리 1).
• 이 접근 방식은 기존의 암시적 미분 방법을 재현하고 하나의 모듈식 시스템 내에서 새로운 방법들을 가능하게 한다.
• 실험을 통해 다중분류 SVM의 하이퍼파라미터 최적화, 데이터셋 증류, 작업 주도 딕셔너리 학습, 분자 동역학 민감도 분석에 대해 효율적인 차분을 보인다.
• 고정점 형식(mirror descent, proximal gradient, projected gradient)을 암시적 미분과 함께 사용하면 실용적이고 확장 가능한 이중 최적화 워크플로우를 얻을 수 있다.
• 이 방법은 여러 이중 작업에서 unrolling에 비해 실행 시간이 우수해 실용적인 속도와 단순성을 보여준다.