Partial Attention Probing (PAP)

• PAP is a mechanism that selectively skips less informative tokens in multi-head attention layers, reducing redundant computation in Transformer models.
• It scores token relevance via partial forward attention, using these scores as gating signals to bypass costly computations in both attention and feedforward layers.
• When integrated with the SPTS framework and low-rank transformation probing, PAP achieves up to 2.46× speedups on 32K-token contexts with only marginal accuracy degradation.

Partial Attention Probing (PAP) is a mechanism designed to enable efficient, principled token skipping in the multi-head attention (MHA) layers of large Transformer-based LLMs during long-context inference. PAP operates by performing partial forward attention computation to score token relevance in the current context, thus determining which tokens are sufficiently informative to participate in further expensive computations. PAP forms a central component of the Self-Predictive Token Skipping (SPTS) framework, which achieves substantial acceleration in inference with negligible loss in accuracy through a combination of targeted skipping at both attention and feedforward layers (Wu et al., 19 Jan 2026).

1. Conceptual Overview

Partial Attention Probing is motivated by the need to dynamically identify and propagate only the most informative tokens in the MHA blocks at inference time, thereby reducing redundant computation and achieving low-latency inference for long sequences. Within SPTS, PAP serves as the attention-layer analogue to Low-rank Transformation Probing (LTP) for the feedforward network (FFN), orchestrating selective retention of tokens based on their attention-based contextual importance. The primary insight underlying PAP is that a substantial fraction of tokens in extended sequences contribute minimally to the output of subsequent layers and can safely be bypassed with little impact on model fidelity.

2. Mechanism and Scoring Principle

In the Transformer architecture, the multi-head attention computes contextually weighted mixtures for each token across all others. PAP introduces a partial computation approach, where only a fraction of the full attention map is computed to quickly estimate how "attended" or relevant each token is within the local context. The resulting per-token relevance scores, denoted $S_n^{\mathrm{MHA}}$, capture the contextual informativeness of each token in relation to the rest of the sequence.

Within the SPTS pipeline, these $S_n^{\mathrm{MHA}}$ scores are then directly used as gating signals to control whether subsequent layers (such as full FFN computation) will be applied to the token or bypassed via the residual connection. This approach enables near-instantaneous identification of high-impact tokens without requiring computation over all possible token interactions in the attention matrix.

3. Integration with Low-rank Transformation Probing and Staged Pruning

PAP does not operate in isolation: it is systematically combined with LTP, which deploys a closed-form, low-rank approximation of the FFN update magnitude per token. Each token's ultimate selection for full computation in the FFN is determined by the product score

$$S_n^{\mathrm{FFN}} = C_n^{\mathrm{FFN}} \cdot S_n^{\mathrm{MHA}}$$

where $C_n^{\mathrm{FFN}}$ is the $\ell_2$-norm of the LTP proxy network’s output for that token, representing a data-driven estimate of the magnitude of its update if the full FFN were to be applied. PAP thus acts as a front-end proxy, estimating the potential contextual impact of a token before more expensive downstream computation is performed.

Multi-Stage Delayed Pruning (MSDP) further leverages these per-token scores, reallocating the skipping budget across multiple layers to optimize the speed–accuracy trade-off dynamically throughout model depth.

4. Quantitative Performance and Empirical Results

Experimental findings demonstrate that PAP achieves substantial reductions in inference time with minimal degradation to prediction accuracy. Specifically, when applied in conjunction with LTP and MSDP, PAP enables speedups of up to 2.46× for prefilling and 2.29× for end-to-end text generation on 32K-token contexts in LLaMA-8B-scale models, all while maintaining state-of-the-art performance on comprehensive evaluation suites such as LongBench (Wu et al., 19 Jan 2026). An ablation study reports an average score of 45.29% on LongBench when using PAP ($S^{\mathrm{MHA}}$) alone for FFN skipping (with 50% token skipping in each FFN). The addition of the LTP proxy to form the composite $S^{\mathrm{FFN}}$ marginally increases average accuracy by 0.6% to 45.89%.

The following table summarizes the cumulative impact of different SPTS components on inference speed (TTFT: time to first token) for 32K-token inference on LLaMA-8B:

Configuration	Relative TTFT	Improvement over Baseline
Full model	1.00×	—
+ PAP only	1.44×	+44%
+ PAP + LTP	2.15×	+49% over PAP only
+ PAP + LTP + MSDP	2.46×	+14% over prior stage

5. Construction, Calibration, and Training-Free Nature

PAP is inherently training-free: no additional gradient-based fine-tuning or reoptimization is required. The per-token attention scores are obtained by partially evaluating the model’s own pretrained attention weights. All calibration occurs offline using a small corpus, ensuring that runtime overhead is negligible and no model weights are modified during inference (Wu et al., 19 Jan 2026).

A plausible implication is that since PAP relies on the model's native pre-trained parameters and representative calibration data, its effectiveness may depend on the alignment between calibration and deployment data distributions. Nonetheless, the training-free property ensures broad applicability and low engineering cost in practical deployment.

6. Practical Considerations and Limitations

While PAP is highly effective at accelerating inference, its performance is influenced by several factors:

• Proxy Fidelity vs. Computational Savings: The partial attention computation must balance fidelity of token relevance estimates against the cost savings from incomplete evaluation. Excessively shallow probing may misclassify important tokens, while over-computation reduces acceleration benefits.
• Task Sensitivity: The distribution of attention scores may vary substantially between tasks. A plausible concern is that in settings where attention is a weak predictor of downstream importance (e.g., certain generation or reasoning tasks), PAP may inadvertently prune tokens that are essential for model performance.
• Interaction with Downstream Pruning: As PAP conditions LTP’s token selection via $S_n^{\mathrm{MHA}}$, propagation of suboptimal scores can compound through the model, potentially amplifying errors if initial attention proxy signals are not robust.

Despite these considerations, empirical evidence from the SPTS framework indicates that PAP enables principled, context-sensitive acceleration of long-context LLM inference, unlocking up to 2–2.5× speedups on large-scale models with negligible loss in output quality (Wu et al., 19 Jan 2026).

7. Relation to Token-Skipping Paradigms in Long-Context Inference

Partial Attention Probing represents a class of token-oriented acceleration strategies that exploit the intrinsic sparsity of contextual importance in long-input scenarios. Unlike static pruning or fixed sparsity patterns, PAP dynamically adapts to the content and structure of each input sequence by utilizing partial, context-aware probes of the multi-head attention landscape. This approach directly addresses prior limitations observed in outdated or heuristic proxy signals, and forms a template for future efficient inference designs in large-scale transformer architectures. The decomposition of attention- and feedforward-layer skipping into specialized, training-free proxy mechanisms—of which PAP is the earliest stage—marks a significant advance toward attaining optimal speed–accuracy trade-offs in practical LLM deployment (Wu et al., 19 Jan 2026).