JFB: Jacobian-Free Backpropagation for Implicit Networks

2021/03/23 by Samy Wu Fung, Fung, Samy Wu, Howard Heaton +9 · 11 citations
Engineering · Physics and Astronomy · #Dam Engineering and Safety #FOS: Computer and information sciences #Machine Learning (cs.LG) #Model Reduction and Neural Networks

paper · pdf · doi:10.48550/arxiv.2103.12803

openalex publication_date 2021/03/23 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

A promising trend in deep learning replaces traditional feedforward networks with implicit networks. Unlike traditional networks, implicit networks solve a fixed point equation to compute inferences. Solving for the fixed point varies in complexity, depending on provided data and an error tolerance. Importantly, implicit networks may be trained with fixed memory costs in stark contrast to feedforward networks, whose memory requirements scale linearly with depth. However, there is no free lunch -- backpropagation through implicit networks often requires solving a costly Jacobian-based equation arising from the implicit function theorem. We propose Jacobian-Free Backpropagation (JFB), a fixed-memory approach that circumvents the need to solve Jacobian-based equations. JFB makes implicit networks faster to train and significantly easier to implement, without sacrificing test accuracy. Our experiments show implicit networks trained with JFB are competitive with feedforward networks and prior implicit networks given the same number of parameters.

Cited by

• Equivariant Deep Equilibrium Models for Imaging Inverse Problems
• Gradient descent for deep equilibrium single-index models
• Bilevel Learning via Inexact Stochastic Gradient Descent
• Learning Regularization Functionals for Inverse Problems: A Comparative Study
• End-to-End Training of High-Dimensional Optimal Control with Implicit Hamiltonians via Jacobian-Free Backpropagation
• Reversible Deep Equilibrium Models
• DEQuify your force field: More efficient simulations using deep equilibrium models
• Self-Composing Neural Operators with Depth and Accuracy Scaling via Adaptive Train-and-Unroll Approach
• Learning Binary Sampling Patterns for Single-Pixel Imaging using Bilevel Optimisation
• Minimizing Surrogate Losses for Decision-Focused Learning using Differentiable Optimization
• Unsupervised Deep Equilibrium Model Learning for Large-Scale Channel Estimation with Performance Guarantees

Related

• Training Implicit Networks for Image Deblurring using Jacobian-Free Backpropagation
• Implicit Deep Learning
• Test like you Train in Implicit Deep Learning
• Efficient Training of Deep Equilibrium Models
• Robust Implicit Backpropagation