Source-linked AI summary
Sparse Token Routing in Efficient Transformers
Sai Krishna Arthanari, JaeHyeong Chang, Chengzhe Sun, Siwei Lyu
TL;DR
The paper asks whether token-routing methods genuinely save computation and whether their importance signals are faithful explanations. Using SEWN, a two-stream Transformer with learned routing and sparse top-k selection, it finds substantial throughput gains, little routing-specific accuracy benefit, and faithfulness that depends on contextual learned ranking rather than hard sparsity alone. The static lexicon prior fails on BoolQ, while contextual gating succeeds in the reported BERT evaluations but has backbone-dependent limitations.
Problem
Adaptive-transformer research lacks separate end-to-end evidence for whether token routing saves compute and whether its importance signal faithfully explains predictions.
Method
SEWN uses a learned gate to route tokens between lightweight and full-capacity streams, with SEWN-sparse hard-selecting top-k tokens and counterfactual masking testing faithfulness.
Results
5.2–8.7× throughput over BERT-base and 2.6–4.4× over DistilBERT is achieved across evaluated datasets, while routing changes accuracy little relative to parameter-matched baselines and learned ranking drives faithfulness.
Takeaways & Limitations
Hard token-selection routing provides real speedups beyond standard distillation, but its importance signal is reliable only under specific content-determined gate designs.
Takeaways & Limitations
The faithfulness advantage of contextual gating over a static prior is backbone-dependent, and the paper lacks direct comparisons with most released token-pruning implementations.
Abstract
from arXiv · showhide
Efficient-transformer research often motivates token pruning and adaptive computation with the claim that not all tokens require equal computational effort. We test this claim end to end using SEWN, a two-stream Transformer that routes tokens through either lightweight or full-capacity processing using a learned gate. Across our experiments, routing introduces negligible accuracy change relative to parameter-matched baselines, while the gate's token-importance signal depends critically on how it is learned. A static lexicon-seeded prior fails a counterfactual faithfulness test on BoolQ, whereas a fully contextual gate achieves highly significant separation ($p<10^{-10}$) on both evaluated tasks without changing task accuracy.
I. INTRODUCTION
The paper separately tests efficiency and interpretability claims behind adaptive token computation using SEWN, rather than assuming routing signals are faithful. It finds that faithfulness can fail despite normal accuracy, depends on gate design, and requires falsification against alternative explanations.
- Adaptive computation bundles an efficiency claim about saving compute and an interpretability claim about using token importance as explanation.
- SEWN makes routing explicit and tests both claims against parameter-matched, non-adaptive baselines and post-hoc explanation methods.
- The lexicon-seeded gate fails BoolQ counterfactual faithfulness despite achieving normal task accuracy.
- Removing the static prior fixes the failure on BERT with little accuracy cost, but the result does not cleanly generalize across backbones or architectures.
- Random selection eliminates significant faithfulness, showing that learned token ranking—not hard sparsity alone—drives the effect.
- The study also reports capacity collapse on long-passage multiple-choice comprehension as a limitation shared by compressed models.
II. RELATED WORK
Related work spans token pruning, token merging, depth-adaptive routing, distillation, and empirical faithfulness evaluation. SEWN connects these areas by learning an explicit gate and evaluating its importance signal under a shared masking protocol.
- Token-pruning methods eliminate tokens using learned significance, thresholds, or reinforcement-learned policies, while ToMe merges tokens based on similarity.
- Mixture-of-Depths routes tokens through variable numbers of Transformer layers, providing a depth-adaptive comparison for SEWN.
- DistilBERT and TinyBERT compress BERT through logit and hidden-state distillation, with DistilBERT serving as the primary efficient baseline.
- Attention transfer distills attention maps and is methodologically close to SEWN’s gate-score distillation.
- Faithfulness research motivates evaluating explanations empirically, while ERASER formalizes mask-and-measure comprehensiveness and sufficiency.
- SEWN computes its routing decision with a per-token MLP gate, available with either a lexicon-derived static prior or no prior.
B. SEWN-sparse: making routing a compute claim
SEWN-sparse turns soft routing into compute savings by sending only gate-selected top-k tokens through the expensive content stream while retaining full-sequence processing in the cheap stream. Random selection isolates whether faithfulness comes from learned ranking or merely from the bottleneck.
- SEWN-sparse hard-selects the top-k tokens by gate importance for the expensive content stream, while the cheap function stream processes the full sequence.
- The gate score splits each token embedding between a lightweight function-word stream and a BERT-initialized content stream before fusion.
- In the selection view, shaded tokens reach the content stream, whereas selected and unselected tokens both reach the cheap function stream.
- Cross-attention lets the k-length content representation query the full-length function representation without paying O(n^2) attention cost in the expensive stream.
- SEWN-lean removes the function stream and fusion, leaving the gate as the only per-token mechanism in a single content stream.
- Random-top-k uses the same architecture but samples content-stream tokens uniformly, isolating learned ranking from the hard-selection bottleneck.
- The experimental setup reports three random seeds and evaluates faithfulness by masking top-, bottom-, and random-ranked token subsets.
V. RESULTS
Routing changes task accuracy little, but the reliability of the importance signal depends strongly on how the gate is learned. The contextual-versus-static-prior result is backbone- and task-dependent.
- 69.68 is the best BoolQ score among SEWN variants, while the full architecture remains statistically indistinguishable from a uniform 50/50 gate at 68.76 ± 0.19.
- Faithfulness differs despite nearly identical accuracy between static-prior and contextual gates.
- 68.85 to 69.20 on BoolQ and 63.66 to 63.80 on SWAG accompany removal of the lexicon prior, at essentially no accuracy cost.
- On RoBERTa, both gates pass on SWAG, but the contextual-gate faithfulness advantage does not replicate on BoolQ.
- The surviving claim is narrower: large hand-built linguistic priors never improve over contextual gates in the tested tasks and sizes.
C. Making sparsity a compute claim: SEWN-sparse and the k-sweep
SEWN-sparse turns routing into actual computation savings by selecting only top-k tokens for the expensive stream. The k-sweep identifies task-dependent operating points that retain faithfulness while balancing speed and accuracy.
- SEWN-sparse hard-selects top-k tokens before the expensive content stream, while the cheap function stream processes the full sequence.
- Every tested k passes the faithfulness test at p < 10^-9, while accuracy plateaus around k = 96 rather than improving toward k = 128.
- k = 96 is used for BoolQ-scale tasks and k = 128 for longer-context PubMedQA; short-sequence tasks use conservatively low k.
- Table IV measures dense SEWN-contextual-gate speed at 2.23× over BERT-base before pruning, establishing architectural savings separately from sparsity.
D. Generalizing to scale: PubMedQA
PubMedQA extends evaluation to roughly three times the training data and sequence length, preserving the efficiency pattern while producing stronger faithfulness effects. SEWN-sparse also leads the joint faithfulness-throughput comparison.
- PubMedQA uses 26,980 training examples and mean sequence length 365, roughly three times BoolQ’s data and token length.
- 6.36× speedup over BERT-base on PubMedQA closely matches BoolQ’s 6.14×, while faithfulness improves rather than degrades at scale.
- The strict faithfulness criterion requires top-k masking to cause larger confidence drops than both bottom-k and random-k masking, with random-k exceeding bottom-k.
- Both reported effect sizes exceed the corresponding BoolQ results, plausibly reflecting a sharper content-versus-filler contrast in longer technical text.
- All three post-hoc methods pass the faithfulness test, including raw attention at essentially normal-forward-pass cost.
- 0.328 top-drop and 0.006 bottom-drop make SEWN-sparse the strongest comparison result and the fastest model on the faithfulness-throughput frontier.
- Integrated gradients costs 46× the compute of the comparison baseline, corresponding at million-example scale to roughly 46 hours versus 7.4 minutes for SEWN-sparse.
F. Task-category breadth
Task behavior varies by category: routing is informative and efficient on several classification tasks, but long-passage comprehension exposes a broader compressed-model capacity limitation. Counterfactual controls further show that faithfulness depends on the ranking signal, not merely hard token exclusion.
- Task-category breadth: Winogrande is uninformative: every model, including full BERT-base, remains within roughly two points of chance at 50%.The paper does not interpret this result as evidence against SEWN specifically.
- Task-category breadth: RACE exposes a capacity limitation: DistilBERT falls 13.1 points relative to BERT-base, while SEWN-sparse reaches 42.97 ± 0.32.The collapse affects compressed models broadly, though the authors cannot fully rule out undertraining from the fixed recipe.
- Task-category breadth: DistilBERT stays close to BERT-base on IMDB, while SEWN-sparse is 5.2× faster but costs 6.51 points of accuracy.The reported IMDB values are −0.91 points for DistilBERT and 33.60ms to 6.42ms for SEWN-sparse.
- Faithfulness controls: p = 0.45: random selection with the identical bottleneck fails to produce significant faithfulness, showing that hard exclusion alone is insufficient.The learned ranking, rather than the mechanical bottleneck alone, accounts for the explanatory signal.
- Faithfulness controls: p = 1.8×10^-13: SEWN-attn-topk passes the faithfulness test, but its 64.60ms latency exceeds SEWN-sparse’s 7.13ms and BERT-base’s 45.31ms.The attention-based ranking requires a full extra frozen BERT-base forward pass.
- Faithfulness controls: Top-drop 0.124 versus random-drop 0.049 and bottom-drop 0.052: Mixture-of-Depths has significant signal but does not match SEWN-sparse’s strict ordering result.Its speedup is 1.20× over BERT-base, compared with SEWN-sparse’s 6.36×.
H. Repairing a broken faithfulness result via distillation
Distillation repairs SEWN-lean’s broken BoolQ faithfulness by transferring the teacher’s stable, content-determined gate ranking. It largely restores the teacher’s accuracy while preserving an already-passing SWAG result.
- Repair via distillation: p: 0.205 to 1.6 × 10^-7: adding logit-KL distillation and gate-score MSE fixes BoolQ’s broken faithfulness result.The loss is L = Ltask + α · Llogit-KD + β · Lgate-MSE, with α = 1.0 and β = 2.0.
- Repair via distillation: 68.33 to 68.95: distillation closes most of SEWN-lean’s accuracy gap to the 69.20 teacher.SWAG’s already-passing faithfulness result remains intact and marginally sharpens.
- Repair via distillation: Distillation transfers a stable, content-determined gate property rather than accuracy alone.This aligns with the finding that faithfulness and accuracy are separable quantities.
VI. DISCUSSION
The discussion narrows the paper’s claims: routing yields substantial efficiency gains, but accuracy improvement and uniform task generalization do not survive scrutiny. Faithfulness is conditional on the gate, backbone, evaluation scope, and model scale.
- Claims that survive: 5.2–8.7× over BERT-base and 2.6–4.4× over DistilBERT: hard top-k routing provides a substantial compute win across tested datasets.The paper reports this gain across a 3× change in data scale and sequence length.
- Claims that survive: SEWN-sparse has the cleanest top-vs-bottom separation, while Mixture-of-Depths shows significant signal without clearing the strict ordering PASS bar.This comparison uses a direct Mixture-of-Depths reimplementation.
- Claims that fail: Routing does not meaningfully improve task accuracy over parameter-matched non-adaptive baselines.The discussion explicitly lists this as a claim that does not survive.
- Claims that fail: The efficiency/faithfulness package is not uniform: it holds for binary classification and sentiment but collapses on long-passage multiple-choice comprehension.The collapse affects every compressed model tested, not only routing-based models.
- Scope boundaries: The study does not directly compare against most released token-pruning implementations, although it directly reimplements Mixture-of-Depths.SEWN-attn-topk supplies a non-learned attention-ranking comparison.
- Scope boundaries: The study evaluates BERT-base-scale backbones only and makes no claim about transfer to billion-parameter models.Routing economics and post-hoc explanation costs could differ at larger scale.
- Scope boundaries: On RoBERTa/BoolQ, the contextual gate fails significance at p = 0.177 despite passing on BERT at p = 1.7 × 10^-10.The authors found no tested prior-size, coverage, or content confound explaining the gap.
- Scope boundaries: Faithfulness is behavioral rather than mechanistic: passing the masking test does not establish unique or true causal explanations.Multiple valid rankings could pass the same test.
VIII. CONCLUSION
The paper tests efficiency and interpretability claims separately rather than assuming them. Efficiency holds with scale-robust speedups but accuracy costs, while faithfulness depends on a contextual, content-stable gate and remains falsifiable.
- Conclusion: The efficiency claim holds with scale-robust speedups, but at an accuracy cost that must be reported by task category.The conclusion emphasizes honest per-category reporting rather than a uniformly positive result.
- Conclusion: Faithfulness depends on a contextual rather than lexicon-anchored gate and on a stable, content-determined selection criterion.The conclusion presents both properties as testable and falsifiable.
APPENDIX
The appendix documents the fixed experimental setup, released result artifacts, and a reproducibility exception for Table VII throughput and cost numbers.
- Experimental setup: All experiments use a single NVIDIA RTX 4090, fixed software versions, three seeds, and a training recipe held constant across variants and datasets unless noted.The seeds are 42, 123, and 456.
- Artifacts: Result files are available in the accompanying code repository for the reported tables and several analysis checks.Listed artifacts cover ablations, counterfactual evaluations, backbone checks, prior sweeps, and related experiments.
- Reproducibility caveat: Table VII batched throughput and cost columns have no saved JSON artifact and require rerunning sewn_throughput_explanations.py for reproduction.The draft recommends rerunning the script before citing those specific numbers in a camera-ready version.