SubspacePath Pruner: Inference-time Pruning via Probe-based Representation–Parameter Coupling
Abstract
Large-scale dedicated application of LLMs in diverse scenarios increasingly demands specialized model inference behavior under strict constraints of accuracy, latency, and memory. However, the heterogeneous and long-tailed nature of real-world specialized scenarios makes it difficult to obtain training data and optimize models. We study a practical inference-time specialization setting: given an LLM base, we compile a reusable, budget-bounded pathway/subnetwork within a specific scenario. Our approach is motivated by an empirical coupling phenomenon: input scenario sets aligned with similar representation subspaces (e.g., domain) in embedding space tend to activate a consistent and sparse set of internal reasoning pathways in model parameter space. To build the bridge between them, we propose probe-based SubspacePath Pruner with two core components: (1) Domain-Basis Synthesis (DBS) constructs a quasi-orthogonal basis of domain axes in embedding space, serving as a stable coordinate system. (2) Probe-based Scenario Pruning (PSP) uses efficient layer-wise linear probes to estimate axis alignment and compute budgeted head-wise pathways for a specific scenario. Experiments on LLaMA-2-13B show 29.3 average Recall on cross-domain tests (vs. 24.7 dense) and 21.6 on cross-dataset tests (vs. 25.5 dense) with 1.27x speedup at ~30% pruning ratio.
Lay Summary
Large language models are powerful, but using them in real-world applications is often expensive because they require large amounts of memory and computation. This becomes especially challenging when a model is repeatedly used in a specific setting, such as answering questions in a particular domain, where we would like the model to be both efficient and reliable without collecting new training data or retraining the model. This paper introduces SubspacePath Pruner, a method that makes a large language model more efficient at inference time. The key idea is that different application scenarios tend to rely on different internal parts of the model. Instead of using the full model for every input, our method first identifies what kind of scenario the model is facing, then keeps the internal components that are most relevant to that scenario and turns off less useful ones. This produces a smaller, reusable version of the model for that scenario. Our approach does not require task labels, model retraining, or changing the original model weights. Experiments across several language models and datasets show that this scenario-aware pruning can reduce computation and memory use while preserving, and sometimes improving, performance, especially when the test data comes from different or mixed domains. These results suggest that large language models can be specialized more practically for real-world deployment by selecting the right internal pathways at inference time.