QMCPACK / qmcpack

Main repository for QMCPACK, an open-source production level many-body ab initio Quantum Monte Carlo code for computing the electronic structure of atoms, molecules, and solids with full performance portable GPU support
http://www.qmcpack.org
Other
295 stars 138 forks source link

Simple implementation of Stochastic Reconfiguration #5017

Closed camelto2 closed 1 month ago

camelto2 commented 4 months ago

Please review the developer documentation on the wiki of this project that contains help and requirements.

Proposed changes

Wanted to out this up to get some eyes on it.

This implements a very simple version of stochastic reconfiguration. The benefit is that we avoid the <Psi_i | H | Psi_j > matrix elements which get costly for a large number of parameters. This stochastic reconfiguration approach basically solves -tau h = S dp for dp, where tau is a small projection time, h is the vector <Psi_i | H | Psi_0>, and S is the overlap matrix <Psi_i | Psi_j>. The simple approach would be to build S, invert, and solve for dp. A better method, employed here, is to use a conjugate gradient method to solve the linear equation above. This avoids having to explicitly build the entire S matrix, and only builds the S*z matrix-vector product at each iteration of the CG algorithm...essentially this is the "Accelerated" SR approach described in https://doi.org/10.1103/PhysRevB.85.045103, described in the text surrounding Eqn. 7.

Preliminary tests using this have allowed us to scale up 100,000+ variational parameters, which was simply not possible with the current implementation of one_shift_only and has a huge speedup over one_shift for large parameter counts.

This is still a WIP because there are a number of things to improve. A few things off the top of my head are

  1. Need to add meaningful tests. This was the result of some exploratory coding to see if it could work. Now that our preliminary testing seems to suggest the method works pretty well and is fast, I need to add proper tests.
  2. Would like to abstract away the CG solver
  3. Probably need a better implementation of the CG solver...I literally implemented whatever was on wikipedia. I have some hard coded CG convergence criteria which may not be sufficient
  4. Right now, it either uses a fixed tau to get the parameter update or it can do a line search with correlated sampling. In some cases, the line search is far superior and converges significantly faster. However, if the correlated sampling weight gets small and stays small throughout the optimization then the energy can blow up and tank the optimization. In those cases, you just need to use the tau as something small, but this can result in many optimization iterations. Each iteration is a lot faster, but you may end up needing O(100) or more to get convergence. Also if the tau is too large, then the optimization can go haywire as well. I think an optimal approach would be to use the line search unless the weights get too small, and just take steps with small tau otherwise. Or some sort of adaptive method which scales the timestep automatically to help accelerate convergence.
  5. probably many other things I'm not thinking of.

I'm going to be on leave for a few weeks, but I wanted to get this up to get some comments/suggestions. This seems to be the approach other QMC codes are using for large parameter count optimizations, so it would be nice to have something like this in the code. But I think it needs a lot of improvements.

What type(s) of changes does this code introduce?

Delete the items that do not apply

Does this introduce a breaking change?

What systems has this change been tested on?

Checklist

Update the following with a yes where the items apply. If you're unsure about any of them, don't hesitate to ask. This is simply a reminder of what we are going to look for before merging your code.

prckent commented 4 months ago

Test this please

prckent commented 1 month ago

Cody: Can we revisit the WIP label of this PR and get it merged? This addition is a great achievement -- several people have been able to successfully use it in real calculations. After fixing the conflict and making any other updates you wish, we can review and make a list of what needs addressing in future PRs. e.g. Some small documentation and a basic test to exercise the code paths.

camelto2 commented 1 month ago

Cody: Can we revisit the WIP label of this PR and get it merged? This addition is a great achievement -- several people have been able to successfully use it in real calculations. After fixing the conflict and making any other updates you wish, we can review and make a list of what needs addressing in future PRs. e.g. Some small documentation and a basic test to exercise the code paths.

Glad to hear that people have tried it out successfully. I will fix the conflict. I might try to refactor part of it and make a smaller PR for the krylov solver which will be easier to write a test for. Then we can incorporate those changes into this optimizer and get it merged if that sounds good.

prckent commented 1 month ago

Since people are using this, best to get it in mainline. So I prefer: fix, merge, & then to iterate testing+refactoring+documenting.

prckent commented 1 month ago

Test this please

ye-luo commented 1 month ago

Since it is a specific flavor of SR, could you make the enum and option name more specific?

prckent commented 1 month ago

@camelto2 I think Ye has a point, but I'll leave it up to you whether to name now or in a later PR. Please leave a comment either way. This capability is obviously evolving and as you point out (thanks) e.g. the CG is worth both refining and abstracting.

ye-luo commented 1 month ago

Test this please