Effective pruning of task-trained recurrent neural networks using noisy fluctuations and connection rescaling

This paper evaluates the pruning rule noise-prune for recurrent neural networks, which uses noisy fluctuations in neural activity to determine the importance of connections. The authors apply noise-prune to recurrent neural networks trained on eight cognitive tasks from the Mod-Cog battery, comparing it against magnitude-based pruning, random pruning, and a non-local Optimal Brain Surgeon (OBS) variant.

The key findings are:

1. Noise-prune preserves task performance: "We show that noise-prune preserves task-performance in task-trained recurrent neural networks, greatly outperforming a strategy that only uses the magnitude of connections and performing on par with or exceeding a non-local strategy that uses second-order information. Specifically, At sparsity levels between 50–60% noise-prune is comparable to the more computationally-intensive recurrent-OBS method... At 80% sparsity, L-NP and S-NP retained 0.378 ± 0.067 and 0.382 ± 0.063 of baseline task accuracy, respectively, compared with 0.137 ± 0.025 for recurrent-OBS and 0.087 ± 0.013 for magnitude pruning."

2. Probabilistic sampling and rescaling is essential: "Rather than deterministically removing connections that fall below a certain threshold importance, noise-prune samples connections to preserve based on their importance and strengthens retained connections to preserve average synaptic strength. We show that this sampling and rescaling is essential to good performance. Removing this step dramatically decreases performance for both variants, to levels well below noise-prune with sampling and rescaling."

3. Optimal rescaling is lower than theory predicts: We show that the optimal empirical degree of rescaling is lower than that predicted by the original theoretical argument. Capping the amplification factors improves performance: "At the capping percentile corresponding to peak task retention, L-NP and S-NP on average retain 0.691 ± 0.038 and 0.730 ± 0.032 of task accuracy, respectively... Both differences are significant (L-NP: P = 6.53 × 10−5; S-NP: P = 8.34 × 10−6)."

The paper concludes that "effective task-preserving pruning of recurrent neural networks under biological constraints can be carried out by a combined strategy of using noisy fluctuations to evaluate the importance of network connections along with using these importance factors to probabilistically sample connections to prune and strengthening the unpruned weights."

Improvements to AI Systems:

1. Biologically-plausible, local pruning for RNNs/LSTMs: Implement noise-prune as a drop-in pruning method for recurrent architectures (e.g., LSTMs, GRUs, transformers with recurrence) that require on-device or energy-efficient deployment. Unlike magnitude pruning, which fails at high sparsity, noise-prune uses intrinsic activation noise (e.g., from dropout or stochastic depth) to compute connection importance locally, enabling 80% sparsity while retaining 38% of task accuracy vs. 9% for magnitude pruning.

2. Probabilistic weight retention with rescaling: Replace deterministic threshold-based pruning with a stochastic sampling step where connections are kept with probability proportional to their noise-derived importance, and surviving weights are amplified by a learned or fixed factor. This prevents catastrophic information loss and preserves average synaptic strength, making the pruned network more robust to distribution shift and adversarial perturbations.

3. Adaptive rescaling cap for stability: Introduce a tunable cap on the amplification factor (e.g., percentile-based capping) to prevent over-amplification of a few surviving weights. This improves peak task retention by 0.69–0.73 of baseline accuracy (vs. 0.38 without capping), reducing variance and improving generalization in low-resource settings.

4. Local, communication-efficient pruning for federated or neuromorphic systems: Since noise-prune requires only local noisy fluctuations (no global second-order information like OBS), it can be applied in distributed AI systems where gradients or Hessians are unavailable or expensive to share. This enables on-device pruning during inference or continual learning without central coordination.

5. Dynamic sparsity scheduling: Use noise-prune’s importance scores to guide gradual sparsification during training (e.g., starting at 50% sparsity and increasing to 80%), allowing the network to re-route information through remaining connections. This yields a smoother accuracy-sparsity trade-off curve, useful for adaptive compute in edge AI.

What the improved AI system can do:

• Achieve high sparsity (80%+) in recurrent models while retaining significantly more task performance than magnitude or even second-order methods, enabling deployment on memory-constrained hardware (e.g., microcontrollers, mobile NPUs).

• Maintain performance under noisy or stochastic conditions (e.g., dropout, analog hardware noise) by leveraging that same noise for pruning, making the system naturally robust to hardware-induced variability.

• Prune models locally and continuously without global gradient information, allowing real-time model compression in federated learning or on-device continual learning scenarios.

• Provide a tunable trade-off between compression ratio and accuracy via the rescaling cap, letting system designers choose a sweet spot for their specific latency/energy budget.