PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse Token Routing in Efficient Transformers

Sai Krishna Arthanari, JaeHyeong Chang, Chengzhe Sun, Siwei Lyu

arXiv:2608.20632v1cs.CL

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.

    Source: §I · p. 1

  • 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.

    Source: A. SEWN core · p. 2

  • 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.

    Source: A. Does routing beat parameter-matched… · p. 3, §VI · p. 8

  • 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.

    Source: §VI · p. 8, §VIII · p. 9

    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.

    Source: §VII · p. 9

Abstract

from arXiv · show

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.Source: §I · p. 1
  • SEWN makes routing explicit and tests both claims against parameter-matched, non-adaptive baselines and post-hoc explanation methods.Source: §I · p. 1
  • The lexicon-seeded gate fails BoolQ counterfactual faithfulness despite achieving normal task accuracy.Source: §I · p. 1
  • Removing the static prior fixes the failure on BERT with little accuracy cost, but the result does not cleanly generalize across backbones or architectures.Source: §I · p. 1
  • Random selection eliminates significant faithfulness, showing that learned token ranking—not hard sparsity alone—drives the effect.Source: §I · p. 2
  • The study also reports capacity collapse on long-passage multiple-choice comprehension as a limitation shared by compressed models.Source: §I · p. 2

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.Source: §II · p. 2
  • Mixture-of-Depths routes tokens through variable numbers of Transformer layers, providing a depth-adaptive comparison for SEWN.Source: §II · p. 2
  • DistilBERT and TinyBERT compress BERT through logit and hidden-state distillation, with DistilBERT serving as the primary efficient baseline.Source: §II · p. 2
  • Attention transfer distills attention maps and is methodologically close to SEWN’s gate-score distillation.Source: §II · p. 2
  • Faithfulness research motivates evaluating explanations empirically, while ERASER formalizes mask-and-measure comprehensiveness and sufficiency.Source: §II · p. 2
  • SEWN computes its routing decision with a per-token MLP gate, available with either a lexicon-derived static prior or no prior.Source: A. SEWN core · p. 2

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.Source: B. SEWN-sparse: making routing a… · p. 2
  • The gate score splits each token embedding between a lightweight function-word stream and a BERT-initialized content stream before fusion.Source: Figure 1 · p. 3
  • In the selection view, shaded tokens reach the content stream, whereas selected and unselected tokens both reach the cheap function stream.Source: Figure 2 · p. 3
  • 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.Source: B. SEWN-sparse: making routing a… · p. 3
  • SEWN-lean removes the function stream and fusion, leaving the gate as the only per-token mechanism in a single content stream.Source: C. SEWN-lean: dropping the second stream · p. 3
  • Random-top-k uses the same architecture but samples content-stream tokens uniformly, isolating learned ranking from the hard-selection bottleneck.Source: D. SEWN-random-topk (ablation only) · p. 3
  • The experimental setup reports three random seeds and evaluates faithfulness by masking top-, bottom-, and random-ranked token subsets.Source: §IV · p. 3

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.

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.

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.

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.Source: F. Task-category breadth · p. 6
  • 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.Source: F. Task-category breadth · p. 6, G. Why is sparse routing’s signal… · p. 8
  • 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.Source: F. Task-category breadth · p. 6

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.Source: G. Why is sparse routing’s signal… · p. 7
  • 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.Source: G. Why is sparse routing’s signal… · p. 7
  • 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×.Source: G. Why is sparse routing’s signal… · pp. 7–8, Figure 6 · p. 8

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.Source: H. Repairing a broken faithfulness… · p. 8
  • 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.Source: H. Repairing a broken faithfulness… · p. 8
  • Distillation transfers a stable, content-determined gate property rather than accuracy alone.This aligns with the finding that faithfulness and accuracy are separable quantities.Source: H. Repairing a broken faithfulness… · p. 8

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.Source: §VI · p. 8
  • 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.Source: §VI · p. 8

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.Source: §VI · p. 8
  • 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.Source: §VI · p. 8

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.Source: §VII · p. 9
  • 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.Source: §VII · p. 9
  • 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.Source: §VII · p. 9
  • 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.Source: §VII · p. 9

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.Source: §VIII · p. 9
  • 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.Source: §VIII · p. 9

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.Source: APPENDIX: REPRODUCIBILITY NOTES · p. 10

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.Source: APPENDIX: REPRODUCIBILITY NOTES · p. 10

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.Source: APPENDIX: REPRODUCIBILITY NOTES · p. 10

Related papers

Loading 2608.20632v1…