Visual token routing cuts computation for real-time AI agents
Visual Token Routing: Learning to keep tokens that matter
Computer Vision and Pattern RecognitionMachine Learning
Summary
Many AI agents use thousands of small image pieces, called visual tokens, to understand their surroundings, but most tokens are either not helpful or repeat the same information. The authors introduce Visual Token Routing, a method that learns to keep only the important tokens for predicting the next action, making the system faster and less costly. Their approach uses a lightweight decision process to pick relevant tokens from current and past images. With minimal training, this method matches full performance while using far fewer tokens, speeding up processing for tasks like video games and long tasks requiring many steps.
What this means in practice
- •For video game developers: Implement AI agents that react quickly using fewer image details to save computing power in real-time gameplay.
- •For mobile app developers: Create efficient AI that processes fewer visual tokens to reduce battery and CPU use during long user interaction sessions.
Authors
Pranjal Aggarwal, Sean Welleck
Abstract
Computer-use agents (CUAs) have recently achieved superhuman performance on several long-horizon benchmarks. However, their computational cost limits their use in long-horizon tasks and real-time domains such as video games. A key bottleneck is representing each screenshot as thousands of visual tokens, which accumulate into long contexts as agents reason over extended interaction histories. We observe that most of these visual tokens are either semantically irrelevant to the agent's next action or temporally redundant across screenshots. Building on this observation, we introduce Visual Token Routing (VTR), a joint training framework for learning sparse visual representations that preserve the information needed for action prediction. In VTR, we train the vision-language model together with a lightweight spatio-temporal router that examines each image region in the current and previous screenshots and decides whether its visual token enters the language model. We optimize the router using an adaptive loss to keep the number of visual tokens as small as possible while minimizing performance degradation. With just 2,000 training steps, VTR matches the performance of the full-resolution model on OSWorld, CUA-World, and video game action prediction, while using 2.5-5.8x fewer visual tokens. These savings reduce inference cost by up to 3.3x and speed up prefill by up to 2.9x. VTR also outperforms lowering the image resolution and other baselines for dropping tokens, by up to 12.8 points on OSWorld. Overall, VTR enables computer-use agents to process only the visual tokens they need, allowing them to be used in long-horizon tasks and real-time domains at a fraction of the computational cost. Code and models are available at https://pranjal2041.github.io/vtr