Efficient Mathematical Reasoning Models via Dynamic Pruning and Knowledge Distillation

TL;DR

This paper introduces a method combining dynamic attention head pruning with knowledge distillation to create lightweight models that efficiently perform mathematical reasoning with minimal accuracy loss.

Contribution

It presents a novel real-time pruning technique based on importance metrics and uses knowledge distillation to retain reasoning capabilities in smaller models.

Findings

Parameters reduced by 18.7% at 30% pruning ratio

Inference speed improved by 27.5%

Accuracy decreased by only 0.7%

Abstract

With the rapid development of deep learning, large language models have shown strong capabilities in complex reasoning tasks such as mathematical equation solving. However, their substantial computational and storage costs hinder practical deployment. This paper proposes a lightweight optimization method that integrates dynamic attention head pruning with knowledge distillation. The approach dynamically evaluates the importance of each attention head in the multi-head attention mechanism using a combination of weight norms and entropy, and prunes redundant heads in real time to reduce computational overhead. To mitigate performance degradation, knowledge distillation transfers information from the original model to the pruned student, enabling the smaller model to preserve reasoning ability. Experiments conducted on both Math23k and ASDiv-A verify the effectiveness of the proposed method. For example, on Math23k with a 30% pruning ratio, parameters are reduced by 18.7%, inference speed is improved by 27.5%, FLOPs are reduced by 19.3%, and accuracy drops only 0.7% (from 84.4% to 83.7%). These results demonstrate that the method achieves substantial efficiency gains while maintaining strong reasoning performance, providing a practical solution for efficient deployment of large language models in mathematical reasoning tasks.