JFB: Jacobian-Free Backpropagation for Implicit Networks

TL;DR

This paper introduces Jacobian-Free Backpropagation (JFB), a novel method that enables efficient, fixed-memory training of implicit networks by avoiding costly Jacobian calculations, making implicit networks more practical and scalable.

Contribution

JFB is the first fixed-memory backpropagation method for implicit networks that simplifies implementation and accelerates training without losing accuracy.

Findings

JFB reduces training time for implicit networks.

Implicit networks trained with JFB match the accuracy of feedforward networks.

JFB enables scalable implicit network training with constant memory usage.

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.