The Cross-Entropy Method, or CEM, is a sampling-based optimizer. It searches for a good action sequence when the planner cannot directly solve the optimization problem.
Planning steps
- We first start with a probability distribution over action sequences
- Sample many candidate sequences (many ways to go from initial state to goal state)
- Use a World Model to predict the consequence of each sequence
- Score each predicted consequence with an Energy-Based Model
- Keep the best candidates, which are called the elite set, this is something like the top-k tokens in a transformer
- Update the action distribution using the elite set
- Repeat the sampling and update process
- Execute the first action from the selected sequence and plan again
Role in V-JEPA 2 and GeoWorld
V-JEPA 2 uses CEM to search for action sequences that bring the predicted future representation close to a visual goal.
GeoWorld also uses CEM. Its main change is the hyperbolic representation and geometry-aware training.
Compared with greedy token selection
CEM is similar as taking the top-1 next token. In the transformer, we choose the top k number of candidates for our next token prediction. Similarly, CEM evaluates a population of complete candidate action sequences.
My way of visualizing CEM
I think of CEM as the optimizer that finds the best path. It proposes several candidate action sequences, passes them into the World Model, and keeps the sequences that produce future representations closest to the goal.
It selects an elite group of complete sequences, updates the sampling distribution, and repeats the search before choosing the final sequence.
This also means that CEM is the search method. GeoWorld changes the latent geometry and the energy used to score each sequence.