XanaduAI / GradDFT

GradDFT is a JAX-based library enabling the differentiable design and experimentation of exchange-correlation functionals using machine learning techniques.
Apache License 2.0
79 stars 7 forks source link

The return of the NaN #72

Open jackbaker1001 opened 1 year ago

jackbaker1001 commented 1 year ago

I've noticed that for a few basis sets, NaN gradients are appearing again when trained using the DIIS SCF loops but not the linear mixing loops.

I think this is likely because degenerate eigenvectors/eigenvalues are being encountered in the calls to jnp.linalg.eigh routines. We can try switching these out for our custom safe_eigh which will hopefully fix the bug.

jackbaker1001 commented 1 year ago

Changing to safe_eigh did not fix things! I'm seeing this with diatomic molecules using 6-31G bases. I.e,

Na2_bl = 3.0789
Na2_mol =  pyscf.M(
        atom = 'Na %.5f 0.0 0.0; Na %.5f 0.0 0.0' % (-Na2_bl/2, Na2_bl/2),
        basis = '6-31G'
)

and gd.energy_and_density_loss

jackbaker1001 commented 1 year ago

although I have now noticed that it also occurs in gd.mse_energy_loss

jackbaker1001 commented 1 year ago

And it is not just the gradients which are NaN, it is also the loss itself!

PabloAMC commented 1 year ago

@jackbaker1001 I can't replicate the issue with the above molecule in the 4th basic example script. The only change I made was including

from .eigenproblem import safe_fock_solver, safe_eigh

in the utils.__init__.py, as safe_eigh would not import otherwise...

jackbaker1001 commented 1 year ago

If you made clips which resolved NaN's here @PabloAMC can you push them to the branch so I can test on my end?

PabloAMC commented 1 year ago

@jackbaker1001 test the new branch I created https://github.com/XanaduAI/GradDFT/tree/72b-the-return-of-the-nan. The only change is the clip constant lowering to 1e-25. I tested a few extra basis sets, but they seem to work fine.

jackbaker1001 commented 1 year ago

@PabloAMC I still get NaNs on this branch running the self-consistent training in the notebook ~examples/intermediate_notebooks/training_methods_03.ipynb when switching the basis sets for all molecules to 6-31G.

Can you tell me if this happens in your local install please?

Screenshot 2023-09-25 at 2 46 32 PM

PabloAMC commented 1 year ago

Did lowering the clipping constant help? I can't replicate the issue on my end with the current one.