How do world models handle multi-modal future distributions when a car at an intersection could turn left, turn right, or go straight, as in driving world models like Wayve’s GAIA?
Answer
The future at an intersection is genuinely multi-modal, so the first thing a world model must avoid is regression to the conditional mean. The minimizer of a squared-error objective is , which means a network that has perfectly learned that left, straight, and right are all plausible will output their average: a path down the median that no driver would ever take. The fix is architectural rather than a loss-weighting trick. A world model is built as a conditional sampler, not a point predictor, so drawing from it returns one internally consistent future at a time and the multimodality appears across draws rather than inside a single blurred output. Three mechanisms dominate in practice: a stochastic latent sampled at every rollout step (Dreamer-style RSSM), ancestral sampling of discrete tokens over a learned video codebook (GAIA-1-style), and iterative denoising from Gaussian noise (GAIA-2-style latent diffusion). Trajectory-level stacks add a fourth: a small set of anchors or mixture components with explicit probabilities, trained with a winner-take-all loss so each head owns one maneuver.
(1) Squared Error Averages Modes: any L2-trained deterministic head converges to the mean of the modes, and at a T-junction that mean is a physically impossible maneuver even though it minimizes the loss.
(2) Sample, Do Not Average: the model defines and multimodality is expressed by independent draws, each of which must remain a single coherent maneuver from first frame to last.
(3) Stochastic Latents: an RSSM splits state into a deterministic recurrent part and a sampled part
(DreamerV3 uses 32 categoricals of 32 classes), so branching happens once per step and is then carried forward consistently.
(4) Discrete Token Sampling: quantize frames into codebook tokens and sample autoregressively with temperature or top-p, which turns mode choice into ordinary categorical sampling at the price of a long token sequence per second of video.
(5) Explicit Mode Heads: anchored components with probabilities
give the planner a calibrated, enumerable set of maneuvers instead of an opaque sampler, at the cost of a fixed mode budget.
(6) Metrics Must Reward Coverage: single-sample L2 rewards mode averaging, so evaluation moves to , maneuver recall, and probability calibration, with
the standard budget on the Waymo Open Motion Dataset.

Figure 1: The same scene under two objectives. The L2 optimum is the pointwise average of the three maneuvers, so it drifts up the middle of the junction and matches none of them, while 12 draws from a stochastic model are individually legal and jointly cover all three modes. Nothing is wrong with the averaged model’s likelihood estimate of the mean; the problem is that the mean is not an admissible trajectory.
Where the randomness enters decides how the model behaves in a rollout. In a latent state-space model the sample is a small categorical latent drawn once per timestep, so a single decision at the junction propagates through the recurrent state and the decoded frames stay consistent for the rest of the horizon; the risk is posterior collapse, where the KL term is tuned so aggressively that the prior stops carrying maneuver information and rollouts become deterministic again. In a discrete autoregressive model the sample is a token, and mode choice is spread over hundreds of tokens per frame, which makes temperature a global blur-versus-diversity knob: too low and every rollout goes straight, too high and lane geometry falls apart. Diffusion models place the randomness in the initial noise vector, giving the best sample fidelity and the most controllable conditioning, but they pay 20 to 50 network evaluations per sample and cannot easily produce a probability for each maneuver. Explicit mixture heads sit at the opposite end: a Wayformer-style model emits Gaussian components with softmax weights in one forward pass of a few milliseconds, which is what a downstream planner actually wants, but six is a hard ceiling on expressible futures.

Figure 2: All three families implement the same idea with different noise sources, and the tree shows why it works: a draw at the first branch commits to a maneuver, later draws only refine speed and gap acceptance, and the leaf probabilities recover the marginal maneuver distribution (0.28 left, 0.45 straight, 0.27 right). A deterministic model collapses this tree to its centroid.
Mathematical Formulation:
Where:
is the future to be predicted (a trajectory of waypoints or a sequence of frames) and
is the squared-error optimum, which is the conditional mean and therefore not generally a valid maneuver.
is the conditioning context: past observations, lane graph, ego speed, and traffic-light state.
indexes mixture components or anchors, with weights
summing to one, means
, and covariances
;
is the standard benchmark budget.
is the stochastic latent whose prior
carries the mode choice; marginalizing it makes
multi-modal even when
is unimodal.
indexes rollout steps or tokens, and the product form is the ancestral sampling factorization used by discrete-token world models.
Why L2 Prefers The Impossible Path (two equally likely 8 m lateral outcomes):
Committing to a real maneuver scores twice as badly as predicting the physically impossible average, which is the whole reason single-output regression is untrustworthy here. Mode-based training removes the incentive by only penalizing the closest component and learning the weights separately, and rollout models remove it by making the loss a likelihood rather than a distance.
Mode-Based Training And Scoring:
The evaluation side matters as much as the model. Reporting a single-sample average displacement error silently rewards mode averaging, so benchmarks score and
over
samples, plus a mAP-style metric that requires the probabilities to be calibrated and not just the geometry to be covered. That combination is deliberate:
alone can be gamed by spraying diverse but implausible samples, while likelihood alone can be won by a model that puts all its mass on “straight” because straight is the majority class at most intersections. In production the planner consumes both, treating each mode as a separate scenario to cost, so a mode that is missing from the sample set is a scenario the planner never considers.

Figure 3: Drawing more samples only helps if the samples differ. A mode-collapsed sampler flattens out near 1.8 m because sample 6 repeats sample 1, while a diverse sampler keeps improving; the bar panel shows the consequence the planner feels, namely how often the maneuver the other car actually performed is present anywhere in the predicted set.
| Property | Anchored mixture head | Latent state-space (RSSM) | Discrete token AR | Latent diffusion |
|---|---|---|---|---|
| Source of multimodality | K anchors with softmax weights | Categorical latent drawn each step | Temperature or top-p token sampling | Initial Gaussian noise vector |
| Cost per future | One forward pass gives all K | T cheap recurrent steps | Hundreds of tokens per frame | 20 to 50 denoising steps |
| Explicit probabilities | Yes, directly usable by a planner | Only via repeated sampling | Sequence likelihood, hard to read | No tractable density |
| Expressible futures | Capped at K (usually 6) | Combinatorial in latent codes | Unbounded, pixel-level detail | Unbounded, highest fidelity |
| Dominant failure | Dead heads under winner-take-all | Posterior collapse, deterministic rollouts | Low temperature collapses to majority mode | Latency, and no calibrated mode weights |
Leave a Reply