Optimizing a variable loss function

[wkretschmer]

I have an ansatz circuit C(p) parametrized by some parameters p, and a loss function L(C, q) that takes in the circuit C and some auxiliary variable q. I have a long list of values [q[1], ..., q[n]], and for each q[i] I wish to optimize L(C(p), q[i]) over the parameters p. Using TNOptimizer with jax backend, the cost of this optimization is dominated by jit compilation of the gradient. My question is: is there any way to avoid compiling the gradient for each q[i]? In principle, it should be possible to have a single gradient function that takes q as input. The trouble is, if I feed in q as a constant to TNOptimizer, the compiled gradient (seemingly) will not vary at all if I go back and modify q. However, I also don't want to feed in q as a parameter to be optimized, because in any given iteration I'm fixing q to q[i] and only optimizing over p.

In case it's still unclear what I mean, below an example based on https://quimb.readthedocs.io/en/latest/examples/ex_tn_train_circuit.html. The Hamiltonian H changes in each iteration, but the ansatz circuit is always the same. Can I restructure this code so that the gradient compilation only happens once, instead of once per iteration?

import quimb as qu
import quimb.tensor as qtn
import jax

def single_qubit_layer(circ, gate_round=None):
    """Apply a parametrizable layer of single qubit ``U3`` gates.
    """
    for i in range(circ.N):
        # initialize with random parameters
        params = qu.randn(3, dist='uniform')
        circ.apply_gate(
            'U3', *params, i, 
            gate_round=gate_round, parametrize=True)
        
def two_qubit_layer(circ, gate2='CZ', reverse=False, gate_round=None):
    """Apply a layer of constant entangling gates.
    """
    regs = range(0, circ.N - 1)
    if reverse:
        regs = reversed(regs)
    
    for i in regs:
        circ.apply_gate(
            gate2, i, i + 1, gate_round=gate_round)

def ansatz_circuit(n, depth, gate2='CZ', **kwargs):
    """Construct a circuit of single qubit and entangling layers.
    """
    circ = qtn.Circuit(n, **kwargs)
    
    for r in range(depth):
        # single qubit gate layer
        single_qubit_layer(circ, gate_round=r)
        
        # alternate between forward and backward CZ layers
        two_qubit_layer(
            circ, gate2=gate2, gate_round=r, reverse=r % 2 == 0)
        
    # add a final single qubit layer
    single_qubit_layer(circ, gate_round=r + 1)
    
    return circ

n = 6
depth = 9
gate2 = 'CZ'

circ = ansatz_circuit(n, depth, gate2=gate2)

V = circ.uni

for iter in range(10):
    # the hamiltonian
    # ****NOTICE: the parameters of the Hamiltonian vary by iteration****
    H = qu.ham_ising(n, jz=1.0 * iter, bx=0.7 * iter, cyclic=False)

    # the propagator for the hamiltonian
    t = 2
    U_dense = qu.expm(-1j * t * H)

    # 'tensorized' version of the unitary propagator
    U = qtn.Tensor(
        data=U_dense.reshape([2] * (2 * n)),
        inds=[f'k{i}' for i in range(n)] + [f'b{i}' for i in range(n)],
        tags={'U_TARGET'}
    )

    def loss(V, U):
        return 1 - abs((V.H & U).contract(all, optimize='auto-hq')) / 2**n

    print(loss(V, U))

    tnopt = qtn.TNOptimizer(
        V,                        # the tensor network we want to optimize
        loss,                     # the function we want to minimize
        loss_constants={'U': U},  # supply U to the loss function as a constant TN
        tags=['U3'],              # only optimize U3 tensors
        autodiff_backend='jax',   # use 'autograd' for non-compiled optimization
        optimizer='L-BFGS-B',     # the optimization algorithm
    )
    jax.config.update("jax_enable_x64", True)
    V_opt = tnopt.optimize_basinhopping(n=500, nhop=10)

[jcmgray]

Hi @wkretschmer. No I think sadly this kind of 'non-optimized variable' is not supported at the moment. It would be a nice feature to have, especially for jax, which indeed has very slow compilation.

My suggestions would be:

• use a different autodiff_backend or try turning jit_fn=False - might be slow however.
• optimize the expression outside of TNOptimizer - see e.g. https://quimb.readthedocs.io/en/latest/examples/ex_quimb_within_jax_flax_optax.html - maybe with the qtn.pack / qtn.unpack functions
• try and add the feature! I think its probably not too difficult (especially if only supporting the jax backend), if a bit fiddly. Basically one would want to pass something like 'loss_variables' along as kwargs to value_and_grad.