NaN during training about ferminet HOT 8 CLOSED

from ferminet.

from ferminet.

Thanks for the quick response!
I could work around this issue by modifying the logsumdet function.
The following should work as a drop-in replacement:

def logsumdet(
        xs: Sequence[jnp.ndarray],
        w: Optional[jnp.ndarray] = None) -> jnp.ndarray:
    # Special case if there is only one electron in any channel
    # We can avoid the log(0) issue by not going into the log domain
    dets = [x.reshape(-1) for x in xs if x.shape[-1] == 1]
    dets = functools.reduce(
        lambda a, b: a*b, dets
    ) if len(dets) > 0 else 1

    slogdets = [slogdet(x) for x in xs if x.shape[-1] > 1]
    maxlogdet = 0
    if len(slogdets) > 0:
        sign_in, logdet = functools.reduce(
            lambda a, b: (a[0]*b[0], a[1]+b[1]), slogdets
        )

        maxlogdet = jnp.max(logdet)
        det = sign_in * dets * jnp.exp(logdet - maxlogdet)
    else:
        det = dets

    if w is None:
        result = jnp.sum(det)
    else:
        result = jnp.dot(det, w)

    sign_out = jnp.sign(result)
    log_out = jnp.log(jnp.abs(result)) + maxlogdet
    return sign_out, log_out

This avoids working in the log domain if it's not necessary.

from ferminet.

from ferminet.

I'm curious about your technical set up (jax, jaxlib, and CUDA versions). The hydrogen atom works fine for me (several thousand iterations), though I agree the jax version is not as careful around this as the TF version.

The energy looks essentially converged by this point so I'm surprised you're hitting a NaN. Given the wavefunction shouldn't(!) contain any nodes, it can only happen far from the nucleus, so an MCMC configuration must have wandered quite a way away. (As an aside, the default step size is way too small for hydrogen -- you should increase it to at least 0.5).

from ferminet.

Sure, here are the details:

While the mean of the energy looked good, it wasn't converged yet (high variance). I checked the position of the electron and it was in a radius of 3A around the nucleus which isn't very afaik.
The error also occurs at higher step sizes.

from ferminet.

This error also occurs for Lithium far from the optimum.

I0306 02:47:37.867998 139629635032896 train.py:461] Step 00093: -6.6119 E_h, pmove=0.76
I0306 02:47:38.068941 139629635032896 train.py:461] Step 00094: -6.6105 E_h, pmove=0.76
I0306 02:47:38.270185 139629635032896 train.py:461] Step 00095: -6.6636 E_h, pmove=0.76
I0306 02:47:38.471880 139629635032896 train.py:461] Step 00096: nan E_h, pmove=0.76
I0306 02:47:38.671543 139629635032896 train.py:461] Step 00097: nan E_h, pmove=0.00
I0306 02:47:38.870149 139629635032896 train.py:461] Step 00098: nan E_h, pmove=0.00

Though, only if one sets full_det to False. As far as I can tell #23 fixes this.

On a side note: Is there a particular reason why full_det defaults to True? Isn't a wavefunction only antisymmetric with respect to permutation of electrons of the same spin? Also, it does not align with the definition of FermiNet in the papers.

from ferminet.

Fixed in #23 .

from ferminet.