TEAM-VLA: Efficient Token Compression for VLA
- The paper shows that TEAM-VLA achieves a balanced trade-off between efficiency and effectiveness by compressing tokens in a training-free manner, yielding a 1.5× speedup while preserving task success rate.
- TEAM-VLA employs a two-stage token manipulation process, using early expansion-based pruning and mid-layer action-guided merging to reduce the quadratic attention cost of Transformers.
- The framework outperforms prior training-free token compression methods on LIBERO benchmarks by leveraging lightweight, inference-time operations without retraining or buffering.
Searching arXiv for the specified paper to ground the article in the cited source. Token Expand-and-Merge-VLA (TEAM-VLA) is a training-free token compression framework for Vision-Language-Action (VLA) models that accelerates inference while preserving task performance. It is designed for pretrained VLA systems in which a policy maps current sensory observations, a language instruction, and robot state to a low-level action, and it targets the computational bottleneck created by the quadratic attention cost of Transformer backbones. TEAM-VLA operates at inference time, without fine-tuning, retraining, or parameter updates, by coupling early token expansion and pruning with mid-layer action-guided token merging in a single feed-forward pass (Ye et al., 10 Dec 2025).
1. Problem setting and design objective
A VLA model maps current sensory observations and a language instruction to a low-level robotic action. The standard architecture has three main parts: a sensory encoder, a language backbone, and an action head. The sensory encoder processes observations into image tokens , converts task instruction into language tokens , and encodes robot proprioceptive state into . The language backbone takes a concatenation of image tokens, language tokens, and state tokens, and performs multi-layer self-attention and cross-modal reasoning. The action head reads the output hidden state at special action tokens and produces the action . Formally, the VLA policy is
The motivation for TEAM-VLA is the computational and latency burden of large VLA backbones. Transformers incur 0 attention cost in the token length 1. In VLAs based on large VLM backbones, each image produces hundreds of visual patch tokens, the language backbone is deep, and real-time closed-loop control at 10–30 Hz or higher requires per-step latency well below approximately 100 ms. Vanilla OpenVLA-OFT on LIBERO tasks runs at 109 ms per step, which is borderline or insufficient for many on-robot scenarios (Ye et al., 10 Dec 2025).
The framework is explicitly motivated against three classes of prior token compression approaches. Cross-attention or similarity-based pruning can produce extremely sparse and noisy maps in early layers because robot instructions are short and sparse. Temporal or streaming methods such as VLA-Cache, SpecPrune-VLA, and TTF-VLA require buffering previous frames 2, assume temporal continuity, introduce extra memory and complexity, and can fail with large viewpoint changes or partial observability. Learned or prior-dependent pruning methods, including differentiable pruning, learned queries, or adapters, require additional training and sometimes architecture changes. TEAM-VLA therefore targets a training-free, observation-only, buffer-free token compression method that works directly on the current frame and can be plugged into a pretrained VLA.
2. Architectural placement and inference pipeline
TEAM-VLA is applied at two points in the model. First, before the language backbone, it performs token expansion and pruning on visual tokens after the sensory encoder. Second, at an intermediate backbone layer, it performs action-guided token merging on deeper tokens.
Its pipeline consists of five stages. Sensory encoding maps 3, 4, and 5 to 6, 7, and 8. Token expansion and pruning then compute cosine similarity between each image patch token and language tokens, choose for each language token its most similar image token as a seed, build a sparse binary mask over the image grid, convolve that mask to obtain a density map, and apply deterministic and stochastic spatial expansion to obtain a dense foreground mask. A small subset of background tokens is then retained by context sampling, and only the foreground and sampled-context tokens are passed into the language backbone. The early backbone layers then process the reduced visual token set together with language and state tokens.
At a selected middle Transformer layer, TEAM-VLA computes similarity between visual tokens and both text tokens and action tokens. It selects the top-9 most relevant visual tokens as a source set 0, defines the remaining visual tokens as a target set 1, computes a similarity matrix and soft assignment from 2 to 3, and merges 4 into 5 through weighted averaging. The merged token set 6 replaces the original visual tokens, and the deeper layers continue from that compact representation. The action head finally decodes actions from the final hidden states as in the original VLA (Ye et al., 10 Dec 2025).
This schedule couples expansion-based pruning and action-guided merging in a single feed-forward pass. The paper states that the result is a balanced trade-off between efficiency and effectiveness, with large token reductions and approximately 1.5× speedup, while maintaining essentially unchanged success rate.
3. Token expansion and context sampling
The token expansion mechanism begins with similarity sampling and sparse foreground seeds. Let 7 be the language embeddings and 8 the visual patch embeddings after the image encoder. The method uses cosine similarity:
9
For each language token, the image token index with maximum cosine similarity is selected as a foreground anchor. The paper notes that instructions often contain many non-object words and that even object phrases may map to only a few patch positions, so the resulting highlighted set is sparse.
Visual tokens come from a fixed grid of patches. With patch size 0 and image size 1, the number of patches per dimension is denoted by
2
with the clarification that the notation in the PDF is slightly mangled, but conceptually corresponds to the patch grid. The similarity-selected anchors are converted into a binary mask
3
where 4 if patch 5 is a seed. To measure local foreground density, TEAM-VLA convolves this mask with a uniform kernel:
6
where
7
with zero padding. The density map 8 records how many seeds lie in each 9 neighborhood.
Deterministic expansion is then applied in dense regions. The dense neighborhoods are defined as
0
where 1 is a density threshold. For each 2, all originally unrelated patches in its local 3 neighborhood are set to 1 in the mask. Sparse neighborhoods are defined as
4
For each 5, TEAM-VLA randomly flips one 0 entry in its neighborhood to 1. The final foreground mask therefore includes seed positions, dense regions around clusters of seeds, and scattered expansions in sparse regions. The convolution with an all-ones kernel can be implemented efficiently, and the entire expansion step costs only 1–2 ms.
The context sampling mechanism complements the foreground mask by retaining a small fraction of background tokens as global context. Let 6 be the context sampling ratio. Over the entire visual token sequence, including both views, TEAM-VLA performs interval-based random sampling of background tokens such that the approximate fraction retained is 7. This preserves coarse scene layout and background cues necessary for spatial understanding. The paper’s interpretation is that expansion reconstructs coherent spatial regions such as object bodies and robot arm segments, while context sampling preserves coarse layout cues. A plausible implication is that the two operations are intended to prevent aggressive early pruning from removing either task-relevant local structure or scene-level context.
4. Action-guided token merging
The merging stage is motivated by the fact that, after early pruning and expansion, the remaining tokens still include both strongly object-relevant foreground tokens and subtle but important action-related cues. TEAM-VLA therefore uses action-guided merging rather than additional deep-layer pruning.
Let 8 be the set of visual tokens at a chosen intermediate layer in the language backbone. TEAM-VLA partitions 9 into a source set 0 and a target set 1. The source set contains the top-2 tokens most relevant to both text and action tokens; the target set contains the remaining visual tokens. These sets satisfy 3 and 4. Selection of 5 is based on similarity between each visual token and a combination of text and action tokens, aggregated in the style of the cosine similarity computation used earlier.
The actual merge is performed through soft bipartite matching. Let
- 6 be the matrix of source visual token features,
- 7 be the matrix of target visual token features,
- 8 be the hidden dimension.
The normalized similarity matrix is
9
with the text stating that this defines a soft matching matrix 0. Applying softmax along the source dimension yields assignment weights:
1
and
2
Each row of 3 gives a probability distribution over source tokens for one target token. Target representations are aggregated to source tokens by
4
and the total soft assignment weight per source token is
5
The updated source tokens are then computed as
6
with the paper noting that the printed formula contains minor typos and that the intended operation is a normalized sum of the original source and aggregated target features.
After merging, only 7 remains as the visual representation at that layer; all target tokens 8 are removed. The stated effect is a reduction in token count from 9 to 0, while preserving semantic information from removed tokens by embedding it into action-critical locations. The paper explicitly contrasts this with pruning, which discards tokens entirely.
5. Integration with OpenVLA-OFT and computational profile
TEAM-VLA is implemented on top of OpenVLA-OFT, described in the paper as an open-source VLA with strong performance of approximately 96–97% success on LIBERO and a large multimodal transformer as language backbone. TEAM-VLA uses the LightVLA codebase and wraps OpenVLA-OFT at inference without modifying any weights (Ye et al., 10 Dec 2025).
The required architectural access points are the visual patch embeddings after the image encoder, the language embeddings, and the action token representations at a chosen mid layer. The external operations are cosine and RMSNorm-based similarities, lightweight 2D convolution over binary masks, token selection and masking, and simple weighted sums. No new layers or trainable modules are added; the changes are procedural manipulations of tokens outside or around Transformer layers.
The inference schedule places expansion and context sampling after the sensory encoder and places merging at a selected layer 1, empirically best around layer 16 in OpenVLA-OFT. This placement is tied to a trade-off between compression timing and representation quality. The paper reports that early expansion-based pruning gives a large quadratic FLOP reduction, while merging gives further reduction at deeper layers when features are already well aligned.
In complexity terms, vanilla VLA attention scales as 2 per layer, where 3 is the total number of tokens. TEAM-VLA reduces 4 first through early visual token pruning and later through mid-layer merging. From Table 1, OpenVLA-OFT uses 100% FLOPs with latency 109 ms, whereas TEAM-VLA uses 39% FLOPs of baseline and reaches 72.1 ms latency. The paper states that this corresponds to approximately 61% FLOP reduction. It does not provide additional asymptotic formulas beyond standard 5, but it explicitly frames the cost reduction as moving from the original token count to reduced counts after expansion and then to 6 tokens after merging.
6. Empirical evaluation, ablations, and limitations
The empirical evaluation uses the LIBERO benchmark with four suites: LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-Long. The metrics are Success Rate (SR), CUDA Latency, and FLOPs. The compared methods are OpenVLA-OFT, SparseVLM, VLA-Cache, EfficientVLA, SpecPrune-VLA, and TEAM-VLA.
In the main quantitative results, OpenVLA-OFT achieves Spatial 97.6 / 109, Object 96.5 / 109, Goal 97.9 / 109, Long 94.5 / 109, with average SR 96.6, average latency 109 ms, and FLOPs 100%. TEAM-VLA achieves Spatial 99.2 / 68.1, Object 96.5 / 74.7, Goal 97.0 / 72.9, and Long 93.8 / 72.8, with average SR 96.6, average latency 72.1 ms, and FLOPs 39%. SparseVLM reports average SR 95.5, latency 85.3 ms, and FLOPs 77%. VLA-Cache reports average SR 96.9, latency 101.7 ms, and FLOPs 83%. EfficientVLA reports average SR 88.9, latency 70.6 ms, and FLOPs 35%. SpecPrune-VLA reports average SR 96.5, latency 75.1 ms, and FLOPs 43%. TEAM-VLA is marked as LM-pre: Yes and Buffer-free: Yes. The paper concludes from these numbers that, among training-free, buffer-free methods, TEAM-VLA offers the best accuracy-latency trade-off (Ye et al., 10 Dec 2025).
A token-count comparison in Fig. 5 reports that TEAM-VLA retains substantially fewer tokens than other training-free approaches while achieving higher performance, and that it prunes on average approximately 432 tokens relative to the baseline while keeping the same success rate.
The ablation studies clarify the contribution of individual components. On LIBERO-Long, the baseline is SR 94.5 and latency 109 ms. Context sampling only, with no expansion, yields SR 80.1, 88.6, 91.7, and 92.1 for 7 and 8, with latencies 67.4, 71.3, 71.9, and 76.9 ms, and pruned-token counts 461, 384, 358, and 307. The “Most Similar Token” variant, with no expansion, yields SR 89.5, latency 71.9, and pruned 358, showing that simply keeping similarity-selected tokens without expansion is suboptimal. Adding token expansion with 9 and 0 gives SR 93.8, latency 76.8, and pruned 302; using 1 gives SR 94.1, latency 82.5, and pruned 249. Adding action-guided merging with expansion and 2 gives, for top-50, SR 88.3 and latency 71.6; for top-80, SR 93.8 and latency 72.8; for top-110, SR 92.8 and latency 73.7; and for top-130, SR 93.8 and latency 74.1. The chosen configuration is top-80.
Layer ablations compare merging with pruning at different layers. For token merging, layer 1 gives SR 72.6 and latency 68.7, layer 4 gives 88.0 and 70.2, layer 8 gives 90.4 and 71.5, and layer 16 gives 93.8 and 72.8. For token pruning at the same layers, the corresponding SR values are 74.2, 90.2, 92.2, and 92.1, with the same latencies. The paper draws two explicit conclusions: merging is better than pruning at deeper layers, and doing heavy compression in very early layers dramatically hurts performance.
The threshold 3 for expansion is also varied. For 4, average SR is 97.2%, average latency is 86.0 ms, and average patches are 342.3. For 5, average SR is 96.6%, average latency is 75.3 ms, and average patches are 196.0. For 6, average SR is 94.6%, average latency is 71.8 ms, and average patches are 151.0. The default is 7.
The implementation details specify that TEAM-VLA is applied to both agent view and wrist view, with batched convolution across views. The reported hyperparameters vary by LIBERO suite: LIBERO-Spatial uses 8 and source merging 9; LIBERO-Goal uses 0 and 1; LIBERO-Long uses 2 and 3; and LIBERO-Object uses 4 and 5. The main experiments use expansion kernel size 6, density threshold 7, and merging around layer 16. Hardware is a single NVIDIA A100 40GB, and the code is released at the public GitHub repository listed in the paper.
The paper notes several limitations. TEAM-VLA is evaluated on OpenVLA-OFT and LIBERO, so behavior on other backbones or substantially different tasks may differ. The method depends on hand-designed heuristics such as 8, 9, 00, and 01, which are manually tuned. It currently uses cosine and RMSNorm-based similarities, and it merges at a fixed middle layer. The paper identifies future directions including learning-based compression that preserves the training-free or low-training-cost spirit, adaptive layer selection and token budgets, and more robust action-aware relevance scoring. This suggests that TEAM-VLA is best understood as a procedural compression framework whose current effectiveness is demonstrated on a specific VLA architecture and benchmark, rather than as an architecture-independent guarantee.