Source-linked AI summary
JFB: Jacobian-Free Backpropagation for Implicit Networks
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, Wotao Yin
TL;DR
Implicit networks compute inferences as fixed points, but training them can involve costly Jacobian-based linear solves. The paper introduces Jacobian-Free Backpropagation (JFB), which avoids those solves, and reports competitive implicit-network results against prior approaches with similar or fewer parameters.
Problem
Method
JFB trains implicit networks without Jacobian-based linear-system solves, using a backpropagation scheme that omits the Jacobian term.
Source: §5 · p. 7Results
JFB substantially outperforms ODE-based networks and MONs with similar or fewer parameters, and is competitive with Multiscale DEQs despite using less than a tenth as many parameters.
Source: Classification · p. 6Takeaways & Limitations
JFB enables implicit-network training with fixed memory costs regardless of depth and is described as easy to code and efficient in backpropagation.
Source: §5 · p. 7With Jacobian-based backpropagation, RΘ cannot contain batch normalization.
Source: §1 · p. 2
Abstract
from arXiv · showhide
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.
1 Why Implicit Networks?
Implicit networks model some outputs as fixed points and can represent feedforward networks while training with constant memory costs. Their practical trade-offs include slower inference and constraints on the operator used to compute fixed points and gradients.
Implicit networks for implicitly defined outputs
- Implicit networks suit problems whose desired output is naturally characterized as a fixed point, including examples from game theory and inverse problems.The paper illustrates this with an implicitly specified scalar relationship for which an explicit solution requires an infinite-series representation.Source: §1 · p. 2
“Infinite depth” with constant memory training
- Implicit networks avoid storing intermediate forward-pass quantities, enabling training with memory costs constant with respect to depth.The paper compares fixed-point solving to an “infinite depth” weight-tied, input-injected feedforward network.Source: §1 · p. 2
No loss of expressiveness
- Implicit networks are at least as expressive as feedforward networks; with affine SΘ and QΘ, their class contains all feedforward networks.The paper says universal approximation properties then follow from those of conventional deep neural models.Source: §1 · p. 2
Architectural limitations
- In practice, fixed-point algorithms constrain RΘ, and guaranteeing the existence and computability of parameter gradients places further constraints on it.The paper gives Jacobian-based backpropagation's exclusion of batch normalization as an example.Source: §1 · p. 2
Slower inference
- Fixed-point iteration often makes inference slower than in a feedforward network because it repeatedly evaluates RΘ.Source: §1 · p. 2
2 Implicit Network Formulation
The formulation maps data into a latent space, finds a fixed point of a data-conditioned operator, and maps that point to an inference. Under the stated contraction assumption, fixed-point iteration converges linearly to a unique point; the section also contrasts solver choices and related implicit formulations.
Implicit Network Formulation
- The implicit network maps data into latent space through QΘ, uses TΘ to define its latent operator, and maps the fixed point to inference space through SΘ.The formulation treats the spaces and mappings generally before making them concrete for each application.Source: §2 · p. 2
Implicit Network with Fixed Point Iteration
- Fixed-point iteration initializes a latent variable, updates it while its residual exceeds ε, and returns SΘ applied to the final iterate.The initialization may be any fixed quantity, such as the zero vector.Source: §2 · p. 3
Convergence
- Under Assumption 2.1, the update sequence converges linearly to the unique fixed point of TΘ(·; d).Assumption 2.1 requires the operator to be Lipschitz and, holding d fixed, a contraction.Source: §2 · p. 3
Alternative Approaches
- The formulation uses simple fixed-point iteration rather than Broyden’s method to avoid Jacobian-related computations in its solver.Broyden’s method updates an approximation to the Jacobian and solves a linear system at each iteration.Source: §2 · p. 3
- Anderson acceleration is reported to find the approximate fixed point faster than vanilla fixed-point iteration and Broyden’s method, while combining it with JFB is left for future work.Source: §2 · p. 3
Other Implicit Formulations
- Neural ODEs use black-box ODE solvers for forward and adjoint-equation computations, trading memory savings for added computation and requiring high-accuracy adjoint solves for a guaranteed descent direction.Source: §2 · p. 3
3 Backpropagation
JFB replaces the Jacobian-based linear solve used in implicit-network backpropagation with a preconditioned gradient. Under stated assumptions, its update is a descent direction, including for sufficiently accurate approximate fixed points.
3 Backpropagation
- JFB omits the Jacobian term and uses a preconditioned gradient, avoiding the linear-system solve used in implicit-network backpropagation.The paper presents this as a simpler alternative to Jacobian-based backpropagation.Source: §3 · pp. 3, 5
- Under the stated assumptions, p_Θ is a descent direction for the loss with respect to Θ.The conditioning portion of one assumption was not enforced in the experiments, which the authors hypothesize is acceptable because worst-case behavior rarely occurs in practice.Source: §3 · pp. 4–5
- For sufficiently accurate approximate fixed points, JFB still yields a descent direction for the loss.The authors state they know of no analogous error-tolerance results in the implicit-depth literature.Source: §3 · p. 5
- Across the reported tests on nearly identical models, JFB was competitive with Jacobian-based gradients, although its different path to minimizers need not yield the same results on nonconvex problems.The comparison used conjugate gradient with a maximum iteration count matched to the maximum forward-propagation depth.Source: §3 · p. 5
- The implementation computes the fixed point without building a gradient graph, then backpropagates through one explicit latent-space layer.The authors describe this as standard backpropagation through an explicit model with one latent-space layer.Source: §3 · p. 5
- JFB is a zeroth-order approximation to the Neumann series; unlike Neumann-RBP, it has a proved descent guarantee under the paper’s assumptions.The paper also notes that Neumann-RBP gives no guidelines for truncating its series.Source: §3 · p. 5
4 Experiments
Experiments compare JFB with implicit-network baselines, Jacobian-based backpropagation, and higher-order Neumann approximations. Across these comparisons, JFB is competitive in test accuracy while offering faster training or broader architecture support.
4 Experiments
- The experiments use ResNet-based networks; weight tuning encouraged contractive behavior without explicitly enforcing the stated assumption.The reported experiments ran on a single NVIDIA TITAN X GPU with 12GB RAM; training batches also satisfied the stated condition when using batch normalization.Source: §4 · p. 5
Classification
- JFB substantially outperforms ODE-based networks and MONs with similar or fewer parameters, and is competitive with Multiscale DEQs using less than one tenth as many parameters.Table 1 reports test accuracy for JFB-trained Implicit ResNets and comparison methods.Source: Table 1 · p. 6, Classification · p. 6
Comparison to Jacobian-based Backpropagation
- On CIFAR-10, JFB trains faster and gives better test accuracy than standard Jacobian-based backpropagation at fixed-point tolerance ϵ = 10−4.The comparison uses comparable network configurations; the Jacobian-based results were not tuned to their best ability, and that method omitted batch normalization because its linear solve did not converge.Source: Comparison to Jacobian-based… · p. 6, Figure 5 · p. 7, Table 2 · p. 7
Higher Order Neumann Approximation
- With fewer Neumann-series terms, training is faster; JFB (k = 0) is fastest and retains competitive test accuracy.Figure 6 compares MNIST training with different truncations of the series used to approximate the inverse Jacobian.Source: Higher Order Neumann Approximation · p. 7, Figure 6 · p. 7
5 Conclusion
The conclusion presents JFB as a simple approach for training implicit networks with fixed memory and efficient backpropagation. Experiments find competitive results, while stochastic settings remain a topic for future analysis.
5 Conclusion
- JFB enables fixed-memory training regardless of depth and efficient backpropagation by removing linear solves at each step.The conclusion also describes the scheme as easy to code and theoretically justified when fixed points are approximately computed.Source: §5 · p. 7
- Experiments show JFB yields competitive results for implicit networks.The authors identify extensions to imaging, phase retrieval, geophysics, and games as future directions.Source: §5 · p. 7
- Analyzing JFB in stochastic settings remains future work.The conclusion also points to extensions that satisfy additional constraints in several application areas.Source: §5 · p. 7
A Proofs
The proofs establish invertibility of JΘ and show that the proposed pΘ direction is a descent direction under the stated assumptions. They also establish the corollary’s existence claim by choosing ε=δ.
Lemma A.1
- The proof establishes (1 − γ)-coercivity of JΘ, then uses its trivial kernel to conclude that JΘ is invertible.The invertibility argument is part of the proof sequence for the stated results.Source: A Proofs · p. 10
Theorem 3.1
- Under Assumptions 2.1, 3.1, 3.2, and 3.3, pΘ forms a descent direction for the loss with respect to Θ.The proof verifies the descent inequality for the given weights and data.Source: A Proofs · pp. 12–13
- The proof identifies pΘ as a preconditioned gradient, with a preconditioner that need not be symmetric.Source: A Proofs · p. 12
Corollary 3.1
- The corollary’s existence claim is established by choosing ε = δ.The supplied proof passage states that this choice makes the referenced condition hold.Source: A Proofs · p. 14
B Classification Accuracy Plots
On MNIST and SVHN, the figures report that JFB is faster and yields better test accuracy than standard Jacobian-based backpropagation. The compared configurations are nearly identical, with latent-space batch normalization used for JFB.
MNIST
- On MNIST, JFB is reported faster and more accurate on test data than standard Jacobian-based backpropagation.The configurations are nearly identical, but JFB uses latent-space batch normalization.Source: Figure 7 · p. 14
SVHN
- On SVHN, JFB is reported faster and more accurate on test data than standard Jacobian-based backpropagation.The architectures/configurations are identical except for latent-space batch normalization used with JFB.Source: Figure 8 · p. 15
Implementation Notes
The implementation solves the Jacobian-based gradient system without explicitly constructing the Jacobian, using normal equations and conjugate gradients. Autograd vector-Jacobian products provide the required multiplications.
Implementation Notes
- The implementation avoids explicitly building J and instead symmetrizes the linear system into normal equations for conjugate gradients.The text notes that explicit Jacobian construction is inefficient and rules out matrix-factorization methods such as LU decomposition.Source: Implementation Notes · p. 15
Coding right-hand-side
- The right-hand side is computed from the loss gradient using autograd vector-Jacobian products.The implementation applies autograd to obtain the multiplication by J transpose.Source: Coding right-hand-side · p. 16
Coding right matrix-vector multiplication by JJ⊤
- Two autograd calls compute multiplication by JJ⊤, which is supplied with the right-hand side to conjugate gradients to solve for w.After the solve, the implementation uses the result to compute the remaining gradient terms.Source: Coding right matrix-vector… · pp. 16–17
- The authors contrast this Jacobian-based implementation with JFB, describing JFB as more straightforward to code.Source: Coding right matrix-vector… · p. 17
D Comparison with Neumann RBP
The section compares JFB with fifth- and tenth-order Neumann-series gradient approximations on SVHN. It also describes a loop-based implementation that accumulates the Neumann sum.
D Comparison with Neumann RBP
- The SVHN comparison evaluates JFB against fifth- and tenth-order Neumann-series gradient approximations.The figure caption identifies the plotted subject as SVHN with different Neumann approximations of the inverse Jacobian, but the supplied passages do not state which method performs better.Source: D Comparison with Neumann RBP · p. 17, Figure 9 · p. 17
Neumann Gradient Implementation
- The Neumann-based gradient is computed by accumulating the series in a loop, following the approach described in Appendix C.Source: Neumann Gradient Implementation · p. 17
E Experimental Settings
Experiments use ResNet-based architectures, batch normalization in the latent component, and dataset-specific fixed-point stopping criteria. The Jacobian-based runs omit latent-space batch normalization because they failed to converge with it.
E Experimental Settings
- Latent-space batch normalization was used in the experiments; the networks never triggered the γ-contractivity warning, while Jacobian-based runs failed to converge with batch normalization and therefore omitted it.All networks used Adam and cross-entropy loss.Source: E Experimental Settings · p. 18
MNIST
- MNIST uses two leaky-ReLU convolutions with max pooling for QΘ and a two-layer ResNet for RΘ, stopping at tolerance 10^-4 or 50 iterations.The learning rate is 10^-4.Source: MNIST · p. 18
SVHN
- SVHN uses three one-layer ResNets for QΘ and a ResNet block with two batch-normalized convolutions for RΘ, stopping at tolerance 10^-4 or 200 iterations.The learning rate is 10^-4 with weight decay 2 × 10^-4.Source: SVHN · p. 18
CIFAR10
- CIFAR10 uses ResNets for QΘ and RΘ, with about 70% of weights in QΘ; JFB stops at tolerance 10^-1 or 50 iterations, versus 10^-4 or 500 for Jacobian-based runs.The tighter Jacobian-based tolerance was used to compute gradients accurately.Source: CIFAR10 · p. 18
F Toy Implicit Example
The toy example studies the fixed point of y = d + y^5 and establishes convergence to a unique solution on a specified interval. It also notes that the solution can be represented with a hypergeometric series.
F Toy Implicit Example
- Figure 10 illustrates the functions y and d + y^5 for d = 1/3.The example concerns solving y = d + y^5 for d in [−1/2, 1/2].Source: Figure 10 · p. 18, F Toy Implicit Example · p. 18
- For d ∈ [−1/2, 1/2], the iteration in Lemma F.1 converges to the unique fixed point in [−10^-1/4, 10^-1/4].The proof establishes contraction on a restricted interval, keeps the iterates inside it, and then proves convergence and uniqueness using a special case of Banach’s fixed point theorem.Source: F Toy Implicit Example · p. 19
- The proof shows the iterates remain in the restricted interval and form a Cauchy sequence, whose limit satisfies the fixed-point equation and is unique.The uniqueness argument uses the contractive property of T(·; d).Source: F Toy Implicit Example · p. 19
Explicit Solution
- The unique root of y = d + y^5 cannot be expressed using only arithmetic operations and roots; a hypergeometric series gives an explicit representation (Abel 1826).Further information on solving quintics with hypergeometric functions is cited to Birkeland (1927) and Ottem (2011).Source: Explicit Solution · p. 19