n-CLU: A Factorized Latent Dynamics Model for Deletable Compositional Generation

A world model trusted with unseen tasks needs to go beyond generalizable predictions. Among other capabilities, dynamic tasks demand that (a) the model outputs are valid under the task's judge (b) the model composes new combinations of factors from the ones it has seen, and (c) the model can forget specific factors when the environment changes permanently. Valid, compositional generation is the primary focus of this work with machine unlearning as a supplementary goal. We present the naive-Causal Learning Unit (n-CLU), and compare against naive baselines (an MLP, an LSTM, a transformer and a continuous-input MLP) on two simulated evaluation tasks: shape-rendering and maze-solving. The n-CLU is designed as a set of factorized, principal subspaces acting on the latent space of a frozen neural-network encoder, allowing composition of encoder-extracted features. The n-CLU is then read by a request-indexed address book, and finally rendered by a frozen decoder. Deleting a training row only downdates the statistic and refits the address book, so the trained n-CLU after deletion equals an n-CLU fitted without those rows. On a held-out set of the Shapes3D dataset, the n-CLU composes new pairs and combinations of every factor. Against matched baselines trained on the same frozen encoder and decoder, the n-CLU produces $\sim 7$ more distinct, valid novel tuples per held-out combination than the strongest baseline, and beats every baseline on every evaluation set of an $n$-way combination task across multiple seeds. To test compositional generation in a harder, world-model aligned task, we created a maze-based benchmark (CMAZE-T) with composable factors and validity rules an agent must follow during navigation, based on the state of the local environment. We also release this small-scale benchmark for valid, compositional generation and plan to expand on it in future works. On CMAZE-T the n-CLU consistently matches the strongest baseline (transformer) with a marginal seed-dependent lead, and we show both quantitative and qualitative results for the same. We finally motivate a Physics-designed extension to the n-CLU, and how it could provide clues towards designing better latent-dynamics based world models.