FLDWM: World Model for learning approximate SO(3) equivariance for scalable and generalizable molecule generation

Equivariant neural networks provide physically accurate representations for 3D molecular generation but introduce computational bottlenecks that limit scalability, while scalable voxel-based methods lack explicit rotational symmetries. To bridge this gap, we introduce the Forced Latent Denoising World Model (FLDWM), the first world-model-based framework for 3D molecule generation. Operating on compressed voxel representations, FLDWM formulates symmetry learning as a latent dynamics problem: the model is trained to invert random SO(3) rotation augmentations, conditioned on the applied rotation via Feature-wise Linear Modulation (FiLM), which we show theoretically forces the denoiser to learn approximate SO(3) equivariance, bounded only by voxel discretization error. On GEOM-Drugs, FLDWM significantly outperforms state-of-the-art non-equivariant (VoxMol, NEBULA, RADM) and equivariant (EDM, MiDi) baselines on stability and validity while retaining fast inference, and demonstrates zero-shot generalization to the out-of-distribution PCQM dataset without fine-tuning. As a practical illustration, seed-conditioned generation from five FDA-approved drugs yields novel lead-optimization-range analogs (ECFP4 $T_c\in [0.4,0.7]$), showing FLDWM's potential to propose alternative drug candidates.