PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse Planning in Visual World Models via Cost Gradients

Yingchen Xu, Edward Grefenstette

arXiv:2610.10274v1cs.LG

TL;DR

Token-based world-model planning repeatedly processes large token grids, and prediction-oriented signals may not identify tokens that matter for control. COSTGRAD selects tokens using planning-cost gradients, then plans on that sparse subset. It matches or exceeds full-token performance on three of four benchmarks, while the results also reveal architecture-dependent differences in selector performance.

  • Problem

    CEM-based planning repeatedly processes large spatial token grids, while prediction error or attention may not reflect which tokens matter for the downstream control objective.

    Source: §1 · p. 1

  • Method

    COSTGRAD uses the planning objective to score goal-conditioned token importance, then performs action search using the selected tokens with a fixed pretrained predictor.

    Source: §1 · p. 2, §2.2 · p. 3

  • Results

    At K=128, COSTGRAD is within 2.0 pp of Full on PointMaze and exceeds Full on the other three benchmarks; sparsity plus reduced CEM yields a 5.2× speedup over Full.

    Source: §3.1 · pp. 4–5

  • Takeaways & Limitations

    Token-selection performance should be assessed jointly with the predictor at the intended token budget.

    Source: §5 · p. 10

    Evaluation covers four simulated environments, one encoder family, and one planner; the selector–architecture finding is tested only for AdaLN and channel-wise concat predictors.

    Source: §5 · p. 10

Abstract

from arXiv · show

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/

1 Introduction

COSTGRAD selects visual tokens using the planning objective because prediction-oriented signals may not identify what matters for control. The introduction reports sparse-planning gains on AdaLN predictors and a failure mode tied to action-conditioning architecture.

Motivation

  • Planning with token-based world models is expensive because CEM repeatedly processes large spatial grids, much of which may describe task-irrelevant background.At 16 × 16 tokens and horizon 6, one planning step performs approximately 10^7 token-level forward computations.Source: §1 · p. 1
  • Attention and prediction error are imperfect proxies for control relevance because prediction quality and downstream planning objectives can differ.An object important for control may be easy to predict because its motion follows directly from the action.Source: §1 · p. 1

COSTGRAD

  • COSTGRAD ranks tokens by the gradient norm of planning cost with respect to each input token, using the current goal to select what matters for control.This is a training-free selector for a fixed pretrained predictor.Source: §1 · pp. 1–2

Contributions

  • At 50% sparsity, COSTGRAD matches or exceeds full-token planning on three of four benchmarks and gives a 2.6× measured speedup; reduced CEM raises total speedup to ∼5× while still exceeding Full.The architecture comparison also finds that pure COSTGRAD loses its advantage over random selection on concat.Source: §1 · p. 2

Architecture-dependent failure mode

  • On matched predictors, COSTGRAD succeeds on AdaLN but loses its advantage on concat, tracking less action-pathway drift than random removal on AdaLN and more on concat.The predictors have comparable full-token planning performance.Source: §1 · p. 2

2 Setup and Method

The setup uses a frozen visual world model and CEM to plan over encoded token grids. COSTGRAD probes the planning objective to select a fixed sparse subset, which is then reused during action search.

2.1 Token-based world model planning

  • The frozen DINOv2 encoder maps each 224 × 224 image to 256 tokens, and the AdaLN-Zero predictor forecasts the next token grid from context and action.CEM searches 300 action sequences over horizon 6 for 30 iterations, requiring about 1.4 × 10^7 token-level forward computations per planning step.Source: §2.1 · p. 2

2.2 COSTGRAD: select tokens, then plan

  • CEM rolls out candidate actions over the full horizon using only selected tokens, evaluates goal distance at those positions, and reselects tokens at the next observation.The probe selects tokens only; CEM samples and optimizes its own actions, and predictor weights remain fixed.Source: §2.2 · p. 3, Figure 1 · p. 3
  • COSTGRAD runs a full-token, zero-action probe and ranks tokens by the norm of planning-cost gradients, retaining the top K as a fixed subset for that planning step.The gradient norm measures each input token’s local effect on planning cost, including its influence on predictions elsewhere.Source: §2.2 · pp. 3–4, Figure 1 · p. 3

2.3 Baselines

  • The comparison includes Full and four token-reduction baselines; each reduced method retains 128 of the predictor’s 256 tokens.Baselines include attention- and prediction-gradient ranking, random selection, and token merging.Source: §2.3 · p. 4

3 Results

On AdaLN-Zero, COSTGRAD preserves or improves full-token planning on most benchmarks while reducing computation. Its selections align with task-relevant regions, but the sparse-planning advantage depends on predictor architecture and tracks action-pathway drift.

3.1 Sparse planning performance and speed on AdaLN-Zero

  • At K=128 (50% tokens), COSTGRAD preserves PointMaze full-token performance (−2.0 pp) and exceeds Full on Wall (+7.3 pp), PushT (+2.8 pp), and MetaWorld (+6.9 pp).Table 1 reports the three-evaluation-seed mean; Full uses 256 tokens and sparse methods use 128.Source: §3.1 · p. 4, Table 1 · p. 5
  • On released JEPA-WMs checkpoints, COSTGRAD preserves full-token performance across all four environments using half the tokens and outperforms Random by 11.3–46.3 pp.The result is reported without retraining.Source: §3.1 · p. 4
  • At 25% tokens, COSTGRAD remains near Full on Wall, PushT, and MetaWorld, while all sparse methods degrade sharply below roughly 25%.The gap to PREDATTN and Random widens as the token budget shrinks.Source: §3.1 · p. 5, Figure 2 · p. 5
  • Combining 50% tokens with 15 CEM iterations reaches 92.0% success at ∼17 s per planning step, a 5.2× speedup over Full (84.7% at 88.6 s).With the standard CEM budget, retaining 50% of tokens gives a 2.6× speedup; COSTGRAD configurations trace the measured non-dominated speed–success frontier.Source: §3.1 · p. 5, Figure 3 · p. 6

3.2 Task relevance and stability of planning-gradient selection

  • At 50% tokens, COSTGRAD concentrates selections on task-relevant regions—including maze corridors, the agent and wall, the pusher and block, and the arm and object—while dropping much background.The same sparse token sets yield qualitatively faithful one-step predictions.Source: §3.2 · p. 6, Figure 4 · p. 7
  • On PushT, many task-critical pusher and T-block tokens are easy to predict yet receive high COSTGRAD scores; such easy-but-critical tokens occur at ∼2× chance density.This illustrates why prediction error may downweight objects that matter for control.Source: §3.2 · p. 6
  • Across equivalent Wall predictors, COSTGRAD selections overlap at IoU > 0.82, versus 0.44 ± 0.02 for attention selections; COSTGRAD rankings also remain stable across 10 probe actions.The reported probe-action ranking correlations are ρ = 0.918 on AdaLN-Zero and 0.958 on matched-concat.Source: §3.2 · pp. 6–7

3.3 Sparse planning depends on selector–architecture compatibility

  • Despite comparable full-token Wall success (85.1% for concat vs. 84.7% for AdaLN-Zero), COSTGRAD beats Full by 7.3 pp on AdaLN-Zero but falls 5.1 pp below Full on concat.Across Wall and PointMaze, COSTGRAD has no consistent advantage over Random on concat.Source: §3.3 · p. 7, Table 2 · p. 8

4 Related Work

Related work distinguishes spatial-token planning from compact-latent approaches and inference-time token reduction for single-pass perception. COSTGRAD instead uses current-goal planning-cost gradients in a fixed predictor, while related studies motivate assessing selection under downstream use and action-conditioning choices.

Related approaches

  • COSTGRAD selects goal-relevant tokens for repeated action-conditioned rollouts in a fixed predictor, complementing sparse training and lower-dimensional latent-planning approaches.Sparse Imagination combines training-time randomized grouped attention with random planning-time selection; compact-latent methods optimize in lower-dimensional spaces.Source: §4 · p. 9

Token selection and importance estimation

  • Unlike perception-oriented token reduction, sparse visual planning must preserve tokens useful across many action-conditioned CEM rollouts.Prior token-reduction methods were primarily developed for single-pass perception.Source: §4 · p. 9

Decision relevance and predictor compatibility

  • Planning-cost gradients connect token selection to downstream decisions, rather than assuming predictive fidelity or attention scores establish decision usefulness.Related work also reports unstable attention-based selections and planning outcomes across Wall predictors with matched validation loss, while planning-gradient selection is more stable in that setting.Source: §4 · p. 9
  • AdaLN supports pure planning-gradient sparsification in the reported experiments, whereas concat can make concentrated subsets unstable and require coverage-preserving selectors.The matched conditioning comparison adds stability under token removal as a criterion for action-conditioning design.Source: §4 · pp. 9–10

5 Discussion and Limitations

Sparse-planning performance depends jointly on control-objective relevance and compatibility between the selector and predictor. The evaluation is bounded to four simulated environments and tested model/planner settings; physical-system evaluation and broader conditioning comparisons remain open.

5 Discussion and Limitations

  • COSTGRAD's measured action-effect drift is lower than Random on AdaLN but higher on concat, tracking its loss of planning advantage on concat.The discussion recommends assessing token-selection performance jointly with the predictor at the intended token budget.Source: §5 · p. 10
  • Four simulated environments, one DINOv2 encoder family, and CEM define the evaluation scope; physical-system evaluation with perception noise, calibration errors, and model mismatch remains future work.Generality beyond the tested AdaLN and channel-wise concat predictors is also untested.Source: §5 · p. 10
  • Full-token prediction quality alone does not establish suitability for sparse control, because token restriction can change responses to candidate actions.The paper recommends evaluating world models under deployment token subsets and action searches.Source: §5 · p. 10

A Environments

  • The four visual planning tasks use an initial frame and goal frame: PointMaze navigation, Wall navigation, PushT manipulation, and MetaWorld Reach.The tasks are continuous-control benchmarks; MetaWorld Reach uses a 3D 7-DoF arm.Source: A Environments · p. 13

B Implementation and hyperparameters

The implementation uses frozen visual features, action-conditioned transformer predictors, and CEM, with selector-specific probes and shared retained-token positions. The appendix specifies predictor training, planning costs, and alternative importance computations.

Encoder and predictor

  • The frozen DINOv2 ViT-S/14 encoder maps each 224 × 224 image to a 16 × 16 grid of 384-dimensional tokens; the AdaLN predictor has six blocks and about 17.6M parameters.The predictor uses 400-dimensional internal embeddings, 16 heads, and action-derived AdaLN modulation parameters.Source: B Implementation and hyperparameters · p. 13
  • Matched-concat replaces AdaLN action conditioning with channel-wise concatenation of a 10-dimensional action vector to each visual token before q/k/v projections.The matched predictor shares the listed architecture and training settings with AdaLN apart from the action pathway.Source: B Implementation and hyperparameters · p. 13

Planner and pretraining

  • CEM plans over horizon 6 using 300 candidate sequences per iteration, 30 iterations, and 10 elites; retrained predictors use 20 epochs of AdamW training.The default cost compares final-step predictions with the goal, with proprioceptive weight 0.1 when applicable.Source: B Implementation and hyperparameters · p. 13

Objective and normalization

  • The sparse objective averages visual and proprioceptive errors over their own feature entries, so reducing retained-token count does not implicitly change the visual term's weight.The proprioceptive term is omitted when proprioception is absent; the selection probe uses a full-grid visual sum of squared errors.Source: B Implementation and hyperparameters · p. 14

Importance estimation

  • PREDGRAD scores tokens using retrospective prediction residuals against a detached encoding from observed history, then applies the selected indices to the current observation and goal.The target is not obtained by executing the probe action.Source: B Implementation and hyperparameters · p. 14

Alternative token reduction and attention

  • ToMe merges the most cosine-similar source–destination pairs and applies the resulting mapping across context frames; PREDATTN equally weights heads and frames without entropy weighting.ToMe rollout keeps surviving positions without rematching, and its visual cost averages uniformly over those positions.Source: B Implementation and hyperparameters · p. 14

C Decoded one-step predictions from COSTGRAD-50% sparse tokens

COSTGRAD-selected 50% token subsets retain qualitatively faithful decoded predictions across all four environments. The accompanying architecture analysis finds that token-removal effects differ between AdaLN and concat, while probe-action changes have limited reported impact on planning success.

C Decoded one-step predictions from COSTGRAD-50% sparse tokens

  • Decoded predictions with COSTGRAD-selected 50% token subsets are qualitatively faithful across all four environments, with missing-token artifacts concentrated in dropped background regions.The reported examples recover agent and goal positions, PushT object poses, and the MetaWorld arm/object configuration.Source: C Decoded one-step predictions from… · p. 14, Figure 5 · p. 15

D Action conditioning and token-removal sensitivity

  • AdaLN-Zero broadcasts action-derived modulation parameters across tokens, while concat adds action features per token before q/k/v projections; token removal can alter attention and output action effects in both.The AdaLN modulation parameters being subset-independent alone does not establish output invariance.Source: D Action conditioning and… · p. 15
  • The selector uses an L2 probe throughout, while the caption reports its largest gain under the L1 CEM cost, where full-token planning performs worst.The supplied table caption identifies this as a Wall planning-success comparison across CEM cost functions.Source: Table 4 · p. 16
  • COSTGRAD-selected/Random drift ratios are 0.29 ± 0.03 for AdaLN-Zero and 1.33 ± 0.17 for matched-concat across three Wall seeds per architecture, with the direction holding for every checkpoint.The diagnostic compares relative spatial distributions of squared action effects over the same retained positions; it is not a guarantee of action-ranking preservation.Source: D Action conditioning and… · pp. 15–16

E Probe action sensitivity (COSTGRAD and PREDATTN)

  • A zero-action probe is the default, and small Gaussian probe noise changes COSTGRAD planning success by at most 5 pp across the four environments.A CEM-mean probe was tested only on Wall and was also within 5 pp; per-token rankings had Spearman ρ > 0.9 across 10 probes on both tested architectures.Source: E Probe action sensitivity (COSTGRAD… · p. 16

F Cost function robustness

COSTGRAD remains useful across the tested Wall planning costs, with its benefit largest for the poorest cost. The L2, cosine, and L1 selection probes also choose substantially overlapping token sets.

F Cost function robustness

  • COSTGRAD’s benefit is robust across planning costs and largest when the cost is poorest.Table 4 varies the CEM planning cost while retaining the L2 selection probe for COSTGRAD in every row.Source: F Cost function robustness · p. 16
  • The Wall ablation compares planning success under L2, cosine-similarity, and L1 CEM costs.COSTGRAD uses the L2 selection probe in every row.Source: F Cost function robustness · p. 16
  • Top-128 Jaccard IoU is 0.917 between L2 and cosine selections and 0.841 between L2 and L1 selections.The diagnostic indicates substantially overlapping token sets across the three probe objectives.Source: F Cost function robustness · p. 16

G Matched compute baseline

On Wall, COSTGRAD with half the CEM iterations retains standard COSTGRAD-50% success while reaching about 5× overall speedup and exceeding Full. A matched Full run with half the iterations achieves a smaller speedup and lower success.

G Matched compute baseline

  • 92.0 ± 5.1% planning success at ∼5× overall wall-clock speedup exceeds the 84.7% Full baseline by +7.3 pp.This is COSTGRAD at 50% tokens with 15 rather than 30 CEM iterations; its mean matches standard COSTGRAD-50% at 92.0 ± 1.8%.Source: G Matched compute baseline · p. 17
  • The half-iteration COSTGRAD run has wider evaluation-seed variation than standard COSTGRAD-50%, but its worst-case seed still beats Full.The reported standard deviations are 5.1 versus 1.8, and the worst-case seed is 88.54%.Source: G Matched compute baseline · p. 17
  • Full at 15 CEM iterations achieves 84.4 ± 1.0% success on Wall for 2× wall-clock speedup.This matched-compute reference is within evaluation-seed noise of standard Full’s 84.7 ± 2.6%.Source: G Matched compute baseline · p. 17
  • COSTGRAD-50%-with-half-iters beats Full-with-half-iters by +7.6 pp (92.0 vs. 84.4) at the same iteration budget.Token selection adds a further 2.6× per-iteration cost reduction; the passage characterizes token selection and increased search as addressing different inefficiencies.Source: G Matched compute baseline · p. 17

H COSTGRAD operating curve

The Wall AdaLN-Zero operating curve suggests that success is roughly flat across several token budgets and often holds when CEM iterations are reduced. The highlighted K=128, 15-iteration point combines high mean success with a 5.2× speedup over Full.

H COSTGRAD operating curve

  • Figure 6 shows every measured COSTGRAD operating point on AdaLN-Zero, sweeping token budgets K ∈ {26, 64, 128, 192} and CEM iteration counts {10, 15, 20, 30}.Source: H COSTGRAD operating curve · p. 17
  • At 30 iterations, success is 85.9%, 92.0%, and 92.7% for K=64, 128, and 192, respectively.The passage describes the marginal benefit of doubling tokens beyond K=64 as small.Source: H COSTGRAD operating curve · p. 17
  • Within each K cluster, reducing iterations from 30 to 15 does not consistently reduce success; lower-iteration measurements scatter around the 30-iteration mean.Reducing iterations mostly slides operating points along the time axis at roughly fixed success, until iteration counts get very low.Source: H COSTGRAD operating curve · p. 17

I Additional diagnostic checks

The additional checks probe selector behavior, its operating point, and anchor mixing on concat predictors. They report stable probe rankings, a collapse when selecting the lowest-gradient tokens, and checkpoint-dependent anchor-mixing outcomes.

I Additional diagnostic checks

  • The concat planning results use K=128 of 256 tokens; Table 5 covers Wall and PointMaze across training seeds, while remaining checks use Wall.The table reports paired COSTGRAD–Random success-rate gaps across three training seeds.Source: I Additional diagnostic checks · p. 17, Table 5 · p. 18
  • Probe-action rank Spearman is 0.918 ± 0.058 on AdaLN-Zero and 0.958 ± 0.039 on matched-concat, with rankings stable under the tested probes.The diagnostic uses 30 observations and 45 action pairs per observation, from one checkpoint per architecture.Source: I Additional diagnostic checks · p. 17
  • Keeping the bottom-K tokens by gradient norm yields 0/60 successful episodes on AdaLN-Zero and 1/96 (1.04%) on matched-concat.The AdaLN-Zero run stopped after 60 of 96 planned episodes; inversion does not recover COSTGRAD’s advantage over Random on concat.Source: I Additional diagnostic checks · p. 17
  • Figure 6’s K=128, 15-iteration Wall operating point gives ∼17 s and 92.0 ± 5.1% success, a 5.2× speedup over Full at higher mean success.Its colors encode K clusters and marker shapes encode iteration count.Source: Figure 6 · p. 18
  • At α=0.8, anchor mixing beats Random on two of three concat checkpoints: 94.1±2.4% vs 79.5±2.1%, 89.6% vs 76.0%, and 76.0% vs 83.3%.The latter two comparisons each use one evaluation seed; results are checkpoint-dependent, and restoring drift toward Random does not guarantee a planning gain.Source: I Additional diagnostic checks · p. 18

J Quadrant analysis methodology

Easy-but-critical tokens are enriched only on PushT; the timing and analytical sections explain the savings and limits of sparse planning, while the public-checkpoint section specifies its comparison setup.

J Quadrant analysis methodology

  • PushT has almost twice the independence-reference count of easy-but-critical tokens, while Wall, MZ, and MW fall below it.The intersection uses bottom-quartile prediction error and top-quartile COSTGRAD score; the independence reference is 16.Source: J Quadrant analysis methodology · p. 18, Table 6 · p. 19

K Timing breakdown

  • 2.60× is the net wall-clock speedup at K = 128 vs. K = 256, exceeding naive 2× scaling as attention compute falls 4× and MLP compute falls 2×.The 27.53 ms forward+backward overhead amortizes to ∼1 ms per CEM iteration over 30 iterations, versus ∼1820 ms per-iteration savings.Source: J Quadrant analysis methodology · p. 19, K Timing breakdown · p. 19

J Quadrant analysis methodology

L Public-checkpoint control

  • Table 8 compares planning success on released checkpoints without retraining, with COSTGRAD and Random using 128 tokens versus Full's 256.Source: Table 8 · p. 19

M An analytical perspective on token selection

  • At a fixed probe with squared losses, exact prediction can make the prediction gradient zero without requiring the planning-cost gradient to vanish.The analysis concerns visual squared-distance cost, while the evaluated planner also includes weighted proprioceptive error.Source: M An analytical perspective on token… · p. 20
  • Keeping the K largest planning-gradient norms minimizes the linearized criterion for independent bounded perturbations to discarded tokens.This local perturbation model does not represent how physical token removal changes attention and subsequent rollouts.Source: M An analytical perspective on token… · p. 20
  • Sparse planning changes both which positions enter the cost and the predictions at retained positions; only action-dependent variation in ES − DS can change candidate rankings.Source: M An analytical perspective on token… · p. 20
Loading 2610.10274v1…