ECAI: Efficient Convolution Activation Inversion for Constant-Memory Convolutional Neural Networks Training

We propose a novel approach that achieves constant activation memory usage during the training of convolutional neural networks (CNNs), addressing a key memory bottleneck in the backward pass. By reconstructing activations required for gradient matrix calculation through the proposed efficient convolution activation inversion (ECAI) rather than storing them in memory during forward pass, it becomes possible to maintain constant activation memory usage across convolution layers. We formulate the activation inversion problem as a set of $n$ systems of linear equations derived from forward convolution operations, and solve them with an accelerated method that achieves $\mathcal{O}(n^2)$ complexity. The proposed approach enables memory-constrained mobile, edge, and embedded devices to perform CNN training without a growth of activation memory over the model capacity while also enhancing training memory efficiency for large-sized images on commercial GPUs. The experimental results demonstrate that the proposed approach maintains constant activation memory by reusing a fixed memory space, improving memory efficiency without degradation in model accuracy. The memory savings achieved by the proposed method increase when using more convolution layers, potentially achieving near-zero activation (e.g., $30\times$ or more activation memory reduction in specific setups). The code implementation is available at an anonymous GitHub.