World models learn better global rules with compositional training
Do World Models Learn Global Understanding?
Machine LearningArtificial Intelligence
Summary
Many AI systems seem to understand facts locally but struggle to use that knowledge more broadly. The authors studied whether AI models can learn global rules from local examples by using specially designed math-based tasks. They found that usual training methods fail to help models apply complex rules across longer scenarios, but a new training style that hides intermediate steps helps them reason better and generalize. This improvement also helps AI models working with real environments and knowledge graphs show better understanding.
What this means in practice
- •For embodied ai developers: Improve AI agents’ reasoning about actions and their consequences by training with compositional paths that promote better constraint propagation.
- •For knowledge graph engineers: Enhance large language models’ relational reasoning on knowledge bases by using compositional training techniques to better infer unseen facts.
Authors
Alexander Detkov, Matt Thomson
Abstract
AI systems often feel brittle and fragmented. A large language model (LLM) may correctly explain a concept but fail to apply it, or follow safety instructions in one context but not another. This behavior suggests a general failure to lift local information to a global understanding. To gain fundamental insight, we frame "understanding" as learning constraints and propagating their consequences. We construct learning tasks on monoid worlds, sets of states connected by action transitions, where observed training transitions and an unseen constraint jointly determine held-out transitions. Measuring generalization tests whether models can learn global constraints from local transitions and propagate their consequences. We consider inverse, commutativity, composition, and periodicity constraints relevant to spatial and semantic structure. Across attention, recurrent, and state-space architectures, next-state training fits the data but fails to propagate non-trivial constraints. Compositional training, which uses identical paths but hides intermediate states from the input, achieves 96% accuracy on inverse, commutativity, and composition constraints across architectures, yields corresponding improvements in geometric generalization of world models trained on embodied environments and relational generalization in Wikidata-finetuned LLMs. How far do models propagate constraints when inferring an unseen fact may depend on first inferring others? We define proof depth d of a held-out transition, measuring the minimum number of inference rounds to infer the transition, and find that model generalization decreases sharply with proof depth. Increasing compositional path length T improves generalization. These results provide a formal way to investigate global understanding in language and world models and demonstrate that compositional training promotes information propagation and integration.