LapMerge: Differentiable Exact-Budget Token Merging in Linear Time
Authors: Magdalena Trędowicz, Aleksandra Nowak, Tomasz Steifer, Łukasz Struski, Szymon Płotka, Jacek Tabor
Organizations: Jagiellonian University in Kraków · Institute of Fundamental Technological Research, Polish Academy of Sciences · Centre for Credible AI, Warsaw University of Technology
Modern vision and language models process ever longer token sequences, and reducing them is the main lever on downstream cost. Token compression, however, faces a fundamental trade-off: selecting tokens discards information, while existing merging methods preserve it only through pairwise token-to-token matching or MN token-to-output interactions. We introduce LapMerge, a differentiable merging operator that resolves this trade-off for ordered sequences: a learned mass distribution partitions the sequence into exactly N equal-mass segments under a Laplace-kernel density, and prefix/suffix scans aggregate each segment, so both the forward and the backward pass run in O(M+N) time. This combination (an exact output budget, end-to-end differentiability, and linear cost) is precisely what previous reducers trade off against one another. The budget can be changed without changing operator parameters, and the learned mass doubles as a token-level map of where the budget is allocated. Across vision and vision-language tasks, LapMerge matches or improves on strong token reducers while being up to two orders of magnitude faster at large sequence lengths; it is particularly effective when token order reflects spatial or temporal locality, and works both as a drop-in reduction module and as a trainable architectural component.
Figures & tables
Figure 1: Overview of the standard LapMerge pipeline. Given an ordered sequence of M input tokens, LapMerge produces exactly N budget-exact, soft-merged output tokens. A learned scorer assigns each input token a mass, represented as a Laplace kernel. The summed density is scanned cumulatively and partitioned into N equal-mass segments, each of which is collapsed into a single output token. The resulting N tokens are then processed by a Transformer classifier. The rightmost panel shows a real histopathology slide with the learned scorer weights overlaid on the tissue; opacity indicates the importance attributed to each patch. These same scores determine the token merging and can therefore be read as a learned routing-mass map: a direct view of where the model allocates its token budget.
Figure 2: The same six tokens cut into N=3 pieces two ways – by token count they are 3-1-2, but by mass all three pieces are comparable. Top : the hard partition from the running sum – barriers can only fall between tokens, so each piece’s mass is merely close to 1/N . Bottom : treating positions as noisy turns each token into a smooth bump; the shaded area under the resulting mixture density (bold) is exactly 1/N per piece (Proposition 1 ). These intervals partition probability mass, not tokens: one token’s bump can span several intervals. The kernel shown is Gaussian, purely for illustration.
Figure 3: Why the Laplace kernel is fast : the tick at each kernel’s peak splits it into a left half and a right half, and Equation 6 aggregates every left half with one prefix sum and every right half with one suffix sum – never a per-pair computation. Five tokens shown: three peaks fall inside [bℓ−1,bℓ) ( xi,xj,xk ) and two fall outside ( xh,xl ); only the shaded, interval-clipped part of each counts toward tℓ (Definition 1 ).
Figure 4: Reducer forward–backward latency on GH200 ( B=1 , d=512 , FP32; five warm-ups, 30 repeats). Dashed: N=1024 ; solid: N=M/16 . LapMerge scales similarly in both regimes; the SDPA poolers retain MN interactions. Open circles precede a 600-s timeout, the filled OOM, the filled square an execution failure. Timing covers the reducer core (LapMerge gets precomputed masses; the poolers project inside the timed region); scorer/routing overheads: Appendix B.4 .
method
CRC
STAD
ABMIL
88.97±2.91
92.68±3.24
CLAM-SB
89.92±4.40
93.75±2.90
CLAM-MB
90.66±1.43
91.62±2.74
DS-MIL
90.33±7.60
87.11±3.35
DTFD-MIL
90.86±1.58
93.22±4.46
TransMIL
81.07±3.90
89.24±4.52
Table 1: Test AUROC (mean ± std, 6 seeds). External anchors are literature numbers, not reproduced here: Kather et al. 2019 is a CNN tile-classifier, not MIL; Wagner et al. 2023 is AttentionMIL over CTransPath features.
Method
10%
30%
FastV ( Chen et al., 2024a )
65.99
80.47
DART ( Wen et al., 2025 )
71.91
83.43
HoloV ( Zou et al., 2026 )
75.84
83.27
PruMerge+ ( Shang et al., 2025 )
76.52
83.69
VisionZip ( Yang et al., 2025 )
78.46
84.73
DivPrune ( Alvar et al., 2025 )
79.27
84.68
Table 2: Drop-in token reduction on POPE (F1, frozen Qwen2.5-VL-7B, no retraining) at matched 10% and 30% visual-token budgets. Baseline rows quote the published evaluation of Zhu et al. (2025) (their Table 1); our rerun of the released selection reproduces its F1 to 0.06 (Appendix D.3 ). LapMerge reuses the VisionSelector scorer at default widths (αr,αa)=(1,1) ; the POPE-tuned variant is a reference ceiling (see Ablation ).
Method
ListOps
Text
Retrieval
Image
Pathfinder
Unreduced
38.10±1.21
65.03±0.66
83.02 ± 0.20
35.94±0.58
49.9±0.1
Performer
38.57±0.37
65.93±0.29
82.2±0.15
43.85±0.13
50.21±0.13
Uniform pooling
29.65±0.25
75.05±0.35
88.12±0.33
42.29±0.65
49.67±0.21
ToMe
35.47±2.36
60.97±0.15
78.85±0.34
36.16±0.19
49.84±0.06
Top- N selection
24.92±9.24
55.68±2.55
74.93±9.13
36.89±0.18
50.78±1.42
TokenLearner
38.78±0.71
65.47±0.26
83.51±0.34
37.34±0.28
49.98±0.19
Table 3: Long Range Arena, end-to-end training, budget 10%. Test accuracy (%), mean ± std over 3 seeds. Unreduced and Performer are full-sequence references. Best result per task in bold ; Pathfinder is at chance level for all methods, so no entry is bolded there.
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
Context
Weight w
Summation domain
Query point
Cell i source
§ A.1 , nodei
r
tokens j
pi
index i
§ A.1 , solve/density
r
tokens j
bℓ
Safeguard A search
§ A.2 , Forward
r˙
tokens j
bℓ
reused from solve
§ A.2 , Reverse
c (below)
barriers ℓ
pj
complement of swapped scan
§ A.3
r⊙x
tokens j
bℓ
reused from solve
Appendix
Table 4: Five applications of the recentered prefix/suffix scans. The reverse row takes the complement of Φc(pj) to obtain the required CDF orientation. Cell indices are determined by a monotone merge search.
experiment
scorer (input → arch)
budget N
merge width α
trained modules
backend
histopathology
feats + xy → MLP(1536+2,128,1)
1024
learned scalar, init 1/N
scorer, width, head
native
LRA
token → linear(dim,1)
10% of M
fixed 1 (rank units)
scorer, backbone, head
native
Cars
token → MLP(384,128,2) →(r,a)
20
learned per-coord. + frame U
scorer, U , width, head
native
arXiv
token → MLP(768, ⋅ ,2) →(r,a)
8 – 64
learned per-coord. + frame U
scorer, U , width, head
native
POPE
frozen VisionSelector scorer
⌈bM⌉
fixed (αr,αa)=(1,1) , index units
none (eval-only)
native
Appendix
Table 5: Configuration recipe per experiment. Routing is LapSum soft top- k at temperature 1 in every row (POPE standardizes the frozen scores per image first); the runtime benchmark instead uses softmax routing, as stated in Appendix B . “Native” is the pure PyTorch implementation; the CUDA-fused backend appears only in the runtime study. Backbones are frozen except on LRA, where the reducer and the backbone are trained jointly from scratch; the row shows the configuration shared by all five LRA tasks (only the input dimensionality varies).
moderate masses
near-hard masses
backend (vs dense-FP64)
small
large
small
large
scan-FP64
4×10−12
3×10−11
3×10−10
5×10−8
dense-FP32
2×10−3
2×10−2
0.16
30
scan-FP32 (eager)
3×10−3
3×10−2
0.37
56
auto-FP32 / compile-FP32
2×10−3
3×10−2
0.16
56
fused-FP64
3×10−13
0.57
2×10−12
3×1014
Appendix
Table 6: Measured backend agreement : maximum relative forward deviation from a dense-FP64 reference on identical inputs. “Small” covers M≤4096 with N≤512 (plus N=M , N>M cases); “large” covers M≥65,536 or N≥1024 at M≥4096 . Gradient deviations follow the same pattern. FP32 rows degrade together with the problem’s conditioning (shared input-precision sensitivity); the fused rows mark that backend’s genuine limits, which the accuracy experiments avoid by using the native implementation.
real slides ( 98 )
tiled ∼106 , N=1024
tiled ∼106 , N=M/16
scan-FP32
8.2×10−4
2.2×10−3
8.3×10−2
fused-FP32
4.6×10−4
5.6×10−2
4.0×10−1
fused-FP64
1.9×10−12
5.6×10−2
4.0×10−1
Appendix
Table 7: Backend agreement on real trained-scorer masses : maximum relative forward deviation from the scan-FP64 reference (which agrees with dense-FP64 to ≤5×10−8 , Table 6 ), over the trained histopathology models’ own (mass, position, width) inputs on all test slides, and over million-token probes tiled from the same profiles. Gradient deviations follow the same pattern. Identical fused-FP32/FP64 values on the tiled probes mark the algorithmic-path limit already documented above.
Figure 5: Matched GH200 forward–backward times for the same LapMerge operator at B=1 , d=512 , FP32. Coordinates are already ordered. Each point is the median of 30 synchronized repeats after five warm-ups. The dense path was only attempted to M=220 . A filled endpoint marks the last success before OOM; a cross marks the last success before a compiler failure. Lines are measured points, not extrapolations.
Figure 6: Incremental peak CUDA allocation for the same matched sweeps. Allocation is measured above the already live input tensors, not as total device occupancy. At fixed N , the dense path also grows linearly with M but has a much larger constant; N=M/16 exposes the different dependence on both sequence dimensions. Endpoint symbols have the meaning of Figure 5 .
N=1024
N=M/16
Implementation
time [ms]
peak [GiB]
time [ms]
peak [GiB]
Explicit dense
128.52
35.03
—
—
Native PyTorch
274.30
22.05
289.54
23.16
torch.compile
214.76
24.05
230.13
25.16
CUDA fused
111.04
2.53
136.09
3.64
Appendix
Table 8: Matched GH200 measurements at M=220 , B=1 , d=512 , FP32. Time is steady-state forward plus backward. Peak is incremental CUDA allocation above live inputs. The explicit dense reference was not swept at N=M/16 ; its omission is a scope limit, not an OOM result.
N=1024
N=M/16
Reducer
time [ms]
peak [GiB]
time [ms]
peak [GiB]
LapMerge (fused)
111.04
2.53
136.09
3.64
TokenLearner (SDPA)
336.16
4.29
7158.16
4.29
PatchMerger (SDPA)
453.95
4.02
14750.05
4.38
Appendix
Table 9: Matched reducer measurements at M=220 on one GH200, B=1 , d=512 , FP32. Times include forward, mean-square loss, and backward; peaks are incremental CUDA allocations above live inputs. Each cell is a median of 30 synchronized repeats after five warm-ups.
M
N
scorer
routing
merge (eager)
merge (fused)
total (fused)
214
1024
0.8
5.7
6.2
6.1
10.4
217
1024
1.3
5.6
33.3
18.5
24.7
217
8192
1.3
5.6
35.8
22.8
26.0
220
1024
7.8
5.3
274.2
111.0
122.5
220
65536
7.8
5.2
291.3
135.6
147.1
223
1024
60.8
7.2
–
898.3
979.8
Appendix
Table 10: End-to-end decomposition of the full LapMerge module (median ms, forward+backward, GH200, B=1 , d=512 , FP32). Components are timed separately; total runs the whole scorer → routing → merge chain with one backward. Dashes: the eager merge at M=223 exceeds memory in this control harness, which keeps extra per-component buffers; the main benchmark covers that point.
PatchMerger
TokenLearner
LapMerge
M
N
published
SDPA
published
SDPA
fused hybrid
65,536
1,024
120.22
28.01
118.52
20.64
10.86
65,536
2,048
123.81
37.45
123.38
22.55
15.10
262,144
1,024
475.24
110.60
473.75
80.46
29.51
262,144
2,048
493.98
151.90
490.19
89.33
29.64
1,048,576
1,024
1932.20
456.79
1918.33
322.97
112.61
Appendix
Table 11: Historical published-versus-SDPA control on one GH200 (ms). All methods use B=1 , d=512 , and FP32. Published PatchMerger and TokenLearner materialize their M×N assignment maps. Their SDPA variants are mathematically equivalent memory-efficient implementations that do not materialize these maps. LapMerge uses its shape-aware dense/fused-CUDA dispatcher. Entries are warmed medians over 20 iterations for the competitors and 100 iterations for LapMerge.
M
d=512
d=1024
d=2048
step ratio ( d=512 )
103
4.8
4.7
7.1
—
104
5.9
6.7
10.7
1.2×
105
22.0
40.4
75.7
3.7×
106
222.5
433.5
oom
10.1×
3⋅106
688.3
⋅
⋅
3.1×
107
oom
⋅
⋅
—
Appendix
Table 12: Independent H100 ( 80 GB) dimensional stress test of the complete reducer, forward plus backward, at N=0.1M . Times are warmed medians in milliseconds from a torch.compile run. This is not the matched GH200 protocol; the final column shows the step ratio at d=512 .
Regime
M
N
Forward [ms]
Forward + backward [ms]
Peak [GiB]
N=1024
222
210
169.09
439.28
10.07
N=1024
223
210
360.85
899.09
20.13
N=1024
224
210
807.38
1878.61
40.24
N=M/16
222
218
150.64
531.17
14.57
N=M/16
223
219
299.83
1067.26
29.14
N=M/16
224
220
—
OOM
—
Appendix
Table 13: Large-shape GH200 probes for the CUDA-fused LapMerge dispatcher, with ordered inputs, B=1 , d=512 , and FP32. The failed ratio cell has no reported latency or allocation; no forward-only measurement is inferred. Peak allocation is incremental above live inputs.
LapMerge one sweep (soft)
operator-free (hard)
Top- N selection
N=10
74.91±0.60
59.51±0.66
71.47±0.31
N=20†
75.10±0.28
67.53±0.26
74.20±0.44
N=40
76.52±0.40
73.25±0.27
77.28±0.13
Appendix
Table 14: Budget dose-response on Stanford Cars (top-1, mean ± sd over three shared seeds, original protocol). † The N=20 row is the capacity-matched frozen-frame configuration of Table 15 ; the full one-sweep single run reaches 75.46 there.
Figure 7: Learned routing mass on Stanford Cars. Each pair shows the input image (left) and the per-token routing mass ri of the trained reducer as a heatmap (right; normalized per image). Mass concentrates on discriminative car parts – grille, badge, and headlight region – while background patches receive near-zero mass. Left pair: the standard 196 -token single-crop configuration of Table 15 ; right pair: a 2×2 multi-crop variant ( 784 tokens) showing the same behaviour on a finer grid.
Reducer
Tokens
Params
Reducer cost
Soft
Hard
to head
(%)
(%)
Dense transformer, no reduction
196
—
none
76.71
—
Top- N selection
20
52 k
O(Md) + top- N
74.10
74.10†
TokenLearner
20
53 k
O(MNd)
73.60
n/a
PatchMerger
20
8 k
O(MNd)
71.83
n/a
ToMe post-encoder merge
20
0
O(M2d)
69.87
n/a
Appendix
Table 15: Matched pre-head reducers on Stanford Cars. “Reducer params” excludes the 2.44 M shared by every row; “reducer cost” is analytic. Soft runs the reducer as trained; hard keeps the N highest-scoring raw tokens and removes it entirely (a head-dependence test, not a hard-trained model). † Selection coincides on both by construction; “n/a” marks poolers, which expose no operator-free path. The ToMe row reaches 20 tokens by iterating bipartite matching roughly four times with no transformer layer in between, well outside its published schedule; read it as training-free similarity merging pushed out of its design regime. The last block is the capacity-matched control (mean ± sd over three shared seeds): the read-out frame U is frozen at identity, so the reducer matches selection’s ≈52 k parameters while sharing its scorer – the merge advantage survives equalized capacity. Protocol in Appendix C.1 .
Reducer (arXiv, frozen Longformer)
Reducer
Output budget N for M=1024
every arm trained at its own budget
params
8
16
32
64
Dense (no reduction)
—
87.6
TokenLearner ( Ryoo et al., 2021 )
102 k
88.28
87.88
88.12
87.68
PatchMerger ( Renggli et al., 2022 )
14 k
88.00
88.28
87.08
87.48
Top- N selection
100 k
81.48
81.44
82.12
81.80
LapMerge (scan coordinate)
0.75 M
87.28
87.68
87.72
87.68
Appendix
Table 16: Budget-matched training on arXiv (single seed, soft deploy): every entry is trained at the budget it is evaluated at, which removes the extrapolation penalty the poolers would otherwise pay and is conservative for us. The scan coordinate is the parameterization of § 3 and is what the text calls LapMerge; the published-coordinate row is an ablation. For scale only: three independent trainings of the published-coordinate arm give 86.24/85.84/86.76 at N=8 , i.e. σ≈0.46 for that one arm; all other entries are single runs. Against LapMerge the poolers lead in five of eight cells, trail in two and tie in one, with a largest gap of 1.0 point.
reducer
raster
col-major
Hilbert
rnd (fixed)
rnd (per img)
LapMerge
48.01±1.34
45.42±0.14
48.96±0.26
32.48±0.89
35.36±0.80
uniform pooling
45.15±0.27
43.67±0.45
46.71±1.14
29.46±0.82
28.11±0.28
ToMe
37.16±0.24
37.33±0.45
37.67±0.16
37.35±0.33
42.85±0.31
TokenLearner
37.08±0.34
—
—
37.08±0.34
—
PatchMerger
26.94±0.34
—
—
26.94±0.34
—
learned top- N selection
37.20±0.88
—
—
37.20±0.87
—
Appendix
Table 17: Permuting the scan, with everything else held fixed. LRA-Image at a 10% budget, mean ± sd over three seeds, sharp width. Down the top block: a locality-preserving order is a locality-preserving order – raster, its transpose and a Hilbert curve agree to within 3.5 points – while a random permutation costs the merge 15.5 , dropping it into the band of the reducers below it. Across the bottom block: those reducers cannot see the scan coordinate, and they return their own numbers to two decimals in both mean and spread, which rules out the permutation leaking through the data, the positions, or the harness. A per-image permutation is gentler than a fixed one for the merge ( −12.7 against −15.5 ) and actually helps ToMe ( +5.7 ): resampling the order every step acts as augmentation, which is why the fixed draw – one the Transformer is free to learn – is the cleaner intervention.
sweep
value
raster
permuted (fixed)
gap
width α (budget 10% )
0.5
48.0±1.2
33.2±0.7
14.8
1
48.0±1.1
33.0±0.9
15.0
2
47.2±0.5
32.3±0.9
14.9
5
44.4±0.1
31.5±1.0
12.9
10=M/N
42.6±0.2
31.4±1.2
11.2
30
41.5±0.2
31.4±0.8
10.2
Appendix
Table 18: How the order effect responds to width and to budget. LRA-Image, LapMerge, mean ± sd over three seeds; every row is the raster scan against one random permutation of the scan coordinate, fixed for the run, with everything else as in Table 17 ; the gap is the difference of means. Across the width it is not flat: sharp kernels ( α≤2 ) hold the full gap, which shrinks as α grows toward M , where the raster arm itself falls to the order-free level (the permuted arm stays between 31 and 33 throughout) – exactly what a positional mechanism should do. The 10%,α=1 configuration appears in both blocks with the same three seeds and returns 48.0±1.1 and 48.8±0.8 (permuted 33.0 and 32.8 ): the difference is run-to-run nondeterminism of the GPU scatter, about the size of the seed spread, so differences under a point between rows are not to be read.
per-anchor softmax width τ
budget
0.05
0.3
0.7
1.5
3.0
10%
79.67
79.75
80.70
80.58
78.61
30%
85.11
85.32
85.26
84.41
80.92
Appendix
Table 19: Per-anchor smoothing control (Recipe-B) : POPE F1 in the shared harness. τ=0.05 recovers hard selection exactly; the best column trails the default score-routed LapMerge of Table 20 at both budgets.
read-out width αa ( αr=1 )
routing width αr ( αa=1 )
budget
select
0.3
0.7
1.0⋆
1.5
3.0
0.5
2.0
10%
79.67
82.01
81.67
81.31
80.69
79.56
80.97
80.70
30%
85.11
85.87
85.87
85.50
84.73
82.92
85.19
84.97
Appendix
Table 20: Width sweep of score-routed LapMerge (POPE F1, frozen Qwen2.5-VL-7B; the same shared harness and metric as Table 2 ). “select” reruns the released selection pipeline under identical conditions and reproduces its published F1 ( 79.73 / 85.05 ) to within 0.06 . The starred column is the dataset-agnostic default used in Table 2 . Decoding is greedy, so each cell is deterministic.
Domain
Dataset
#Classes
Train
Test
Histopathology
TCGA-CRC (COAD+READ)
2
365 slides †
TCGA-STAD
2
303 slides †
LRA
ListOps
10
96,000
2,000
Text
2
25,000
25,000
Retrieval
2
147,086
17,437
Image
10
45,000
10,000
Appendix
Table 21: Summary of the datasets used in our experiments.
Figure 8: Per-patch scorer weight of LapMerge for six TCGA-CRC test slides (three MSI-H, three MSS), one column per slide: true-color thumbnail with true label and predicted p(MSI-H) (top) vs. the desaturated scorer-weight overlay (bottom; opacity tied to weight, so only the regions the scorer relies on stand out). The banded columns (TCGA-D5-6540, TCGA-F4-6570) are gaps in patch coverage between several narrow tissue fragments on those slides – confirmed from the raw patch coordinates – not a scorer artefact.
Token merging combines redundant transformer tokens to reduce sequence length, attention costs, and inference time while preserving task-relevant information.