Reproduce “Efficient and Modular Implicit Differentiation” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-López, Fabian Pedregosa, Jean-Philippe Vert
Google Research developed an “automatic implicit differentiation” framework for JAX, enabling efficient and modular gradient computation through solutions of optimization problems. The framework automates the historically tedious process of implicit differentiation, achieving up to 4x faster computation for dataset distillation and significant memory savings for hyperparameter optimization compared to unrolling methods.
Key insights
- Automatic differentiation can be applied to the optimality conditions of an optimization problem, effectively automating the derivation of implicit gradients.
- Decoupling the definition of optimality conditions from the actual solver allows for modularity and the use of any state-of-the-art, optimized solver.
- Theoretical bounds show that the error on the estimated Jacobian is at most proportional to the error in the approximate optimal solution, validating the practicality of approximate solvers.
- Fixed-point iterations, common in many optimization algorithms, can be elegantly reformulated as root-finding problems for the implicit differentiation framework.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.