Feature Learning for the High Dimensional Stationary Schödinger Equation with Deep Ritz Method

2026-07-07

AI Summary

This work addresses the challenge of solving high-dimensional stationary Schrödinger equations under Neumann boundary conditions by investigating the feature learning mechanism of neural networks within the deep Ritz framework. By constructing single-exponential and two-neuron multi-exponential models and integrating a regularized Ritz energy functional with Riemannian gradient descent, the authors theoretically prove that the algorithm converges to an approximate global optimum in $O(\log(1/\varepsilon))$ iterations. The study reveals an optimal feature alignment mechanism for single-exponential source terms and, for the first time, characterizes the emergence pattern of the second eigenfeature in the two-neuron model as a function of the regularization parameter. Numerical experiments corroborate the theoretical predictions, demonstrating the efficacy of the proposed approach.

Abstract

This paper investigates feature learning within the framework of the deep Ritz method for solving the stationary Schrödinger equation with Neumann boundary conditions. We first analyze the convergence of Riemannian gradient descent in an agnostic setting, where the hypothesis function is restricted to a single-index model while the PDE solution is arbitrary. We prove that gradient descent reaches an approximate global minimum: after T = O(log(1/ε)) iterations, the loss is within εof a constant multiple of the optimal loss. We then examine the loss landscape when the source term of the PDE itself follows a single-index model, considering hypothesis functions given by either a single-index model or a two-neuron multi-index model. In the single-index case, we show that the minimum Ritz energy is attained at the feature vector aligned with that of the source term. In the two-neuron case, we study the landscape of regularized Ritz losses and characterize how a second feature emerges, given that the first feature is aligned with the source, as the regularization parameter varies. Finally, numerical experiments are presented to validate the feature emergence theory in the two-neuron setting.

Problem

• feature learning
• stationary Schrödinger equation
• deep Ritz method
• single-index model
• loss landscape

Innovation

• feature learning
• deep Ritz method
• single-index model
• Riemannian gradient descent
• regularized loss landscape