WFM-TTS — Test-Time Scaling for World Foundation Models

Paper. Wenyan Cong, Hanqing Zhu, Peihao Wang, Bangya Liu, Dejia Xu, Kevin Wang, David Z. Pan, Yan Wang, Zhiwen Fan, Zhangyang Wang. Can Test-Time Scaling Improve World Foundation Model? COLM 2025. UT Austin · UW-Madison · NVIDIA. [arXiv] · [project page] · [code]


0. The Picture in One Paragraph

World Foundation Models (WFMs) — video-generative models that predict future visual states given past observations and control inputs — have become load-bearing components in Physical AI stacks for autonomous driving and robotics. But they’re expensive to make better: pretraining costs are enormous, and post-training improvements are bottlenecked by curated paired data. This paper asks whether test-time compute scaling is a third path — one that has been transformative for LLM reasoning but has never been systematically applied to WFMs. The answer is yes: the authors introduce WFM-TTS, the first test-time scaling framework for WFMs. The central engineering challenge is that naive multi-candidate search is impractical over continuous video tokens with a diffusion-based decoder. WFM-TTS solves this with a replacement efficient tokenizer decoder that is ~9,000× faster than diffusion decoding, then applies beam search + top-K pruning guided by rule-based physical-quality rewards. A new five-dimension evaluation toolkit (3D consistency, temporal consistency, spatial awareness, perceptual quality, text-video alignment) provides the measurement infrastructure the field was missing. The headline result: a 4B-parameter COSMOS model with WFM-TTS matches or outperforms a 12B-parameter baseline, and test-time scaling laws hold — more compute monotonically improves output quality.


1. The Problem — WFMs Are Hard to Scale, Easy to Evaluate Poorly

1.1 Why WFMs matter

A World Foundation Model takes a sequence of past observations (e.g., camera frames from an autonomous vehicle) and ego-action commands and predicts the next video frames — what the world will look like if the agent does those things. For physical AI:

  • Autonomous driving: the vehicle can “imagine” its trajectory before executing it.
  • Robotics: the robot can simulate proposed motions to check for collisions or physical plausibility before committing.

This is the “planning by imagination” paradigm, and it scales directly with how good the WFM is.

1.2 The scaling wall

WFMs are trained on massive video datasets with large compute budgets. Making them better has two levers:

Lever Cost Constraint
More pretraining data / compute Enormous Diminishing returns; data quality ceiling
Post-training fine-tuning Moderate Bounded by curated paired (video, action) data
Test-time compute scaling Per-query Unexplored for WFMs

Test-time scaling — the idea behind LLM chain-of-thought and OpenAI o1 — trades inference compute for quality without retraining. It works for LLMs because you can generate multiple candidates and use reward models to pick the best. For WFMs the obvious extension is blocked by two problems: (1) diffusion-based decoders are far too slow for multi-candidate generation, and (2) no standardized reward or evaluation toolkit exists to score candidate videos systematically.

WFM-TTS addresses both.

1.3 The missing evaluation infrastructure

A secondary problem: there’s no agreed metric for WFM output quality. FID and FVD measure frame quality and video quality respectively, but they don’t capture physical plausibility — whether the generated world is physically consistent in 3D, whether objects maintain their spatial relationships, whether the generated scene is temporally coherent across frames. Without those measures, you can’t run principled reward-guided search, and you can’t compare WFMs systematically.


2. WFM-TTS — The Framework

WFM-TTS has three components: an efficient decoder, a search procedure, and a reward signal. The backbone model is NVIDIA COSMOS — an autoregressive WFM that operates over discrete video tokens.

2.1 The efficient tokenizer decoder

The standard COSMOS decoding pipeline reconstructs video frames from latent tokens using a diffusion-based decoder. This is accurate but extremely slow — slow enough that generating even a handful of candidate videos for comparison is computationally infeasible in a test-time search loop.

WFM-TTS replaces the diffusion decoder with an efficient tokenizer decoder that:

  • Maintains consistent quality trends relative to the diffusion decoder (same ordering of candidates by quality).
  • Runs approximately 9,000× faster.

This speedup is what unlocks beam search. Without it, test-time scaling for WFMs is a dead end.

2.2 Probability-based top-K pruning

At each generation step, the autoregressive WFM produces a distribution over possible next tokens. WFM-TTS generates K candidate continuations and prunes the lowest-scoring ones by token probability. Only the top-K beams survive to the next step.

This is the standard beam-search pruning strategy, now made feasible by the fast decoder.

The search runs over the full discrete token space of the WFM:

  1. Expand each surviving beam to K children.
  2. Score each child using the efficient decoder + reward function.
  3. Keep the top-K children; discard the rest.
  4. Repeat until the full video sequence is generated.

The output is the highest-scoring complete video sequence found during search.

2.4 Rule-based rewards

The reward signal that guides beam selection is deliberately simple: rule-based physical quality metrics derived from measurable properties of the generated video. Examples include:

  • Aesthetic quality score.
  • Object permanence / temporal consistency score.
  • 3D geometric consistency (via CUT3R, a 3D reconstruction model).

The paper compares rule-based rewards against learned preference models and finds that rule-based rewards consistently win on both stability and alignment with qualitative human preference. Learned rewards introduce noise; rule-based rewards are deterministic and directly tied to measurable physical properties.


3. The WFM Evaluation Toolkit

Alongside the search framework, WFM-TTS introduces a five-dimension evaluation toolkit that serves as both the reward signal during search and the benchmark for comparing WFMs broadly.

Dimension What it measures How
3D consistency Geometric coherence of the generated scene CUT3R 3D reconstruction from the video
Temporal consistency Smooth transitions; no object teleportation Frame-to-frame coherence metrics
Spatial relationship awareness Correct relative positions of objects Object detection + spatial reasoning
Perceptual quality Visual fidelity and aesthetics VBench / VideoScore
Text-video alignment Generated video matches the action conditioning Cross-modal similarity

Standard metrics (FID, FVD) are evaluated separately on the nuScenes and Waymo distributions, covering 900 test input videos (150 nuScenes scenes + 750 Waymo clips from 150 scenes × 5 samples).

The toolkit is modular and is released as open-source infrastructure for WFM research — filling the measurement gap that had been blocking principled comparison of WFMs.


4. Results

4.1 Test-time scaling law holds for WFMs

The paper’s central empirical finding:

Performance on every evaluation metric increases monotonically as test-time compute (number of beam-search iterations, candidate count) increases — even in a compute-optimal (FLOPs-matched) setting.

This is the WFM analog of the LLM test-time scaling law. It’s not obvious that it would hold, because WFMs operate over continuous-valued visual predictions rather than discrete language tokens. It does hold.

4.2 4B matches 12B

The most striking headline:

COSMOS-4B + WFM-TTS ≥ COSMOS-12B (baseline, no TTS)

A model with 3× fewer parameters, given more inference-time compute via WFM-TTS, matches or outperforms the larger model on overall quality. Human evaluators prefer the WFM-TTS-enhanced outputs.

This is the same compute-optimal crossover point seen in LLM scaling — you can substitute inference compute for model capacity, and at some point the substitution is favorable.

4.3 Quantitative gains

Overall improvements of 6–12% across evaluation dimensions relative to greedy/baseline decoding (no test-time search).

4.4 Rule-based > preference-based rewards

Across all settings, rule-based reward functions outperform learned preference models for beam search guidance. The preference models introduce calibration noise; the rule-based rewards are stable and physically grounded.


5. The Broader Context

5.1 Connection to LLM test-time scaling

The LLM test-time scaling research established that “more compute at inference = better outputs” when search is structured — not random sampling but principled search with evaluation signals. WFM-TTS is the first paper to extend this to physical world modeling, establishing that the same principle applies across modalities.

The key enabling difference: language tokens are cheap to decode; video frames from diffusion are not. The 9,000× speedup in the efficient decoder is therefore the structural prerequisite for the entire framework, the way softmax temperature is the prerequisite for language model sampling.

5.2 Relation to DreamGen and DreamZero

In the robotics world-model lineage reviewed on this site (DreamGen, DreamZero), the WFM is used to generate synthetic trajectories or as a policy backbone. WFM-TTS is upstream of both those use cases: it makes the WFM better at the video-generation task itself, without changing the architecture or training. Better WFM → better synthetic trajectories for DreamGen → better downstream policies. Better WFM → better predictive features for VPP-style backbones.

This is the “test-time compute as a lever for WFM quality” argument; the downstream benefits to the robotics pipeline are a corollary.


6. Why It Matters

Three reasons:

  1. It establishes test-time scaling as a first-class lever for WFMs. Until this paper, test-time compute scaling was LLM territory. WFM-TTS shows that physical world modeling obeys the same scaling law, provided you solve the decoder-speed bottleneck. The implication: the “more inference compute → better quality” recipe now spans language reasoning and physical simulation.
  2. The efficient decoder is the enabling contribution. A 9,000× speedup that maintains quality ordering is nontrivial engineering. Without it, beam search over WFM token spaces is impractical and the rest of the framework doesn’t exist. This is a reusable component for anyone working with autoregressive tokenized WFMs.
  3. The evaluation toolkit fills a genuine gap. FID and FVD were borrowed from image generation and video generation respectively — they don’t measure physical plausibility. A toolkit that captures 3D consistency, temporal coherence, and spatial relationship awareness is necessary infrastructure for the WFM research community.

7. Limitations Worth Knowing

  • COSMOS-specific. All experiments use NVIDIA COSMOS at 4B and 12B scale. Whether the approach extends to diffusion-based WFMs (rather than autoregressive tokenized WFMs) requires different search designs.
  • Autonomous driving focus. Evaluation datasets are nuScenes and Waymo — driving-centric. Robotics scenarios are mentioned but not quantitatively evaluated. The rewards and toolkit metrics may need domain adaptation for manipulation tasks.
  • Rule-based rewards have a quality ceiling. They are stable and physically grounded, but they measure specific observable properties. Tasks requiring richer goal-conditioned evaluation (did the agent accomplish a specific behavior?) may need learned rewards anyway.
  • Beam search still multiplies inference cost. Even with the 9,000× speedup, keeping K beams active multiplies per-query compute by K. At large K, the practical deployment cost is non-trivial. The paper doesn’t fully analyze latency-quality Pareto curves under production constraints.
  • Offline evaluation. Quality is assessed on fixed test sets from nuScenes/Waymo. Closed-loop driving evaluations — where the generated future feeds back into planning decisions — are not included.

8. The Takeaway for a First Reader

If you remember three things:

  1. World Foundation Models are expensive to improve by retraining. WFM-TTS asks whether test-time compute scaling is an alternative. It is: test-time scaling laws hold for WFMs (monotonically improving quality with compute) even in a compute-optimal (FLOPs-matched) regime.
  2. The key enabling technology is a ~9,000× faster efficient tokenizer decoder that replaces slow diffusion decoding and makes multi-candidate beam search over WFM token spaces practical. On top of this, rule-based physical-quality rewards guide beam selection more stably than learned preference models.
  3. The headline result: COSMOS-4B + WFM-TTS ≥ COSMOS-12B (baseline), with overall 6–12% quality gains. The companion five-dimension WFM evaluation toolkit (3D consistency, temporal consistency, spatial awareness, perceptual quality, text-video alignment) fills a measurement gap in the field.

That’s the arc: WFMs can’t be improved cheaply by retraining → test-time scaling is the unexplored lever → fast decoder makes beam search feasible → rule-based rewards guide search → smaller model beats larger model on physical-quality metrics.


References




    Enjoy Reading This Article?

    Here are some more articles you might like to read next:

  • The Expressive Power of Transformers with Chain of Thought
  • Landscape of Thoughts — Visualizing Where LLM Reasoning Actually Goes
  • Magellan — Guided MCTS for Escaping the Gravity Wells of LLM Creativity
  • PriorZero — Injecting LLM Priors into MuZero-Style World Models at the MCTS Root
  • SuperThoughts — Reasoning Tokens in Superposition