Low-Rank Transformation Probing (LTP)

• Low-Rank Transformation Probing (LTP) is a method that constructs a low-rank proxy network to estimate per-token FFN updates in Transformer models.
• It employs structured pruning and SVD to efficiently determine which tokens need full processing, thereby reducing unnecessary computations.
• Integrated within the Self-Predictive Token Skipping framework, LTP achieves significant speedups with minimal accuracy loss during long-context inference.

Low-rank Transformation Probing (LTP) is a training-free strategy for efficient inference in deep Transformer models, specifically enabling the selective skipping of Feed-Forward Network (FFN) computations for tokens that are predicted to change little under the block’s transformation. Developed as a core component of the Self-Predictive Token Skipping (SPTS) framework, LTP constructs a low-rank proxy network that estimates the FFN-induced change for each token and uses this estimate, combined with an attention-derived importance score, to inform token selection in long-context LLM inference. The method exploits structured pruning and singular value decomposition (SVD) of FFN weights, yielding substantial reductions in inference cost while maintaining high fidelity to original model outputs (Wu et al., 19 Jan 2026).

1. Theoretical Motivation

In Transformer architecture, each FFN block transforms input hidden states $x_n \in \mathbb{R}^D$ individually:

$$F_{\mathrm{FFN}}(x_n) = W_2 \cdot \sigma(W_1 x_n + b_1) + b_2$$

where $W_1 \in \mathbb{R}^{D \times D_{\mathrm{ff}}}$, $W_2 \in \mathbb{R}^{D_{\mathrm{ff}} \times D}$, and $\sigma$ is a nonlinearity. Empirically, deep Transformer layers exhibit sparse representational changes: most tokens incur negligible updates, while a small subset undergo high-magnitude changes and carry new semantic information. Efficient inference thus hinges on identifying these “active” tokens to selectively apply full FFN computation, bypassing needless evaluation for the rest. LTP enables on-the-fly estimation of the FFN update magnitude $\|\!F_{\mathrm{FFN}}(x_n)\!\|_2$ without incurring the full computational cost, allowing the system to rank and skip tokens accordingly.

2. Proxy Network Architecture and Low-Rank Factorization

The proxy construction begins with a reformulation of the FFN:

$$F_{\mathrm{FFN}}(X) = \bigl[\sigma(X W_{\mathrm{gate}} + b_{\mathrm{gate}}) \odot (X W_{\mathrm{up}} + b_{\mathrm{up}})\bigr] W_{\mathrm{down}} + b_{\mathrm{down}}$$

where $X \in \mathbb{R}^{N \times D}$, and FFN weights are expressed as $$W_{\mathrm{gate}}, W_{\mathrm{up}}, W_{\mathrm{down}} \in \mathbb{R}^{D \times D_{\mathrm{ff}}}$$. LTP executes proxy construction via two mechanisms:

1. Channel Selection (Structured Pruning): Activations for each hidden state are probed on a calibration set $G$. For channel $j$,

$$z(x) = \bigl|\sigma(x W_{\mathrm{gate}} + b_{\mathrm{gate}}) \odot (x W_{\mathrm{up}} + b_{\mathrm{up}})\bigr| \in \mathbb{R}^{D_{\mathrm{ff}}}$$

$$I_j = \mathrm{Mean}\Bigl(\mathrm{TopK}\{z(x)[j]\ |\ x \in G\},\ \rho|G|\Bigr)$$

The top $D_{\mathrm{low}}$ channels $C$ are retained; weights and biases are truncated accordingly.

2. Low-Rank SVD Factorization: For each truncated FFN weight matrix $W' \in \mathbb{R}^{D \times D_{\mathrm{low}}}$, a rank-$r$ SVD is computed:

$W' \approx U V$

with $U \in \mathbb{R}^{D \times r}$, $V \in \mathbb{R}^{r \times D_{\mathrm{low}}}$ and $r \ll \min(D, D_{\mathrm{low}})$. This yields decompositions for gate, up, and down branches.

The resulting low-rank proxy computes:

$$\mathrm{Gate}(X) = \sigma(X U_{\mathrm{gate}} V_{\mathrm{gate}} + b'_{\mathrm{gate}})$$

$$\mathrm{Up}(X) = X U_{\mathrm{up}} V_{\mathrm{up}} + b'_{\mathrm{up}}$$

$$f(X) = [\mathrm{Gate}(X) \odot \mathrm{Up}(X)] (V_{\mathrm{down}}^T U_{\mathrm{down}}^T) + b'_{\mathrm{down}}$$

Cost is $\mathcal{O}(N D r + N r D_{\mathrm{low}})$, with $r$ and $D_{\mathrm{low}}$ set much smaller than $D_{\mathrm{ff}}$.

3. Inference-Time Algorithmic Workflow

For each deep FFN block with $N$ tokens:

1. The importance score $S^{\mathrm{MHA}}_n$ is computed from the previous multi-head attention block.
2. The low-rank proxy $f$ is applied to all tokens. For $n = 1 \ldots N$, compute the FFN update proxy magnitude:

$C^{\mathrm{FFN}}_n = \|f(x_n)\|_2$

1. Condition the transformation score by attention:

$$S^{\mathrm{FFN}}_n = C^{\mathrm{FFN}}_n \cdot S^{\mathrm{MHA}}_n$$

1. Select the $M$ highest-scoring tokens—indices $T_{\mathrm{active}}$.
2. Only $X[T_{\mathrm{active}}]$ undergo the full FFN computation, producing:

$$\hat{Y}[T_{\mathrm{active}}] = X[T_{\mathrm{active}}] + F_{\mathrm{FFN}}(X[T_{\mathrm{active}}])$$

$\hat{Y}[\text{others}] = X[\text{others}]$

This greedy, layerwise procedure yields fine-grained token-skipping at minimal additional computational overhead.

4. Training-Free Proxy Construction and Integration

LTP employs a single-shot, training-free construction workflow:

• Original FFN weights are extracted.
• A small, representative calibration set $G$ is run to compute per-channel saliency.
• Structured pruning selects top-$D_{\mathrm{low}}$ channels.
• SVD is applied to each truncated matrix for low-rank decomposition.
• The resulting factors $U_i, V_i, b'_i$ are stored.
• At inference, the proxy network operates solely with precomputed low-rank linear maps and elementwise operations.

This workflow does not require any additional finetuning, gradient computation, or alteration of the base model, rendering LTP compatible with off-the-shelf LLM deployments.

5. Empirical Speed-Accuracy Trade-Offs

On LLaMA-3.1-8B-Instruct with up to 32K-token contexts, SPTS—which incorporates PAP, LTP, and Multi-Stage Delayed Pruning (MSDP)—achieves:

Metric	SPTS (PAP+LTP+MSDP)	Baseline (FTP)
TTFT Speedup	2.46×	~1.32×
E2E Speedup	2.29×	—
Avg. LongBench Drop	0.36% (47.98→47.62)	5.92% (absolute)

Ablation with only FFN skipping at 50% token budget shows that attention-only heuristics yield average scores of 45.29%, while LTP’s conditioned proxy raises this to 45.89% (+0.6%). This demonstrates substantial improvements in speed-accuracy trade-off and validates LTP's fine-grained selection mechanism.

6. Practical Considerations and Limitations

• Design Parameters: The selection of $D_{\mathrm{low}}$ and rank $r$ impacts the computational/accuracy balance; typically $r \approx 128$ is effective on 8B-parameter models.
• Calibration Set: Calibration requires only a few hundred sequences, provided they are representative of typical hidden states.
• Overhead Scaling: Proxy overhead is linear in $N D r$; for very large context windows, overhead can become non-negligible, but the MSDP mechanism alleviates this by progressive pruning in stages.
• Assumptions: LTP presumes that individual per-token FFN updates dominate the representational change, which may not hold for tasks requiring intricate token interactions.
• Greedy Selection: The procedure is greedy at the layer level, with no guarantee of global optimality; nevertheless, consistency across multiple models and tasks suggests robustness.

LTP transforms estimation of token-specific FFN updates—a computation previously expensive and prohibitive—into a lightweight one-shot process using optimized subnetwork proxies. This mechanism supports dynamic, high-precision FFN skipping at runtime, facilitating efficient long-context inference without any training or finetuning overhead (Wu et al., 19 Jan 2026).