Moonwalk: Inverse-Forward Differentiation
This addresses scalability issues for deep learning practitioners by enabling more efficient training of invertible networks with reduced memory usage.
The paper tackles the memory consumption problem of backpropagation in invertible networks by proposing a forward-mode gradient computation method, Moonwalk, which reduces memory footprint and achieves computation time comparable to backpropagation with linear time complexity and empirical speed-ups of several orders of magnitude.
Backpropagation, while effective for gradient computation, falls short in addressing memory consumption, limiting scalability. This work explores forward-mode gradient computation as an alternative in invertible networks, showing its potential to reduce the memory footprint without substantial drawbacks. We introduce a novel technique based on a vector-inverse-Jacobian product that accelerates the computation of forward gradients while retaining the advantages of memory reduction and preserving the fidelity of true gradients. Our method, Moonwalk, has a time complexity linear in the depth of the network, unlike the quadratic time complexity of naïve forward, and empirically reduces computation time by several orders of magnitude without allocating more memory. We further accelerate Moonwalk by combining it with reverse-mode differentiation to achieve time complexity comparable with backpropagation while maintaining a much smaller memory footprint. Finally, we showcase the robustness of our method across several architecture choices. Moonwalk is the first forward-based method to compute true gradients in invertible networks in computation time comparable to backpropagation and using significantly less memory.