[Paper Review] Fixed Point Networks: Implicit Depth Models with Jacobian-Free Backprop.

Samy Wu Fung, Howard Heaton+4 more|arXiv (Cornell University)|Mar 23, 2021

Model Reduction and Neural Networks36 references15 citations

TL;DR

This paper introduces Fixed Point Networks (FPNs), a novel implicit depth model that guarantees convergence to a unique fixed point via network weights and input data. It proposes a Jacobian-free backpropagation (JFB) scheme that eliminates the need to solve costly Jacobian equations, enabling fixed memory training with lower computational cost while achieving state-of-the-art performance in classification tasks.

ABSTRACT

A growing trend in deep learning replaces fixed depth models by approximations of the limit as network depth approaches infinity. This approach uses a portion of network weights to prescribe behavior by defining a limit condition. This makes network depth implicit, varying based on the provided data and an error tolerance. Moreover, existing implicit models can be implemented and trained with fixed memory costs in exchange for additional computational costs. In particular, backpropagation through implicit depth models requires solving a Jacobian-based equation arising from the implicit function theorem. We propose fixed point networks (FPNs), a simple setup for implicit depth learning that guarantees convergence of forward propagation to a unique limit defined by network weights and input data. Our key contribution is to provide a new Jacobian-free backpropagation (JFB) scheme that circumvents the need to solve Jacobian-based equations while maintaining fixed memory costs. This makes FPNs much cheaper to train and easy to implement. Our numerical examples yield state of the art classification results for implicit depth models and outperform corresponding explicit models.

Motivation & Objective

• To address the computational inefficiency of existing implicit depth models that require solving complex Jacobian-based equations during backpropagation.
• To develop a training method for implicit depth models that maintains fixed memory usage without incurring high computational overhead.
• To ensure convergence to a unique fixed point in forward propagation through a stable, weight- and data-dependent mechanism.
• To enable practical deployment of implicit depth models by simplifying implementation and reducing training cost.

Proposed method

• FPNs define a forward pass as the limit of iterative inference, where network weights and input data jointly determine a unique fixed point.
• The method leverages the contraction mapping theorem to guarantee convergence of the iterative process to a unique solution.
• A novel Jacobian-free backpropagation (JFB) scheme is introduced, avoiding explicit Jacobian computation by using a recursive gradient update based on the fixed point iteration.
• The JFB scheme maintains constant memory usage by storing only the final fixed point and intermediate activations, not the full computation graph.
• Backward pass is computed via a recursive formula derived from the implicit function theorem, without solving linear systems involving the Jacobian.
• The framework is compatible with standard deep learning optimization tools and can be implemented with minimal modifications to existing training pipelines.

Experimental results

Research questions

• RQ1 Can implicit depth models be trained efficiently without solving computationally expensive Jacobian-based equations during backpropagation?
• RQ2 Does a fixed point iteration scheme in deep networks guarantee convergence to a unique solution under practical conditions?
• RQ3 Can a Jacobian-free backpropagation method maintain fixed memory usage while reducing training cost compared to existing implicit models?
• RQ4 How does the performance of FPNs compare to explicit depth models and other implicit depth architectures on standard benchmarks?

Key findings

• FPNs achieve state-of-the-art classification accuracy among implicit depth models, outperforming both explicit and implicit baselines.
• The proposed Jacobian-free backpropagation (JFB) scheme eliminates the need to solve linear systems involving the Jacobian, significantly reducing training cost.
• The method maintains fixed memory usage throughout training, enabling efficient training on large-scale datasets.
• Convergence to a unique fixed point is guaranteed under mild conditions, ensuring stable and predictable forward inference.
• The implementation of FPNs is simple and compatible with standard deep learning frameworks, facilitating widespread adoption.
• Numerical experiments confirm that FPNs achieve superior performance with lower computational overhead compared to existing implicit depth models.