Application-Specific Component-Aware Structured Pruning of Deep Neural Networks in Control via Soft Coefficient Optimization

TL;DR

This paper presents a novel structured pruning framework for deep neural networks used in control applications, optimizing importance metrics with application-specific constraints and balancing model compression with task performance.

Contribution

It introduces a new importance metric calculation framework and two coefficient optimization approaches tailored for application-specific neural network pruning.

Findings

Effective model size reduction while maintaining control performance.

Gradient descent-based coefficient optimization improves compression-performance trade-off.

Validated on MNIST autoencoder and TDMPC agent with positive results.

Abstract

Deep neural networks (DNNs) offer significant flexibility and robust performance. This makes them ideal for building not only system models but also advanced neural network controllers (NNCs). However, their high complexity and computational needs often limit their use. Various model compression strategies have been developed over the past few decades to address these issues. These strategies are effective for general DNNs but do not directly apply to NNCs. NNCs need both size reduction and the retention of key application-specific performance features. In structured pruning, which removes groups of related elements, standard importance metrics often fail to protect these critical characteristics. In this paper, we introduce a novel framework for calculating importance metrics in pruning groups. This framework not only shrinks the model size but also considers various application-specific constraints. To find the best pruning coefficient for each group, we evaluate two approaches. The first approach involves simple exploration through grid search. The second utilizes gradient descent optimization, aiming to balance compression and task performance. We test our method in two use cases: one on an MNIST autoencoder and the other on a Temporal Difference Model Predictive Control (TDMPC) agent. Results show that the method effectively maintains application-relevant performance while achieving a significant reduction in model size.