See What Matters: Differentiable Grid Sample Pruning for Generalizable Vision-Language-Action Model
Abstract
Vision-Language-Action (VLA) models have shown remarkable promise in robotics manipulation, yet their high computational cost hinders real-time deployment. Existing token pruning methods suffer from a fundamental trade-off: aggressive compression using pruning inevitably discards critical geometric details like contact points, leading to severe performance degradation. This forces a compromise, limiting the achievable compression rate and thus the potential speedup. We argue that breaking this trade-off requires rethinking compression as a geometry-aware, continuous token resampling in the vision encoder. To this end, we propose the Differentiable Grid Sampler (GridS), a plug-and-play module that performs task-aware, continuous resampling of visual tokens in VLA. By adaptively predicting a minimal set of salient coordinates and extracting features via differentiable interpolation, GridS preserves essential spatial information while achieving drastic compression (with fewer than 10\% original visual tokens). Experiments on both LIBERO benchmark and a real robotic platform demonstrate that validating the lowest feasible visual token count reported to date, GridS achieves a 76\% reduction in FLOPs with no degradation in the success rate. The code is available at https://github.com/Fediory/Grid-Sampler.
Lay Summary
Robots rely on AI to see and interact with the world, but processing every detail in an image requires massive computational power. To speed things up, previous methods simply deleted rigid blocks of the image, which often erased crucial details like the exact spot needed to grasp an object. We introduce the Differentiable Grid Sampler (GridS), a tool that changes how robot AI processes vision. Instead of blindly dropping image patches, GridS scans the scene and smoothly zooms in to sample only the related interactive areas. This focused approach allows robots to understand their environment using less than 10% of the original image data, cutting computational costs by 76% without sacrificing performance. By forcing the AI to ignore distracting background clutter, GridS also makes robots nearly 29% more successful in new and unfamiliar situations. This paves the way for faster, smarter robots that can run on everyday computer hardware.