Neural SHAKE: geometric constraints in neural differential equations

Abstract

Generating accurate molecular conformations hinges on sampling effectively from a high-dimensional space of atomic arrangements, which grows exponentially with system size. To ensure physically valid geometries and increase the likelihood of reaching low-energy conformations, it is us ful to incorporate prior physicsbased information by recasting them as geometric constraints that naturally arise as nonlinear constraint satisfaction problems. In this work, we propose an approach to embed these strict constraints into neural differential equations, leveraging the denoising diffusion framework. By projecting the stochastic generative dynamics onto a manifold defined by constraint sets, our method enforces exact feasibility at each step, unlike alternative approaches that merely impose soft constraints through probabilistic guidance. This technique generates lower-energy molecular conformations, enables more efficient subspace exploration, and formally subsumes classifier-guidance-type methods by treating geometric constraints as strict algebraic conditions within the diffusion process.

Neural SHAKE formulates exact manifold‑projected score‑based diffusion : each reverse-SDEincrement is orthogonally projected, via a Lagrange-multiplier solve, onto the constraint surfaceσₐ(x)=0 for a = 1,…, A, with A the number of independent constraints and thus the manifold’scodimension . This projection preserves global SE(3) symmetry and enforces constraints tosolver tolerance. It induces a well-posed surface Fokker–Planck flow on the (3 N − A)-dimensional manifold, while a coarea/Fixman Jacobian carries the ambient 3 N-dimensionaldensity to a normalized density on that manifold, preserving probability mass after the dimensionality reduction.