Training DNNs in O(1) memory with MEM-DFA using Random Matrices

TL;DR

This paper introduces MEM-DFA, a training method for deep neural networks that reduces memory usage to a constant, enabling training of very deep models with minimal memory overhead by leveraging random matrices and biologically plausible feedback mechanisms.

Contribution

The paper proposes MEM-DFA, a novel memory-efficient training algorithm that maintains constant memory usage regardless of network depth, based on feedback alignment with random matrices.

Findings

MEM-DFA significantly reduces memory consumption compared to BP, FA, and DFA.

Experimental results on MNIST and CIFAR-10 validate theoretical memory savings.

MEM-DFA incurs only a small constant increase in computational cost.

Abstract

This work presents a method for reducing memory consumption to a constant complexity when training deep neural networks. The algorithm is based on the more biologically plausible alternatives of the backpropagation (BP): direct feedback alignment (DFA) and feedback alignment (FA), which use random matrices to propagate error. The proposed method, memory-efficient direct feedback alignment (MEM-DFA), uses higher independence of layers in DFA and allows avoiding storing at once all activation vectors, unlike standard BP, FA, and DFA. Thus, our algorithm's memory usage is constant regardless of the number of layers in a neural network. The method increases the computational cost only by a constant factor of one extra forward pass. The MEM-DFA, BP, FA, and DFA were evaluated along with their memory profiles on MNIST and CIFAR-10 datasets on various neural network models. Our experiments agree with our theoretical results and show a significant decrease in the memory cost of MEM-DFA compared to the other algorithms.