ARLP: Automatic multi-agent transformer reinforcement learning pruner for one-shot neural network pruning

Abstract

Overparameterized Neural Networks demonstrate state-of-the-art performance; however, the escalating demand for more compact and energy-efficient neural networks has arisen to facilitate the deployment of machine learning applications on devices with limited computational resources. A prevalent approach employs various pruning techniques. However, hyperparameters, such as the pruning ratio for each layer in pruning techniques, are typically set by human experts and often lack optimization. In this paper, we therefore propose a novel method named “Automatic Multi-Agent Transformer Reinforcement Learning Pruner” (ARLP). ARLP leverages a transformer-based multi-agent reinforcement learning controller to autonomously prune networks at initialization by extracting network meta-features. This autonomous process eliminates the need for human intervention in determining the optimal pruning ratio for each layer. The meta-features are derived from various zero-cost pruning-at-initialization proxies to perform One-shot Pruning. Extensive experimental results demonstrate that ARLP outperforms other state-of-the-art methods, establishing its efficacy in achieving superior performance.