Other inputs to the model than the system configuration

[sunwoo-kim]

Right now, there is no way to provide other parameters to the model (say, RBM), than the spin configuration at each iteration.

I can kind of do it by hand, by creating a new model at each iteration, setting those values, than creating a new vstate, etc., and then doing sgd by hand, but because it is jitted each time, it's very slow.

Is something like this possible?

[PhilipVinc]

I don't exactly understand what you want to do.. could you write it down in formulas or pseudo code?

[sunwoo-kim]

Sorry about the brevity.

Currently, we can only define models with the configurations (ex. spins) as input. Then the variational state interface and the sampler interface samples this function using Monte Carlo.

class MyRBMModel:
    def __call__(self, sigma):
        # do stuff, return log_amplitude
        return f(sigma)

model = MyRBMModel()
vstate = nk.vqs.MCState(model, sampler)
H = sum(sigmaz(hi, i) for in range(L))
E, E_grad = vstate.expect_and_grad(H)

But I would like to be able to calculate the expect_and_grad whilst providing other inputs, like this:

class MyRBMModel:
    def __call__(self, sigma, other_params):
        # do stuff, return log_amplitude
        return f(sigma, other_params)

model = MyRBMModel()
vstate = nk.vqs.MCState(model, sampler)

for i in range(n_iters):
    other_params = 5*i
    H = sum(other_param*sigmaz(hi, j) for j in range(L))
    E, E_grad = vstate.expect_and_grad(H, other_params)

    vstate.parameters = jax.tree_map(...E_grad...)

Right now, what I do is something like:

class MyRBMModel:
    def __call__(self, sigma):
        # do stuff, return log_amplitude
        return f(sigma, self.params)

    def set_params(self, params):
        self.params  = params

para = None

for i in range(n_iters):
    other_params = 0.5*i
    model = MyRBMModel()
    model.set_params(other_params)
    vstate = nk.vqs.MCState(model, sampler)
    H = sum(other_params*sigmaz(hi, j) for j in range(L))
    E, E_grad = vstate.expect_and_grad(H, other_params)

    vstate.parameters = jax.tree_map(...E_grad...)
    para = vstate.parameters.copy()
    del vstate
    del sampler
    del model

If I don't do the delete step, something goes wrong, I think it is due to just-in-time compliation. But it is also extremely slow, compared to a normal loop. I would like to provide input parameters to the model at each iteration, which seems like it's not possible right now.

Let me know if this is still not clear.

[PhilipVinc]

Ah I see! And do you want to compute the derivatives with respect to those other parameters or not?
If this is just some extra variables but you don't want to derivate wrt to it, you can use the model_state interface from flax..

https://flax.readthedocs.io/en/latest/guides/state_params.html

So in your flax model you should define something somehow like

class Jastrow(nn.Module):
    r"""
    Jastrow wave function :math:`\Psi(s) = \exp(\sum_{ij} s_i W_{ij} s_j)`.

    The W matrix is stored as a non-symmetric matrix, and symmetrized
    during computation by doing :code:`W = W + W.T` in the computation.
    """

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

    @nn.compact
    def __call__(self, x_in: Array):
        # standard jastrow Ansatz
        nv = x_in.shape[-1]
        kernel = self.param("kernel", self.kernel_init, (nv, nv), self.param_dtype)
        kernel = kernel + kernel.T
        kernel, x_in = promote_dtype(kernel, x_in, dtype=None)
        y = jnp.einsum("...i,ij,...j", x_in, kernel, x_in)

        # other params
        is_initialized = self.has_variable('other_params', 'something')
        something = self.variable('other_params', 'something', jnp.zeros, ())

        return y ** something

And this parameter will be accessible from MCState.model_state, and changing it won't cause recompilation.

[sunwoo-kim]

Thanks a lot for this, it is definitely much faster than before.

I guess this is expected, but it's still a bit slower than when using driver.run() on the vanilla case, i.e. when I don't change the parameters, or if I implement the loop myself. Could I run some things I am doing by you, to see if there are faster ways of doing them?

Updating the model variable; I couldn't touch them directly, the way I'm able to do it currently is

    something = 5.
    model_state = flax.core.unfreeze(vstate.model_state)
    model_state['other_params']['something'] = something
    vstate.model_state = model_state
    
    E, E_grad = vstate.expect_and_grad(H_g)
   ...

But this freezing and unfreezing business I think perhaps slows things down.

[PhilipVinc]

Could I run some things I am doing by you, to see if there are faster ways of doing them?

sure. Though without a full reproducer is hard to give definitive help...

Updating the model variable; I couldn't touch them directly, the way I'm able to do it currently is

It's the way you are supposed to do it for now.
The only minor thing I have to say is that this way you will be needlessly moving 5 to a jax array every time you execute a function such as expect, expect_and_grad or the qgt. To avoid doing it multiple times and doing it only once, you could just do the following:

    something = jnp.array(5, dtype=float)
    model_state['other_params']['something'] = something
    ...

this way jax will not have to move it to a jax-managed array multiple times, but only once.
Otherwise this looks good.

But this freezing and unfreezing business I think perhaps slows things down.

Not really. For such a simple function like the one I wrote above, and for small systems, yes, it might be (but there is no alternative).
However in any realistic simulation this constant cost will soon disappear in the time you spend evaluating your model/derivatives.