Component-Aware Pruning Framework for Neural Network Controllers via Gradient-Based Importance Estimation

TL;DR

The paper tackles the challenge of compressing multi-component neural network controllers by introducing a component-aware pruning framework that leverages three gradient-based importance scores: $I_g^{grad}$, $I_g^{Fisher}$, and $I_g^{Bayes}$. These metrics are computed online during training and integrated through a per-group regularization schedule with temporal smoothing, enabling data-driven pruning decisions. Experiments on a MNIST autoencoder and a TD-MPC controller reveal that group importance is highly dynamic and architecture-dependent, with coupling groups not universally critical and encoder/dynamics groups often taking precedence depending on capacity. This approach yields interpretable pruning guidance and preserves control performance under aggressive compression, offering a practical pathway for deploying compressed NNCs in resource-constrained, embedded systems.

Abstract

The transition from monolithic to multi-component neural architectures in advanced neural network controllers poses substantial challenges due to the high computational complexity of the latter. Conventional model compression techniques for complexity reduction, such as structured pruning based on norm-based metrics to estimate the relative importance of distinct parameter groups, often fail to capture functional significance. This paper introduces a component-aware pruning framework that utilizes gradient information to compute three distinct importance metrics during training: Gradient Accumulation, Fisher Information, and Bayesian Uncertainty. Experimental results with an autoencoder and a TD-MPC agent demonstrate that the proposed framework reveals critical structural dependencies and dynamic shifts in importance that static heuristics often miss, supporting more informed compression decisions.

Figures (4)

• Figure 1: Group importance scores compared across various Autoencoder architectures with latent dimension sizes ($d \in \{8, 64, 256, 512\}$) after 110 training epochs. Columns indicate estimation methods (Gradient, Fisher, Bayesian), while rows denote different architectures.
• Figure 2: Extended training dynamics over 200 epochs for the 512-dimensional model. The observed crossovers in importance rankings, especially in the Gradient and Fisher metrics, challenge the static importance assumption stated in Hypothesis 3.
• Figure 3: Evolution of group importance for TD-MPC using Gradient Accumulation (left) and Fisher Information (right). Encoder groups (green/orange) consistently exhibit greater importance than coupling layers, which contradicts Hypothesis 1.
• Figure 4: Bayesian importance estimates for TD-MPC. The pronounced decay of the reward model's importance (purple line) indicates that group criticality is dynamic and changes throughout training.