ONCE: Optimized Neural Domain Decomposition for Scalable PDE Solving

Neural surrogates enable efficient PDE solving for scientific discovery and engineering design, yet generalizing to larger domains and unseen geometries remains challenging due to the limited domain and geometric distributions observed during training. Domain decomposition methods (DDMs) provide a promising route to overcome this bottleneck by iteratively coupling the locally trained solver across subdomains. However, the local solver is typically trained under reference states and repeatedly evaluated on states generated by its own predictions. This training-inference mismatch can compromise global accuracy even when the iteration converges. To address these challenges, we present ONCE, a novel neural DDM framework for scalable PDE solving that unites optimized Schwarz methods (OSMs) with Trajectory-aligned Tuning (TAT). Specifically, TAT collects intermediate interface states from the pre-trained model rollouts on the original training cases and fine-tunes the local solver against the corresponding reference subdomain solutions, while retaining reference-state supervision to preserve consistency. This procedure requires neither additional numerical PDE solves nor backpropagation through the iterations. Our theoretical analysis characterizes fixed-point convergence of neural DDMs and establishes interface error bounds that account for trajectory supervision and distribution shift. Experimentally, ONCE consistently outperforms existing methods across benchmarks spanning varying domain scales and unseen geometries, demonstrating its superior accuracy and scalability.