From Sequential Refinement to One-Shot Global Selection: Monotone Adaptive Tokenization

Adaptive tokenization reduces the cost of transformers on high-resolution and high-dimensional data by allocating fine-grained tokens to informative regions and coarse-grained tokens elsewhere, typically through sequential Top-$1$ refinement over a quadtree or octree. This refinement can itself become a bottleneck: when run on the GPU, it accounts for 53–80% of each training step on our 2D benchmarks, exceeding the runtime of the model itself, because its $K$ dependent rounds must run sequentially, each dominated by kernel-launch overhead. We observe that when importance scores are monotone along the hierarchy, with every ancestor scoring no lower than its descendants, the hierarchy satisfies the heap order property; as in classical heap selection, $K$ sequential Top-$1$ refinements then select exactly the $K$ highest-scoring regions. Monotone Adaptive Tokenization (MAT) therefore replaces sequential refinement with one global Top-$K$ selection while preserving the exact Top-$1$ allocation under monotone scores; non-monotone metrics are made monotone through bottom-up sum or max reductions. Across high-resolution 2D segmentation, large-scale 3D point clouds, and multichannel weather data, MAT achieves up to $3.73\times$ and $2.56\times$ higher training and inference throughput than sequential refinement on 2D data, $1.18\times$ and $2.06\times$ over a dense-voxel baseline in 3D with substantially lower memory usage, and $1.19\times$ and $1.66\times$ over sequential refinement on weather data, while preserving downstream quality. Across the tested token budgets and attention backends, the speedups reach up to $6.3\times$ in training and $19.2\times$ in inference.Our code is publicly available at https://anonymous.4open.science/r/Monotone_Adaptive_Tokenization-76FC.