Residual Learning of Few-Body Molecular Wave Functions with Symmetry-Aware Representations

Ultracold polar molecules are a highly controllable platform for engineering quantum matter, but predicting their few-body properties remains computationally challenging. Their anisotropic, long-range interactions induce intricate low-symmetry correlations, leading to high-dimensional eigenvalue problems that conventional grid-based solvers are largely restricted to two-molecule systems. Neural wave-function methods offer a promising alternative: they already solve many-body problems accurately, from lattice spin models to ab initio electronic structure via FermiNet and PauliNet, but existing work on continuous-coordinate few-body systems has mainly considered structureless particles with short-range, isotropic interactions, leaving realistic molecular interactions comparatively unexplored. We address this gap with a physics-informed neural wave-function method whose variational ansatz is derived from the structure of the underlying Hamiltonian. We encode translational invariance, angular structure, and asymptotic decay analytically, leaving a neural residual to learn only the remaining correlations. Spherical-harmonic decomposition reduces the problem to coupled radial channels represented by a shared multilayer perceptron, while a WKB-inspired envelope enforces the correct short- and long-range behavior. Training uses a fixed physics-informed Monte Carlo grid, avoiding repeated resampling from an evolving model distribution. We benchmark the method on two-body field-linked molecules against a high-accuracy grid-based reference solver, mapped discrete variable representation (MDVR) across Rabi couplings from 50 to 1000~MHz. The model recovers ground- and excited-state energies with mean absolute relative errors of $3\times10^{-3}$ and $1.1\times10^{-1}$, respectively, with the larger excited-state error concentrated in the weakly bound regime. Controlled comparisons show that the spherical decomposition, shared parametrization, and asymptotic envelope substantially improve optimization efficiency and stability: a direct full-coordinate representation requires roughly $15\times$ larger batches and $100\times$ longer training on the same hardware. Beyond this core result, the same construction extends systematically to more particles through Jacobi coordinates and tensor products of spherical harmonics. Two-body systems reduce to a single radial coordinate, while three-body systems introduce a second Jacobi coordinate and therefore require a two-dimensional radial residual. We derive the corresponding three-body Hamiltonian and wave-function ansatz and validate the required two-dimensional numerical machinery in the two-body setting, where the resulting solver reproduces the MDVR reference energies to within about $1\%$ on average. Together, these results suggest that embedding known physical structure into neural wave-function representations is a promising route toward first-principles simulation of more complex correlated molecular systems.