Improving Convolutional Neural Network Wave Functions Optimization

Convolutional neural networks (CNN) have become a staple of modern machine learning research, where they are ideally suited for learning local features with spatial inputs such as images. This has made them an obvious candidate for variational wave functions to represent 2D systems. However, their optimization with state of the art variational Monte Carlo (VMC) methods relies on natural gradient descent, which becomes intractable for a large number of variational parameters. This is due to requiring inverting a matrix whose dimensions scale with the number of parameters. We propose an approximate natural gradient method that minimizes an upper bound on the exact natural gradient descent residual. The method relies on grouping the parameters of the CNN by dependency and inverting sub-matrices of dimension equal to the number of parameters in that group. We Implement our method on a deep CNN with Res-Block architecture and complex valued parameters. We benchmark our results on the frustrated 2D J1-J2 Heisenberg model on the square lattice.