Token-Selective Propagation in Deep Models

• Token-Selective Propagation (TSP) is a family of mechanisms that selectively forward salient tokens to lower quadratic attention costs while maintaining task-relevant accuracy.
• It encompasses various approaches—including FastKV, TPC-ViT, MambaMixer, GTP-ViT, and ToSA—that differ in where and how token selection is computed and applied.
• Empirical results across LLMs, vision transformers, and time series models demonstrate TSP’s effectiveness in accelerating processing and reducing FLOPs with minimal performance drops.

Searching arXiv for the cited TSP-related papers to ground the article in current literature. Token-Selective Propagation (TSP) denotes a family of token-selection and token-routing mechanisms that reduce sequence-level computation by allowing only a subset of tokens, or token-derived states, to participate in expensive downstream operations while attempting to preserve task-relevant information. In the recent literature, the term does not identify a single canonical algorithm. In FastKV, TSP is a late-layer decoder mechanism that propagates only salient context tokens beyond a designated layer while decoupling propagated hidden states from KV-cache retention (Jo et al., 3 Feb 2025). In MambaMixer, the role corresponds to a Selective Token Mixer that gates tokenwise signals before a data-dependent state-space scan (Behrouz et al., 2024). In vision transformers, related formulations include token pausing and restarting in TPC-ViT (Zhu, 2024), graph-based propagation from pruned to retained tokens in GTP-ViT (Xu et al., 2023), and top-$K$ token selection with layer bypass and re-concatenation in ToSA (Singh et al., 2024).

1. Scope, terminology, and recurring design pattern

Across these works, TSP addresses a common systems problem: full token participation imposes quadratic attention cost in transformers or otherwise forces all tokens through a uniform computation path. The proposed remedies differ in where the selection signal is computed, whether dropped tokens can re-enter, and whether information is discarded, cached, summarized, or bypassed.

Paper	Selection signal	Propagation rule
FastKV	attention from the last $N_{\mathrm{obs}}$ “window” tokens	top $\lfloor \rho N_I \rfloor$ tokens plus window tokens are forwarded after $\ell_{\mathrm{tsp}}$
MambaMixer	$G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$	gated token stream is processed by a selective SSM
TPC-ViT	pause probability $p_i^k$, restart probability $r_i^k$	tokens may pause in layer $i$ and restart later
GTP-ViT	score $\mathrm{Score}_i=\gamma_i\psi_i$	pruned-token features are propagated to kept neighbors
ToSA	predicted next-layer attention maps $\hat A_{i+1}^h$	top-$N_{\mathrm{obs}}$0 tokens attend; the rest bypass the layer

A central empirical premise recurs across the papers: token relevance is depth-dependent rather than static. FastKV states that early decoder layers have highly unstable attention focus, whereas deeper layers exhibit stabilization of critical context (Jo et al., 3 Feb 2025). TPC-ViT explicitly argues that “tokens that are redundant in one layer can be useful in later layers” (Zhu, 2024). ToSA and GTP-ViT similarly avoid equating non-attended or pruned tokens with useless tokens: ToSA keeps skipped tokens in the representation via bypass, and GTP-ViT injects pruned-token information into retained tokens through graph propagation (Singh et al., 2024, Xu et al., 2023).

2. FastKV: late-layer Token-Selective Propagation in decoder LLMs

In FastKV, TSP is the mechanism that reduces prefill computation by forwarding only salient tokens after a dedicated TSP layer $N_{\mathrm{obs}}$1, while KV cache compression is performed independently at retention rate $N_{\mathrm{obs}}$2 (Jo et al., 3 Feb 2025). The saliency score is computed from the attention received by each token from the most recent $N_{\mathrm{obs}}$3 query positions. For head $N_{\mathrm{obs}}$4 and token $N_{\mathrm{obs}}$5,

$N_{\mathrm{obs}}$6

and the head-averaged saliency is

$N_{\mathrm{obs}}$7

Tokens are ranked by $N_{\mathrm{obs}}$8, and the top $N_{\mathrm{obs}}$9 tokens are selected, with all $\lfloor \rho N_I \rfloor$0 window tokens always included in the propagation set.

The TSP layer index is chosen as the earliest layer at which TSP remains close to the full-context baseline in the final hidden states. With $\lfloor \rho N_I \rfloor$1 as the full-context final hidden state for calibration input $\lfloor \rho N_I \rfloor$2, and $\lfloor \rho N_I \rfloor$3 the counterpart obtained when TSP is applied at layer $\lfloor \rho N_I \rfloor$4, FastKV chooses

$\lfloor \rho N_I \rfloor$5

subject to $\lfloor \rho N_I \rfloor$6 being as small as possible. In practice, the sweep is performed in middle layers, for example $\lfloor \rho N_I \rfloor$7 in a 32-layer model.

Algorithmically, the model processes the full prompt of length $\lfloor \rho N_I \rfloor$8 up to $\lfloor \rho N_I \rfloor$9. At every layer, KVCompress is applied to $\ell_{\mathrm{tsp}}$0 at retention rate $\ell_{\mathrm{tsp}}$1. When $\ell_{\mathrm{tsp}}$2, TSP computes $\ell_{\mathrm{tsp}}$3, selects the propagation set, and HiddenCompress forms a reduced hidden state of length approximately $\ell_{\mathrm{tsp}}$4. All later layers then operate only on that reduced hidden state. The resulting prefill complexity drops from

$\ell_{\mathrm{tsp}}$5

to

$\ell_{\mathrm{tsp}}$6

The per-token decoding cost becomes $\ell_{\mathrm{tsp}}$7 rather than $\ell_{\mathrm{tsp}}$8.

The two hyper-parameters serve different budgets. The TSP rate $\ell_{\mathrm{tsp}}$9 sets the fraction of propagated tokens and therefore the prefill compute rate; smaller $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$0 increases prefill speedup but risks discarding critical tokens. The KV retention rate $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$1 sets the fraction of KV pairs stored per layer for decoding; smaller $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$2 reduces memory bandwidth and generation cost. FastKV emphasizes that $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$3 and $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$4 are independently tunable, unlike GemFilter and PyramidInfer where $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$5 by design.

Empirically, on LLaMA-3.1-8B-Instruct with $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$6 tokens, FastKV reports prefill speedup up to $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$7 at $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$8 and $G=\sigma(\mathrm{Conv}_{1D}(\mathrm{Linear}_1(x)))$9, decoding speedup up to $p_i^k$0 at $p_i^k$1, and LongBench average accuracy within $p_i^k$2 of the full-context baseline, with one example reported as $p_i^k$3 at $p_i^k$4. Ablations show that accuracy plateaus around $p_i^k$5, that large accuracy loss occurs if $p_i^k$6, and that prefill savings vanish if $p_i^k$7.

3. MambaMixer: selective token propagation through a data-dependent SSM

In MambaMixer, the component corresponding to TSP is the Selective Token Mixer, which operates on an input tensor $p_i^k$8 with expansion dimension $p_i^k$9 and uses data-dependent weights together with a selective SSM (Behrouz et al., 2024). Its unidirectional layer-$r_i^k$0 update is: $r_i^k$1

$r_i^k$2

$r_i^k$3

$r_i^k$4

$r_i^k$5

$r_i^k$6

$r_i^k$7

Here the primary selection signal is the gate tensor

$r_i^k$8

while $r_i^k$9, $i$0, and $i$1 provide data-dependent selection of the continuous-time SSM. The stated interpretation is that tokens are first scored by depth-wise $i$2D convolution plus sigmoid, then those scores are used to mask or compress the sequence before the selective SSM, and the MLP path adds non-recurrent mixing.

MambaMixer does not isolate token selection from the rest of the block. Instead, the Selective Token Mixer and the Selective Channel Mixer are coupled by weighted averaging across layers. If $i$3 and $i$4 denote the outputs of the $i$5-th token and channel mixers, then the next inputs are

$i$6

$i$7

where $i$8 are scalar learnable weights and $i$9. This gives later layers direct access to early features and to any earlier token- or channel-filtered representation.

The per-block complexity is reported as approximately $\mathrm{Score}_i=\gamma_i\psi_i$0 for the Token Mixer and $\mathrm{Score}_i=\gamma_i\psi_i$1 for the Channel Mixer, for total time $\mathrm{Score}_i=\gamma_i\psi_i$2. Working memory per block is approximately $\mathrm{Score}_i=\gamma_i\psi_i$3, with the text emphasizing linear dependence on $\mathrm{Score}_i=\gamma_i\psi_i$4 and $\mathrm{Score}_i=\gamma_i\psi_i$5. The implementation uses the GPU-friendly parallel associative scan for the recurrent S6 block, and memory is described as scaling roughly linearly in $\mathrm{Score}_i=\gamma_i\psi_i$6 and $\mathrm{Score}_i=\gamma_i\psi_i$7, without an $\mathrm{Score}_i=\gamma_i\psi_i$8 term as in self-attention.

As a proof of concept, MambaMixer instantiates Vision MambaMixer (ViM2) and Time Series MambaMixer (TSM2). On ImageNet-1K, ViM2-Tiny is reported at approximately $\mathrm{Score}_i=\gamma_i\psi_i$9M parameters and $\hat A_{i+1}^h$0G FLOPs with $\hat A_{i+1}^h$1 top-1 accuracy, ViM2-Small at approximately $\hat A_{i+1}^h$2M and $\hat A_{i+1}^h$3G with $\hat A_{i+1}^h$4, and ViM2-Base at approximately $\hat A_{i+1}^h$5M and $\hat A_{i+1}^h$6G with $\hat A_{i+1}^h$7. The comparison quoted in the summary includes ViM-Small at $\hat A_{i+1}^h$8M and approximately $\hat A_{i+1}^h$9G with $N_{\mathrm{obs}}$00, and swin-Tiny at $N_{\mathrm{obs}}$01M and approximately $N_{\mathrm{obs}}$02G with $N_{\mathrm{obs}}$03. For long-series forecasting, TSM2 is reported to outperform prior transformers and MLP mixers on $N_{\mathrm{obs}}$04 public benchmarks, often reducing MSE by $N_{\mathrm{obs}}$05–$N_{\mathrm{obs}}$06 while using $N_{\mathrm{obs}}$07–$N_{\mathrm{obs}}$08 fewer FLOPs.

4. TPC-ViT: pause, restart, and cumulative halting

TPC-ViT formulates TSP as an adaptive halting mechanism in which a token can be paused in one layer and restarted in a later layer, rather than being discarded permanently (Zhu, 2024). For token $N_{\mathrm{obs}}$09 at layer $N_{\mathrm{obs}}$10, two Bernoulli variables are introduced: $N_{\mathrm{obs}}$11 for pause, with $N_{\mathrm{obs}}$12, and $N_{\mathrm{obs}}$13 for restart, with $N_{\mathrm{obs}}$14. The break probability is

$N_{\mathrm{obs}}$15

The probabilities $N_{\mathrm{obs}}$16 and $N_{\mathrm{obs}}$17 are derived from two scalar components of the token representation $N_{\mathrm{obs}}$18, with no new parameters except learned scale $N_{\mathrm{obs}}$19 and shift $N_{\mathrm{obs}}$20: $N_{\mathrm{obs}}$21 Cumulative break determines halting. Once

$N_{\mathrm{obs}}$22

token $N_{\mathrm{obs}}$23 is removed at layer $N_{\mathrm{obs}}$24.

TPC-ViT adds two stabilizing components. First, a smoothing mechanism regularizes $N_{\mathrm{obs}}$25 toward the layerwise mean over the active token set $N_{\mathrm{obs}}$26. With

$N_{\mathrm{obs}}$27

the smoothed break probability is

$N_{\mathrm{obs}}$28

An equivalent penalty is

$N_{\mathrm{obs}}$29

Second, a model stabilizer injects local bias by sparse attention: for each query $N_{\mathrm{obs}}$30, only the $N_{\mathrm{obs}}$31 nearest keys by Euclidean distance are attended, with

$N_{\mathrm{obs}}$32

Training combines classification loss, a ponder loss that encourages early stopping, a layer-level distribution loss that encourages a bell-shaped halting distribution, and the smoothing regularizer: $N_{\mathrm{obs}}$33 The simplified algorithm initializes the token mask, cumulative break, and remainder; applies StabilizerAttention at each layer; computes $N_{\mathrm{obs}}$34, $N_{\mathrm{obs}}$35, and $N_{\mathrm{obs}}$36; smooths them; updates cumulative break; and uses the resulting mask in the following self-attention.

On ImageNet-1K, the reported results include DeiT-S $N_{\mathrm{obs}}$37 TPC-DeiT-S, with top-1 accuracy $N_{\mathrm{obs}}$38, GFLOPs $N_{\mathrm{obs}}$39, throughput at $N_{\mathrm{obs}}$40 $N_{\mathrm{obs}}$41, and average tokens per layer from $N_{\mathrm{obs}}$42 to approximately $N_{\mathrm{obs}}$43. For DeiT-T, accuracy is reported as $N_{\mathrm{obs}}$44 and GFLOPs as $N_{\mathrm{obs}}$45. For Swin-T, the summary reports $N_{\mathrm{obs}}$46 with GFLOPs $N_{\mathrm{obs}}$47. For LV-ViT-S, it reports $N_{\mathrm{obs}}$48 with GFLOPs $N_{\mathrm{obs}}$49. Ablations show that on DeiT-S, the stabilizer improves top-1 from $N_{\mathrm{obs}}$50 to $N_{\mathrm{obs}}$51, and the token distribution regularizer improves TPC-DeiT-S from $N_{\mathrm{obs}}$52 to $N_{\mathrm{obs}}$53. On DeiT-T, $N_{\mathrm{obs}}$54 yields the best reported top-1 accuracy of $N_{\mathrm{obs}}$55, compared with $N_{\mathrm{obs}}$56 at $N_{\mathrm{obs}}$57.

5. GTP-ViT: graph-based propagation after token selection

GTP-ViT instantiates token propagation as graph summarization: less significant tokens are not simply removed, but their features are propagated to spatially and semantically connected retained tokens (Xu et al., 2023). After patch embedding, the input is $N_{\mathrm{obs}}$58, where each token $N_{\mathrm{obs}}$59 becomes a node $N_{\mathrm{obs}}$60 in a weighted graph $N_{\mathrm{obs}}$61.

Two sparse graphs are built and then combined. The spatial graph uses adjacency $N_{\mathrm{obs}}$62 with $N_{\mathrm{obs}}$63 when tokens $N_{\mathrm{obs}}$64 and $N_{\mathrm{obs}}$65 correspond to spatially adjacent patches and $N_{\mathrm{obs}}$66 otherwise. The semantic graph uses cosine similarity in the initial embedding $N_{\mathrm{obs}}$67: $N_{\mathrm{obs}}$68 If $N_{\mathrm{obs}}$69 is the $N_{\mathrm{obs}}$70-th largest cosine similarity between $N_{\mathrm{obs}}$71 and all other tokens, then

$N_{\mathrm{obs}}$72

The mixed adjacency is

$N_{\mathrm{obs}}$73

with $N_{\mathrm{obs}}$74.

Token importance is computed from the multi-head attention map in a transformer block. With a permutation-invariant reducer $N_{\mathrm{obs}}$75 across heads, regeneration difficulty is

$N_{\mathrm{obs}}$76

broadcasting ability is

$N_{\mathrm{obs}}$77

and the combined score is

$N_{\mathrm{obs}}$78

If $N_{\mathrm{obs}}$79 tokens are to be pruned from $N_{\mathrm{obs}}$80, GTP sorts tokens by $N_{\mathrm{obs}}$81, excludes the $N_{\mathrm{obs}}$82 token, retains the top $N_{\mathrm{obs}}$83 tokens as $N_{\mathrm{obs}}$84, and marks the bottom $N_{\mathrm{obs}}$85 as propagated tokens $N_{\mathrm{obs}}$86.

Propagation is then a single sparse-matrix update. If $N_{\mathrm{obs}}$87 is the kept-by-pruned block of $N_{\mathrm{obs}}$88, the summarized tokens are

$N_{\mathrm{obs}}$89

with $N_{\mathrm{obs}}$90 as a small hyper-parameter. This avoids one-to-one token matching: each kept token can absorb information from all adjacent pruned tokens.

The complexity analysis contrasts GTP with vanilla MHSA and token merging. Vanilla MHSA costs $N_{\mathrm{obs}}$91 per layer. After pruning to $N_{\mathrm{obs}}$92 tokens, self-attention costs

$N_{\mathrm{obs}}$93

plus propagation cost

$N_{\mathrm{obs}}$94

The summary states that token-selection cost $N_{\mathrm{obs}}$95 is negligible and that, with $N_{\mathrm{obs}}$96, GTP reduces quadratic attention cost by approximately $N_{\mathrm{obs}}$97 per layer.

On ImageNet-1K, for DeiT-S with baseline top-1 $N_{\mathrm{obs}}$98, pruning $N_{\mathrm{obs}}$99 tokens per layer yields $\lfloor \rho N_I \rfloor$00 GMACs, a $\lfloor \rho N_I \rfloor$01 reduction from $\lfloor \rho N_I \rfloor$02, throughput $\lfloor \rho N_I \rfloor$03, and top-1 $\lfloor \rho N_I \rfloor$04, corresponding to $\lfloor \rho N_I \rfloor$05. For DeiT-B with baseline $\lfloor \rho N_I \rfloor$06, the same $\lfloor \rho N_I \rfloor$07 yields $\lfloor \rho N_I \rfloor$08 GMACs, $\lfloor \rho N_I \rfloor$09 reduction, throughput $\lfloor \rho N_I \rfloor$10, and top-1 $\lfloor \rho N_I \rfloor$11, also $\lfloor \rho N_I \rfloor$12. The summary further states that at approximately $\lfloor \rho N_I \rfloor$13 GMACs, GTP on DeiT-S achieves $\lfloor \rho N_I \rfloor$14 without finetuning, outperforming ToMe at $\lfloor \rho N_I \rfloor$15 and EViT at $\lfloor \rho N_I \rfloor$16 at similar or higher throughput.

6. ToSA: selective attention with bypassed tokens and full-token reconstruction

ToSA implements token selectivity on pairs of consecutive ViT layers, where the first layer remains standard and the second is replaced by a ToSA transformer layer (Singh et al., 2024). Between them, a token selector takes the unnormalized attention maps

$\lfloor \rho N_I \rfloor$17

predicts the next-layer attention maps $\lfloor \rho N_I \rfloor$18, and extracts token-importance scores from those predictions.

For each head $\lfloor \rho N_I \rfloor$19,

$\lfloor \rho N_I \rfloor$20

For token index $\lfloor \rho N_I \rfloor$21, the importance score is the column sum

$\lfloor \rho N_I \rfloor$22

after which the top

$\lfloor \rho N_I \rfloor$23

tokens are selected as attention tokens. The remaining tokens are skip tokens.

If $\lfloor \rho N_I \rfloor$24 is the output of the first layer, the selected and skipped subsets are

$\lfloor \rho N_I \rfloor$25

Only $\lfloor \rho N_I \rfloor$26 participates in the second layer’s self-attention: $\lfloor \rho N_I \rfloor$27

$\lfloor \rho N_I \rfloor$28

The skip tokens bypass the layer unchanged,

$\lfloor \rho N_I \rfloor$29

and the full token set is reconstructed by concatenation: $\lfloor \rho N_I \rfloor$30 This preservation of one feature per original patch is the stated reason ToSA remains compatible with dense prediction tasks.

Training of the selector uses a frozen pretrained ViT to provide ground-truth attention maps $\lfloor \rho N_I \rfloor$31, with selector loss

$\lfloor \rho N_I \rfloor$32

Only the small convolutional selector is updated at that stage; the full ToSA-modified network is then finetuned end-to-end.

The FLOPs reduction follows directly from replacing $\lfloor \rho N_I \rfloor$33-token attention by $\lfloor \rho N_I \rfloor$34-token attention in each ToSA layer. A standard multi-head self-attention layer costs approximately $\lfloor \rho N_I \rfloor$35, whereas a ToSA layer costs approximately $\lfloor \rho N_I \rfloor$36. If ToSA is applied to $\lfloor \rho N_I \rfloor$37 out of $\lfloor \rho N_I \rfloor$38 self-attention layers, the total self-attention cost becomes

$\lfloor \rho N_I \rfloor$39

for a relative reduction

$\lfloor \rho N_I \rfloor$40

The summary gives an example for DeiT-Tiny with $\lfloor \rho N_I \rfloor$41, $\lfloor \rho N_I \rfloor$42, and $\lfloor \rho N_I \rfloor$43, yielding a $\lfloor \rho N_I \rfloor$44 self-attention FLOPs reduction.

The reported empirical configuration applies ToSA to DeiT-Tiny layers $\lfloor \rho N_I \rfloor$45 with $\lfloor \rho N_I \rfloor$46. On ImageNet-1K, the baseline DeiT-Tiny has top-1 $\lfloor \rho N_I \rfloor$47 and $\lfloor \rho N_I \rfloor$48 GFlops, whereas the ToSA version has top-1 $\lfloor \rho N_I \rfloor$49 and $\lfloor \rho N_I \rfloor$50 GFlops, corresponding to $\lfloor \rho N_I \rfloor$51 points and $\lfloor \rho N_I \rfloor$52 GFlops. On NYU Depth V2 with a NeWCRFs decoder, the baseline DeiT-Tiny encoder yields $\lfloor \rho N_I \rfloor$53 and $\lfloor \rho N_I \rfloor$54, while the ToSA encoder yields $\lfloor \rho N_I \rfloor$55 and $\lfloor \rho N_I \rfloor$56. The summary states that KITTI results are virtually identical.

7. Comparative interpretation and points of clarification

The literature shows that TSP is not synonymous with a single operation such as irreversible token pruning. FastKV performs one-time token selection after full-context early processing and continues later layers on a reduced hidden state, while separately compressing the KV cache (Jo et al., 3 Feb 2025). TPC-ViT allows pausing and later restarting, making token exclusion explicitly reversible at the architectural level (Zhu, 2024). GTP-ViT propagates features from removed tokens to retained tokens rather than bypassing them unchanged (Xu et al., 2023). ToSA bypasses skipped tokens and re-concatenates them, so token count in the representation remains constant even when only a subset participates in self-attention (Singh et al., 2024). MambaMixer, by contrast, uses token selection as a gating operation inside a selective SSM rather than as an attention-pruning rule (Behrouz et al., 2024).

A second recurring distinction is between selection for computation and selection for memory. FastKV makes this explicit by decoupling the TSP rate $\lfloor \rho N_I \rfloor$57 from the KV retention rate $\lfloor \rho N_I \rfloor$58, arguing that prefill compute reduction need not be tied to the decoding KV budget (Jo et al., 3 Feb 2025). This suggests a broader interpretation of TSP as a resource-allocation mechanism that can target hidden-state propagation, attention participation, state-space scanning, or cache storage, depending on the model family.

A third point concerns when token saliency should be measured. FastKV chooses the earliest layer at which critical-token sets have stabilized, and reports substantial degradation when TSP is applied too early. TPC-ViT states that redundancy at one layer does not imply redundancy at later layers. ToSA predicts next-layer attention from current attention, and GTP-ViT computes importance within each block from attention-derived criteria. Taken together, these results suggest that token selectivity is fundamentally depth-aware, and that methods assuming a fixed early pruning boundary or permanent redundancy may be mismatched to the observed dynamics.

Finally, the surveyed works reflect two distinct preservation strategies once full participation is curtailed. One strategy preserves information by carrying fewer but more salient units forward, as in FastKV and MambaMixer. The other preserves information by redistributing or retaining the representation of omitted tokens, as in GTP-ViT and ToSA. TPC-ViT occupies an intermediate position, since paused tokens are removed from immediate computation yet remain eligible for later restart. This suggests that the phrase “Token-Selective Propagation” functions less as a single named primitive than as a family resemblance term for architectures that make token flow conditional, depth-dependent, and explicitly tied to efficiency.