JFB: Jacobian-Free Backpropagation for Implicit Networks
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley J. Osher, Wotao Yin
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers28
- A-ViT: Adaptive Tokens for Efficient Vision TransformerHongxu Yin, Arash Vahdat, José M. Álvarez, Arun Mallya et al.CVPR 2022 · 288 citations
- Online Training Through Time for Spiking Neural NetworksMingqing Xiao, Qingyan Meng, Zongpeng Zhang, Di He et al.NeurIPS 2022 · 121 citations
- Online Deep Equilibrium Learning for Regularization by DenoisingJiaming Liu, Xiaojian Xu, Weijie Gan, Shirin Shoushtari et al.NeurIPS 2022 · 35 citations
- Making Scalable Meta Learning PracticalSang Keun Choe, Sanket Vaibhav Mehta, Hwijeen Ahn, Willie Neiswanger et al.NeurIPS 2023 · 28 citations
- Deep Equilibrium Based Neural Operators for Steady-State PDEsTanya Marwah, Ashwini Pokle, J. Zico Kolter, Zachary C. Lipton et al.NeurIPS 2023 · 27 citations
Builds on4
- Multiscale Deep Equilibrium ModelsShaojie Bai, Vladlen Koltun, J. Zico KolterNeurIPS 2020 · 272 citations
- OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal TransportDerek Onken, Samy Wu Fung, Xingjian Li, Lars RuthottoAAAI 2021 · 210 citations
- Is Attention Better Than Matrix Decomposition?Zhengyang Geng, Meng-Hao Guo, Hongxu Chen, Xia Li et al.ICLR 2021 · 171 citations
- Almost Surely Stable Deep DynamicsNathan P. Lawrence, Philip D. Loewen, Michael G. Forbes, Johan U. Backström et al.NeurIPS 2020 · 28 citations
Related papers
- : Implicit Layers for Implicit RepresentationsZhichun Huang, Shaojie Bai, J. Zico KolterNeurIPS 2021 · 5 citations
- Stabilizing Equilibrium Models by Jacobian RegularizationShaojie Bai, Vladlen Koltun, J. Zico KolterICML 2021 · 80 citations
- SHINE: SHaring the INverse Estimate from the forward pass for bi-level optimization and implicit modelsZaccharie Ramzi, Florian Mannel, Shaojie Bai, Jean-Luc Starck et al.ICLR 2022 · 35 citations
- Neural Deep Equilibrium SolversShaojie Bai, Vladlen Koltun, J. Zico KolterICLR 2022 · 36 citations
- One-step differentiation of iterative algorithmsJérôme Bolte, Edouard Pauwels, Samuel VaiterNeurIPS 2023 · 36 citations
