WM-Diff: World Model-Guided Planning for Angular Diffusion in Protein Backbone Generation

Internal-coordinate diffusion provides a translation and rotation-invariant representation for protein backbone generation, but existing samplers remain largely myopic by optimizing only the current denoising state without considering future structural consequences. Due to serial forward kinematics, small angular perturbations can accumulate into large Cartesian deviations and reduce sampling stability. We introduce WM-Diff, a world model-guided planning framework that equips angular diffusion with lookahead capability through horizon-conditioned future prediction. WM-Diff learns future steric cost-to-go from deterministic reverse-diffusion rollouts of a frozen backbone generator rather than instantaneous structural penalties. During sampling, the world model provides differentiable future-cost feedback to steer trajectories toward lower downstream steric cost. On an unconditional protein backbone generation benchmark, WM-Diff improves self-consistency designability from 27.5% to 52.0% while reducing steric clashes from 15.0 to 1.8 without post-hoc structural relaxation. Mechanistic analyses show that future-cost prediction provides a more effective planning signal than instantaneous objectives while preserving structural diversity. These results demonstrate the effectiveness of world model-guided planning for protein backbone generation.