Introduction

Drawing inspiration from the mechanism of microglia selectively pruning synapses in the brain, researchers developed a dynamic attention head pruning system. On the Phi-3-Mini model, this system achieves 20-30% attention head pruning with minimal accuracy loss while improving inference latency by 10-15%, providing a new optimization approach to the inference cost problem of large language models.

Biology-Inspired Dynamic Pruning Approach

During human brain development, microglia selectively eliminate low-activity synapses to optimize information transmission efficiency. This mechanism inspired researchers to propose a dynamic pruning paradigm: unlike static weight pruning after training, the model adaptively decides which attention heads to skip during inference based on input complexity—aggressive pruning for simple queries, and more resources reserved for complex reasoning.

System Architecture and Training Strategy

Three-Layer Collaborative Design

1. Activation Monitoring Layer: Captures hidden states and attention weights via PyTorch hooks to provide decision-making basis.
2. MicrogliaAgent: A lightweight MLP that receives statistical features (L2 norm of hidden states, entropy of attention distribution) and outputs 0-1 soft mask values (facilitating gradient backpropagation).
3. Masked Attention Layer: Applies masks to suppress attention head outputs, achieving computational savings at the hardware level.

Curriculum Learning Strategy

At the initial stage of training, set a low pruning pressure parameter alpha (0.01) to retain almost all heads; as training progresses, increase alpha to 0.3, forcing the Agent to improve pruning ratio while maintaining accuracy to avoid model collapse.

Experimental Validation and Toolchain

Phi-3-Mini Experimental Results

• 20-30% of attention heads can be safely pruned with only minimal drop in GSM8K accuracy;
• Actual inference latency improved by 10-15% (measured wall-clock time via CUDA events);
• Structured pruning can be mapped to hardware acceleration.

Toolchain and Multi-Model Support

Three Jupyter Notebooks are provided: Quick Demo (20-30 minutes), Strict Experiment (2-3 hours), and Complete Pipeline (3-4 hours); supports Qwen2.5-3B-Instruct, demonstrating cross-model generality.

Limitations and Future Exploration

Current Limitations

• The Agent network introduces a small additional overhead (less than 5% parameter increase);
• Validated only on encoder-decoder structured instruction-tuned models; pure decoder base models and multimodal scenarios remain to be explored.

Future Directions

• Explore 'hard pruning' (binarizing soft masks) to gain greater hardware acceleration;
• Extend to more model types and scenarios.

Conclusion

Microglia Pruning integrates pruning into the inference process, enabling input-adaptive allocation of computational resources. It is an innovative application of the cross-disciplinary idea of 'biological inspiration + machine learning'. The project provides a complete pip package and Colab notebooks; developers can reproduce core results with only consumer-grade GPUs, offering a feasible path to solving large model deployment challenges.