Not listed ".Gaussian" model?

[Daniel-Haas-B]

Hey everyone,
Maybe this is a simple question but I have just started using NetKet and also applying ML to quantum problems, so here it goes.

In this example of finding the ground-state energy of a harmonic oscillator, the model ´nk.models.Gaussian(param_dtype=float)´ slightly confuses me. This model is not listed amongst the other models in the documentation, and therefore the only explanation we get for it is from the class docstring:

class Gaussian(nn.Module):
    r"""
    Multivariate Gaussian function with mean 0 and parametrised covariance matrix
    :math:`\Sigma_{ij}`.

    The wavefunction is given by the formula: :math:`\Psi(x) = \exp(\sum_{ij} x_i \Sigma_{ij} x_j)`.
    The (positive definite) :math:`\Sigma_{ij} = AA^T` matrix is stored as
    non-positive definite matrix A.
    """

    param_dtype: DType = jnp.float64
    """The dtype of the weights."""
    kernel_init: NNInitFunc = normal(stddev=1.0)
    """Initializer for the weights."""

    @nn.compact
    def __call__(self, x_in: Array):
        nv = x_in.shape[-1]

        kernel = self.param("kernel", self.kernel_init, (nv, nv), self.param_dtype)
        kernel = jnp.dot(kernel.T, kernel)

        kernel, x_in = promote_dtype(kernel, x_in, dtype=None)
        y = -0.5 * jnp.einsum("...i,ij,...j", x_in, kernel, x_in)

        return y

Again, I am very new to all this, so I don`t really get what is going on. Is it something similar to an RBM? If not, how does it compare? What are we really optimizing from the trial wave function here?

Furthermore, why is this model not listed? Should this be an issue? I appreciate any guidance about this.
Cheers.

[PhilipVinc]

Hi Daniel,

The fact that the model is not listed in the documentation is an oversight (we need to add its name to a list used to generate the website docs and it's easy to forget to do it..).

Besides, the documentation is generated from the docstrings, so it contains the same information.

About the model: this is literally just a multivariate Gaussian. So it corresponds to

$$ \psi(x_1,\dots,x_n) = \exp(\sum_{i,j} x_i K_{i,j} x_j) $$

And you optimize the parameters $A_{i,j}$ such that $K_{ij}=\sum_l A_{il}A{lj}$(technically, we could directly optimize the weights K, but then K would not automatically be positive semi definite and so we opt to parametrise its square root A).

This model is very different from an rbm. first, it has no nonlinearity, so its very simple and you should not expect to work well beyond extremely simple problems. Second, its particularly geared for systems with continuous degrees of freedom instead of spins.