Relative Kinetic Utility for Reasoning-Aware Structural Pruning in Large Language Models

TL;DR

This paper introduces RKU, a novel framework for structural pruning in large language models that preserves reasoning capabilities at high sparsity levels by focusing on critical structural pathways.

Contribution

The paper proposes RKU, a continuous kinetic integral approach with Fisher normalization, to improve reasoning accuracy in pruned LLMs at high sparsity, surpassing existing methods.

Findings

RKU improves reasoning accuracy at 40% sparsity on GSM8K.

RKU outperforms baseline methods in high-sparsity regimes.

RKU better preserves reasoning representations under out-of-distribution tests.

Abstract

Chain-of-Thought (CoT) prompting symbolized a huge improvement of reasoning capabilities of Large Language Models (LLMs). However, scaling up test-time computation yields extensive CoT sequences, introducing severe inference latency and key-value (KV) cache memory bottlenecks. While structural pruning offers a fundamental, hardware-aware solution to alleviate static parameter burdens, existing magnitude-based methods may cut off the neurons of CoT: by over-indexing on discrete cross-entropy objectives, these heuristics fall into a \textit{magnitude trap}: they prioritize high-frequency, low-information syntactic tokens and trigger a disappointing reasoning collapse at high sparsities (e.g., 40\%). To overcome this topological phase transition, we propose \textsc{Relative Kinetic Utility} (RKU), a novel theoretical framework that elevates discrete pruning to a continuous kinetic integral over the depth manifold of the model based on Alternating Gradient Flow(AGF). By modifying it with Fisher trace normalization, RKU acts as a lightweight curvature-aware normalization to isolate \textit{kinetic spikes} -- the fundamental structural pathways responsible for high-curvature logical routing. Extensive experiments on Qwen-2.5-7B and LLaMA-3-8B improves performance in the high-sparsity regime around 40\%. RKU attains 13.34\% accuracy on GSM8K at 40\% sparsity, outperforming the strongest baseline, and appears to better preserve reasoning-relevant representations under out-of-distribution evaluation.