Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper addresses the computational bottleneck of token-based latent planning in visual world models by introducing COSTGRAD, a training-free method that prunes spatial tokens based on the gradient norm of the planning cost. By prioritizing tokens crucial for downstream control rather than just prediction, COSTGRAD achieves significant speedups—up to $5\times$ with combined sparsity and reduced search—while matching or exceeding full-token planning performance on most continuous-control benchmarks. The effectiveness of COSTGRAD is shown to be architecture-dependent, highlighting the importance of selector-architecture compatibility for efficient sparse world-model planning.
Planning in visual world models can be accelerated by $5\times$ by pruning $50\%$ of spatial tokens based on their influence on the planning cost gradient, without sacrificing performance on most control tasks.
Token-based world models enable fine-grained latent planning, but repeatedly processing large spatial token grids makes action search expensive. We introduce COSTGRAD, a training-free, goal-conditioned selector that ranks spatial tokens by the gradient norm of the planning cost with respect to each input token. By deriving importance from the downstream control objective, COSTGRAD targets tokens that matter for planning rather than merely for prediction. On AdaLN-conditioned predictors at $50\%$ sparsity, COSTGRAD matches or exceeds full-token planning on three of four continuous-control benchmarks, while giving a measured $2.6\times$ wall-clock speedup per environment planning step. Combining token sparsity with reduced CEM search increases this to a $\sim 5\times$ total speedup while still exceeding the full-token baseline. We also identify an architecture-dependent failure mode: in a matched AdaLN-vs-concat comparison, concat maintains comparable full-token performance but pure COSTGRAD loses its advantage over random selection. This difference tracks action-pathway drift: gradient-selected removal produces less drift than random removal on AdaLN, but more on concat. These results highlight selector-architecture compatibility as a design axis for sparse world-model planning. Project page and demos: https://ycxuyingchen.github.io/costgrad/