comparing energies

[dambuck]

Hello,

This might be just me overseeing something obvious but here is a quick question. I still had some nk2 code on my cluster because I could not upgrade to nk3 for quite a while. I wanted to do a quick comparison of the two versions just by running heisenberg1d with mpi. Having the same setting for both codes, one in nk2 and the other in nk3. For some reason the old code seems to perform slightly better. Of course it fluctuates a bit depending if the optimization went good or not. But eg. the errors from the expect functions are always bigger in the nk3 case.

for example:
nk3 -0.44382965-0.00010739j \pm 0.00018117
nk2 -0.44470358278178024+3.9935193506671996e-07j) \pm 7.701715053615179e-06

Am I just missing some obvious settings? the nk3 does run faster though, nearly 2x. I wasnt sure if this is due to some settings I missed which might also explain the difference in the results.

import netket as nk
import numpy as np


L = 20
J = [1., 1.]


Sx = np.array([[0, 1], [1, 0]])*0.5
Sy = np.array([[0, -1j], [1j, 0]])*0.5
Sz = np.array([[1, 0], [0, -1]])*0.5

diag=np.kron(Sz, Sz)
offdiag=np.kron(Sx, Sx)+np.kron(Sy, Sy)

edge_colors = []
for i in range(L):
    edge_colors.append([i, (i+1)%L, 1])

g = nk.graph.Graph(n_nodes=L, edges=edge_colors)



bond_operator = [
                (diag+offdiag).tolist(),]
bond_color = [1]

hi = nk.hilbert.Spin(s=0.5, N=g.n_nodes, total_sz=0)

ha= nk.operator.GraphOperator(hi,
               bond_ops=bond_operator, bond_ops_colors=bond_color, graph=g, dtype=complex)


# RBM Spin Machine
ma = nk.models.RBM(
    alpha=1,
    use_visible_bias=False,
    use_hidden_bias=True,
    dtype=complex
)

# Metropolis Local Sampling
sa = nk.sampler.MetropolisExchange(hi, graph=g, n_chains=16)

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.01)
sr = nk.optimizer.SR(diag_shift=0.01)

# Variational State
vs = nk.vqs.MCState(sa, ma, n_samples=1000)

print('PARAMS: ', vs.n_parameters)

# Variational monte carlo driver
gs = nk.VMC(ha, op, variational_state=vs, preconditioner=sr)


print(f"# variational parameters: {gs.state.n_parameters}")


gs.run(n_iter=2000, out="test")


vs.n_samples = 10000
samp=vs.expect(ha)

print('Ener', [samp.mean/L, samp.error_of_mean/L])

J = [1., 1.]

N=20

Sx = np.array([[0, 1], [1, 0]])*0.5
Sy = np.array([[0, -1j], [1j, 0]])*0.5
Sz = np.array([[1, 0], [0, -1]])*0.5

diag=np.kron(Sz, Sz)
offdiag=np.kron(Sx, Sx)+np.kron(Sy, Sy)



edge_colors = []
for i in range(N):
    edge_colors.append([i, (i+1)%N, 1])


g = nk.graph.CustomGraph(edge_colors)

bond_operator = [
    (diag+offdiag).tolist(),
]

bond_color = [1]

# Spin based Hilbert Space
hi = nk.hilbert.Spin(s=0.5, total_sz=0, graph=g)
# Custom Hamiltonian operator
ham = nk.operator.GraphOperator(hi, bondops=bond_operator, bondops_colors=bond_color)

psi = nk.machine.RbmSpin(hilbert=hi, alpha=1,use_visible_bias=False,
    use_hidden_bias=True, )

# Optimizer
op = nk.optimizer.Sgd(learning_rate=0.01)

# Sampler
sa = nk.sampler.MetropolisExchange(machine=psi, n_chains=16)
psi.init_random_parameters(seed=12, sigma=0.01)

# Stochastic reconfiguration
vmc = nk.variational.Vmc(
    hamiltonian=ham,
    sampler=sa,
    optimizer=op,
    n_samples=1000,
    diag_shift=0.01,
    method='Sr')


print('number params:::', len(psi.parameters))


vmc.run(out='test', n_iter=2000)

stat=nk.variational.estimate_expectations(
            ops=ham, sampler=sa, n_samples=10000     )

print(stat.mean/N, stat.error_of_mean/N)
                                                                                                                                                                                         95,1          Bot

[gcarleo]

if you find a better variational state (i.e. with lower energy) then also the energy variance will be smaller, thus the statistical error will be smaller (at parity of number of samples), so this aspect is not too strange I think.
Assuming you are using exactly the same models in nk2 and 3 (complex parameters, same choice of hidden units, hidden/visible bias etc) then I believe the only main difference might be the way parameters are initialized. If I remember correctly, in nk2 the initializers had a larger sigma (there was a "bug" and we were using variance instead of standard deviation, making the initialization larger than expected). I would suggest you try increasing the sigma in the random inizialization of the parameters in nk3

[PhilipVinc]

Exactly. In the same line, the default initializer in nk3 follows flax, therefore the bias is zero-initialised while in nk2 all parameters where a gaussian. To reproduce the same behaviour of nk2 you should run vs.init_parameters(nn.initializers.normal(0.01)) (and be careful in that nk2 had that bug so you need to square the value in one of the two).

The majority of the time the issue is the variance used in nk3, which is too small.

Also, another difference is the solver used for SR: in nk2 i think svd was used by default, while nk3 uses cg unless you specify another solver.

It would be interesting to see a plot comparing the two oiptimisations once you change the initialiser.