Preparing exited states with Netket

[jackaraz]

Hi, I've been trying to prepare excited states for a certain model with NetKet. I came across various examples of code implementing penalty-based state preparation, such as the code attached to this discussion (as well as corresponding PR), fidelity implementation here and here, etc, and I tried to prepare a code out of these. For most cases, although I'm obtaining reasonable energy values for both ground and excited states, it turns out that my fidelity (computed without VMC) is significantly off for the excited state (perfect for the ground state).

So I came across 2308.16848, where they are taking the Slater determinant of a larger Hilbert space to ensure ground and excited states are orthogonal (eq 1-2 and 15).

As far as I understand, I can implement it in Netket by doubling my Hamiltonian, where two chunks are independent from each other and acting on different portions of the general Hilbert space, which is doubled as well:

hilbert_space = SpinOrbitalFermions(norbitals, n_fermions=nfermions) * hilbert_space = SpinOrbitalFermions(norbitals, n_fermions=nfermions)
c = lambda site: destroy(hilbert_space, site)
cdag = lambda site: create(hilbert_space, site)
hamiltonian=0.
for shift in range(0, norbitals + 1, norbitals):
    for isite in range(shift, norbitals+shift-1):
          q1 = isite
          q2 = isite + 1
          hamiltonian -= (0.5j) * (cdag(q1) @ c(q2) - cdag(q2) @ c(q1))

In order to take care of the two Hilbert spaces separately, I created the following model

class LogNeuralBackflowDouble(nnx.Module):
    def __init__(
        self,
        hilbert,
        hidden_units: list[int],
        *,
        rngs: nnx.Rngs,
    ):
        key1, key2 = jr.split(rngs.params())

        self.model1 = ...
        self.model2 = ...

    def __call__(self, n: jax.Array) -> jax.Array:
        hdim = jnp.shape(n)[-1]
        n1 = n[:, : hdim // 2]
        n2 = n[:, hdim // 2 :]

        psi1_n1 = jnp.expand_dims(jnp.exp(self.model1(n1)), 1)
        psi1_n2 = jnp.expand_dims(jnp.exp(self.model1(n2)), 1)
        psi2_n1 = jnp.expand_dims(jnp.exp(self.model2(n1)), 1)
        psi2_n2 = jnp.expand_dims(jnp.exp(self.model2(n2)), 1)

        @partial(jnp.vectorize, signature="(n),(n),(n),(n)->()")
        def log_sd(psi1_n1, psi1_n2, psi2_n1, psi2_n2):
            h_full = jnp.block([[psi1_n1, psi2_n1], [psi1_n2, psi2_n2]])
            return nk.jax.logdet_cmplx(h_full)

        return log_sd(psi1_n1, psi1_n2, psi2_n1, psi2_n2)

where the model takes only half of the Hilbert space, and each model is essentially a backflow network with different parameters.

I followed the standard training procedure with the doubled Hilbert space and split the two models into two separate MC states after training, ensuring that each new MC state contains half of the sampler_state.σ. While this produces a perfect ground state, I found that the excited state is again quite off, though the two states are not identical. In fact, the overlap between the two states is quite similar to the results I get with the penalty-based implementation.

It seems I am unable to ensure orthogonality between two states, and I cannot find a way to force it either. Do you, by any chance, have an implementation for finding excited states that I could look into? Or could you point me towards a resource that I can implement?

Thanks
Best regards

[PhilipVinc]

@Adrien-Kahn by any chance do you have a code to do the determinant state ground state optimization using G and Gh, or you only have the code for bridge?