Numerical stability during autodifferentiation of eigh about graddft HOT 9 CLOSED

I am investigating whether this google/jax#5461 (comment) could be a simpler yet effective solution. It seems to allow for 2 training iterations before breaking down, (having set up scf iterations = 2). However, it returns erroneous eigenvalues, as it ends up not conserving the charge.

from graddft.

I have created an example to play around with this problem
https://github.com/XanaduAI/DiffDFT/blob/main/examples/basic_examples/example_neural_scf_training.py
It should be deleted if we can't solve it on time. Changing the number of diis iterations in https://github.com/XanaduAI/DiffDFT/blob/9bbab42cfe95bbb6031c45613b5ab718ca9e9e5f/grad_dft/evaluate.py#L524
also has some minor effect.

from graddft.

I have realized that https://gist.github.com/jackd/99e012090a56637b8dd8bb037374900e provides a definition of the derivatives by hand which is probably being used https://github.com/XanaduAI/DiffDFT/blob/9bbab42cfe95bbb6031c45613b5ab718ca9e9e5f/grad_dft/external/eigh_impl.py#L84

from graddft.

I'm going to experiment a bit with this now, but I think we can just do:

import jax.numpy as jnp
def generalized_eigh(A, B):
    L = jnp.linalg.cholesky(B)
    L_inv = jnp.linalg.inv(L)
    A_redo = L_inv.dot(A).dot(L_inv.T)
    return jnp.linalg.eigh( A_redo )

from graddft.

I'm going to experiment a bit with this now, but I think we can just do:

import jax.numpy as jnp
def generalized_eigh(A, B):
    L = jnp.linalg.cholesky(B)
    L_inv = jnp.linalg.inv(L)
    A_redo = L_inv.dot(A).dot(L_inv.T)
    return jnp.linalg.eigh( A_redo )

Ah sorry this is the solution you already posted. Let me see if there is a better way of doing this...

from graddft.

Ok @PabloAMC I think this code fixes the problem:

def generalized_to_standard_eig(A, B):
    L = np.linalg.cholesky(B)
    L_inv = np.linalg.inv(L)
    C = L_inv @ A @ L_inv.T
    eigenvalues, eigenvectors_transformed = np.linalg.eigh(C)
    eigenvectors_original = L_inv.T @ eigenvectors_transformed
    return eigenvalues, eigenvectors_original

Basically, the eigenvectors were not transformed back to the original basis.

from graddft.

You are right, that works, thanks @jackbaker1001. Unfortunately, it still gives errors when backpropagating even when using this version of the generalized eigenvalue problem. It still gives nan errors when I set cycles = 1 in the additional example https://github.com/XanaduAI/DiffDFT/blob/main/examples/basic_examples/example_neural_scf_training.py I included in the basic folder and change the number of cycles to 1, instead of 0, in https://github.com/XanaduAI/DiffDFT/blob/81444ce067a809e3e69bffdf81b4e7c1b5ca4d2b/examples/basic_examples/example_neural_scf_training.py#L138

from graddft.

Ok @PabloAMC so because this issue was about a differentiable eigh implementation, which we now have, I will put in a PR for this and merge. I will keep the custom eigh implementation from before in the code just in case in comes in handy later. It may be removed later also.

I will make a new issue regarding differentiating through the SCF iterator (different implementations of DIIS presently) as this is now the problem!

from graddft.

I think we should reopen this. The problem is numerical stability during the calculation of grads, not a lack of a differentiable implementation. And that does not seem solved yet.

from graddft.