DART-ing Through the Drift: Dynamic Tracing of Knowledge Neurons for Adaptive Inference-Time Pruning

TL;DR

DART provides a training-free, runtime approach to adaptively prune FFN neurons in LLMs by tracking context shifts via attention dynamics and triggering re-pruning to combat knowledge drift. The method combines layer-aware sparsity budgeting with an online detection mechanism that realigns the active neuron subspace as semantic context evolves, delivering significant gains over static and some dynamic baselines while maintaining near-dense performance under modest memory and FLOP overhead. Empirically, it achieves up to 14.5% accuracy gains at 70% FFN sparsity on LLAMA-3.1-8B and substantially higher ROUGE-L scores on summarization, demonstrating robust generalization across zero-shot and domain-specific tasks and long-context generation. The work has practical impact for deploying long-context LLMs in resource-constrained environments, offering a scalable, runtime solution with publicly available code.

Abstract

Large Language Models (LLMs) exhibit substantial parameter redundancy, particularly in Feed-Forward Networks (FFNs). Existing pruning methods suffer from two primary limitations. First, reliance on dataset-specific calibration introduces significant data dependency and computational overhead. Second, being predominantly static, they fail to account for the evolving subset of knowledge neurons in LLMs during autoregressive generation as the context evolves. To address this, we introduce DART, i.e., Dynamic Attention-Guided Runtime Tracing), a lightweight, training-free method that performs on-the-fly context-based pruning. DART monitors shifts in attention score distributions to infer context changes, dynamically updating neuron-level masks to retain salient parameters. Across ten benchmarks, DART outperforms prior dynamic baseline, achieving accuracy gains of up to 14.5% on LLAMA-3.1-8B at 70% FFN sparsity. Furthermore, DART achieves up to 3x better ROUGE-L scores with respect to static-masked pruning on summarization tasks, with its performance comparable to the original dense models. We conclusively demonstrate that the proposed framework effectively adapts to diverse semantic contexts, preserves model capabilities across both general and domain-specific tasks while running at less than 10MBs of memory for LLAMA-3.1-8B(16GBs) with 0.1% FLOPs overhead. The code is available at https://github.com/seeder-research/DART.

Figures (11)

• Figure 1: Variation in attention output distributions across different input contexts. This distributional shift necessitates the activation of distinct neuron subsets.
• Figure 2: Layerwise sensitivity of LLaMA-3.2-3B on the MMLU marketing subject.
• Figure 3: Cosine similarity of embeddings pre- and post-FFN layers in LLaMA-3.2-3B for the MMLU marketing subject.
• Figure 4: Comparison of unmasked and masked model outputs for a multi-expert dialogue. The pruning mask generated for a specific topic biases the output of another topic.
• Figure 5: Cosine similarity of the attention vector over time for unpruned and pruned FFN at the last layer.