A Layer-wise Training and Pruning Method for Memory Efficient On-chip Learning Hardware

Dongwoo Lew, Jongsun Park

Abstract: In this paper, we present a layer-wise on-chip learning and pruning algorithm to enable the training of deep neural networks (DNNs) that are too large to be directly trained using the limited memory resources of edge devices. To perform on-chip learning of large DNN in a limited memory space, the proposed algorithm performs pruning in the middle of the training by initializing only the first layer of the network and training that layer using direct random target projection (DRTP) algorithm. After training of the first layer is completed, a simple magnitude based pruning is applied to create empty memory space for the second layer and this process is repeated until the last layer. The proposed algorithm allows training larger DNN on a smaller physical memory compared to conventional algorithms. The proposed algorithm achieves 95.06% classification accuracy with a network of <tex>$28 \times 28-500-500-10$</tex> structure for the MNIST dataset while using 39.5% smaller physical weight memory compared to conventional learning methods.