Learning a Consensus Sub-Network with Polarization Regularization and One Pass Training

TL;DR

This paper introduces a one-pass training and pruning method that efficiently finds sparse sub-networks with minimal energy use and negligible accuracy loss, suitable for green AI applications.

Contribution

It proposes a novel, one-pass pruning strategy using a binary gating module and polarization loss, enabling simultaneous training and pruning for energy-efficient neural networks.

Findings

Removes 50% of connections with less than 1% accuracy loss.

Achieves comparable performance to full networks with reduced computational cost.

Demonstrates lower accuracy drop than existing pruning methods.

Abstract

The subject of green AI has been gaining attention within the deep learning community given the recent trend of ever larger and more complex neural network models. Existing solutions for reducing the computational load of training at inference time usually involve pruning the network parameters. Pruning schemes often create extra overhead either by iterative training and fine-tuning for static pruning or repeated computation of a dynamic pruning graph. We propose a new parameter pruning strategy for learning a lighter-weight sub-network that minimizes the energy cost while maintaining comparable performance to the fully parameterised network on given downstream tasks. Our proposed pruning scheme is green-oriented, as it only requires a one-off training to discover the optimal static sub-networks by dynamic pruning methods. The pruning scheme consists of a binary gating module and a polarizing loss function to uncover sub-networks with user-defined sparsity. Our method enables pruning and training simultaneously, which saves energy in both the training and inference phases and avoids extra computational overhead from gating modules at inference time. Our results on CIFAR-10, CIFAR-100, and Tiny Imagenet suggest that our scheme can remove 50% of connections in deep networks with <1% reduction in classification accuracy. Compared to other related pruning methods, our method demonstrates a lower drop in accuracy for equivalent reductions in computational cost.