# Moonwalk: Inverse-Forward Differentiation

arXiv: https://arxiv.org/abs/2402.14212 · HTML: https://mlsanity.com/paper/2402.14212/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 60 references that are in the map.
- First posted on arXiv 22 Feb 2024.

## Builds on (most cited first, 22 of 22)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Glow: Generative Flow with Invertible 1x1 Convolutions](https://mlsanity.com/paper/1807.03039/) (2018)
- [NICE: Non-linear Independent Components Estimation](https://mlsanity.com/paper/1410.8516/) (2014)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [Invertible Residual Networks](https://mlsanity.com/paper/1811.00995/) (2018)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [The Reversible Residual Network: Backpropagation Without Storing Activations](https://mlsanity.com/paper/1707.04585/) (2017)
- [Decoupled Neural Interfaces using Synthetic Gradients](https://mlsanity.com/paper/1608.05343/) (2016)
- [i-RevNet: Deep Invertible Networks](https://mlsanity.com/paper/1802.07088/) (2018)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [In-Place Activated BatchNorm for Memory-Optimized Training of DNNs](https://mlsanity.com/paper/1712.02616/) (2017)
- [Memory-Efficient Backpropagation Through Time](https://mlsanity.com/paper/1606.03401/) (2016)
- [Scaling Forward Gradient With Local Losses](https://mlsanity.com/paper/2210.03310/) (2022)
- [Reversible Vision Transformers](https://mlsanity.com/paper/2302.04869/) (2023)
- [Reversible Recurrent Neural Networks](https://mlsanity.com/paper/1810.10999/) (2018)
- [Backprop with Approximate Activations for Memory-efficient Network Training](https://mlsanity.com/paper/1901.07988/) (2019)
- [Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction](https://mlsanity.com/paper/2202.00441/) (2022)
- [Rockmate: an Efficient, Fast, Automatic and Generic Tool for Re-materialization in PyTorch](https://mlsanity.com/paper/2307.01236/) (2023)

## Built on by (most cited first, 1 of 1)

- [Projected Forward Gradient-Guided Frank-Wolfe Algorithm via Variance Reduction](https://mlsanity.com/paper/2403.12511/) (2024)

## Abstract (arXiv)

Backpropagation's main limitation is its need to store intermediate activations (residuals) during the forward pass, which restricts the depth of trainable networks. This raises a fundamental question: can we avoid storing these activations? We address this by revisiting the structure of gradient computation. Backpropagation computes gradients through a sequence of vector-Jacobian products, an operation that is generally irreversible. The lost information lies in the cokernel of each layer's Jacobian. We define submersive networks -- networks whose layer Jacobians have trivial cokernels -- in which gradients can be reconstructed exactly in a forward sweep without storing activations. For non-submersive layers, we introduce fragmental gradient checkpointing, which records only the minimal subset of residuals necessary to restore the cotangents erased by the Jacobian. Central to our approach is a novel operator, the vector-inverse-Jacobian product (vijp), which inverts gradient flow outside the cokernel. Our mixed-mode algorithm first computes input gradients with a memory-efficient reverse pass, then reconstructs parameter gradients in a forward sweep using the vijp, eliminating the need to store activations. We implement this method in Moonwalk and show that it matches backpropagation's runtime while training networks more than twice as deep under the same memory budget.