Error in optimization with gate_split

[rafaeleb]

Hello all,

I am currently implementing an optimization procedure for an MPS, where the loss function is written as

`def loss(psi):

all_terms = [] 
#kinetic terms
for i in range(l):
    phi = cp.copy(psi)
    all_terms.append(phi.gate(jnp.array(0.5 * D2), l, contract = True))

#potential terms
for i in range(l):
    phi = cp.copy(psi)
    all_terms.append(phi.gate(jnp.array(0.5 * X2), l, contract = True))

#two-body interactions
for i in range(l-1):
    phi = cp.copy(psi)
    all_terms.append(phi.gate_split(jnp.array(g*Xx), (i, i+1)))

final_psi = np.sum(np.array(all_terms))
return psi.H @ final_psi #returns the expectation value of the energy <psi|H|psi>`

works works for an individual calculation of the loss, but when running the optimizer

tnopt = TNOptimizer(psi, loss_fn = loss, norm_fn = norm_fn, optimizer = "adam", callback = callback, autodiff_backend = 'jax')
I get an error from JAX (just pasting the last lines):

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/0)> The problem arose with the boolfunction.

I could not find the problematic function but it's related to the application of .gate_split, since the rest of the loss function does not raise errors. Any way I can solve this issue?

Thanks,

Rafael

[jcmgray]

Hi @rafaeleb, you might be able to fix this by using a fixed max_bond and cutoff=0.0 in the gate_split function. The problem currently is that the shape of the traced computational graph depends on the cutoff. Also, you should be able to use torch which doesn't require a fixed computational graph.

[rafaeleb]

Hi Johnnie,

Thanks for the reply. The addition of the compression options did not fix the problem, but it looks like it can be avoided in my particular case by not using gate_split at all (the gates can be decomposed from the start). But it would be nice to know if the gate_split function can be modified is some way to agree with the jax optimizer backend.

BR

[EmilianoG-byte]

Hi! I am having a somewhat similar problem where I am trying to calculate the jax.grad of a scalar function that takes as an input a jax array of parameters to be used in a tensor network.

In my case, since I am not trying to JIT anything, it seems like there is a problem when calculating the gradient, as this returns all nan. When using the jax.debug_nans() it shows me that the error is caused by the truncation using SVD and it shows it as FloatingPointError when calling:

U, s, VH = do("linalg.svd", x)

in quimb/tensor/decomp.py:223.

I tried setting the cuttoff=0.0 and a certain max_bond as you suggested in other posts but that didn't solve the issue.

Any leads would be very appreciated @jcmgray. I can also open a new issue if you reckon this is too different.

[jcmgray]

So a known problem with the way the gradient of the SVD is typically computed is that it depends on the inverse of the singular value differences, I wonder if your matrix has a degenerate spectrum?

[EmilianoG-byte]

Thank you for the fast answer! That's a good point. I now checked first that all the singular values are unique along all bonds of the MPS as:

[len(jnp.unique(evolved_mps.singular_values(i=i))) == len(evolved_mps.singular_values(i=i)) for i in range(1, num_qubits)]

And I got all of them to be true.

I also tried different compression methods for tensor_split, such as eig and lu, which also produce nan and then iterative methods like 'svds' and isvd which fail due to the use of np.asarray(a) on jax traced arrays (quimb/quimb/linalg/rand_linalg.py:18).

Is there anything else I could try? I would have thought that differentiating through compression of TN is a relatively standard procedure but I haven't found any helpful answer so far, unfortunately.

[jcmgray]

It might be worth checking the spectrum at the point the svd is called, not just of the MPS tensors, and also checking for only approximate degneracy, e.g. up to ~1e-6 for single precision, which jax defaults to. Beyond that, what is the norm of you MPS? This blowing up is a common cause of nan.

[EmilianoG-byte]

You were absolutely right about the degeneracy! I have tested a smaller system size (N=10) to better track the singualar values, and there is indeed a lot of degeneracy across different partitions.

I am now looking into a way to still truncate the bond dimension (so not QR unfortunately), while avoiding this differentiatiblity issue with degenerate spectrum. I saw here: jax-ml/jax#2311 there is a "dirty fix", but implies going deeper into the stack of jax so I am still figuring out if this is the easiest way to go.

Any other leads are welcome!

[jcmgray]

You could try using single site fitting with a fixed number of sweeps to compress the MPS. See https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/tensor_1d_compress/index.html#quimb.tensor.tensor_1d_compress.tensor_network_1d_compress_fit. It require supplying a MPS with jax generated arrays (to be jittable) as the initial guess if zipup doesn't work (due to the same svd problems).