Reusing the contraction path for reparameterized circuits

[khidanov]

Hi! I have a question regarding reusing the contraction path for reparameterized circuits. In my problem I need to calculate expectation values/amplitudes multiple times for the same parameterized circuit but with different sets of parameter values. Since the structure of the tensor network remains the same for these different sets of parameter values, I would think that performing local simplifications and finding the optimal contraction path could be done for one set of parameter values, and reusing the obtained contraction path for other sets of parameter values should speed up the contraction calculation. However, I do not see this to be the case, and I am not sure why.

Here is an illustrative example for a parameterized circuit example from the Quimb tutorials:

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

Here I define some observable which corresponds to some Pauli string:

pauli_str = 'IIZXII'
where = [i for i,p in enumerate(pauli_str) if p!= 'I']
paulis = [p for i,p in enumerate(pauli_str) if p!= 'I']

operator = qu.pauli(paulis[0])
for i in range(1,len(where)):
    operator = operator & qu.pauli(paulis[i]) 

Here is the timing estimates for different things:

n = 6
depth = 9
gate2 = 'CZ'

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

#rehearsing
%time reh = circ.local_expectation_rehearse(operator, where, optimize='greedy',simplify_sequence='ADCRS')

#performing contractions using a path found via rehearsing
%time reh['tn'].contract(all, optimize=reh['tree'], output_inds=())

#performing contractions using a path found via rehearsing
%time circ.local_expectation(operator, where, optimize=reh['tree'],simplify_sequence='ADCRS')

#setting new random parameter values for the circuit
old_params_dict = circ.get_params()
new_params_dict = dict()
for i,key in enumerate(old_params_dict.keys()):
    new_params_dict[key]= qu.randn(3, dist='uniform')
circ.set_params(new_params_dict)

#performing contractions using a path found via rehearsing
%time circ.local_expectation(operator, where, optimize=reh['tree'],simplify_sequence='ADCRS');

The first timing output is for rehearsing the contraction. The second one is for contracting the "rehearsed" TN. The third is for contracting the TN using the obtained contraction path. And the fourth is for contracting a reparameterized TN using the obtained contraction path. On my machine these timings are given as t1≈42ms, t2≈2.6ms, t3≈11ms, t4≈122ms, respectively.

First of all, I am not entirely sure why t2 and t3 are different. My guess is that t3 involves an extra simplification step that has been already done for t2. Second and most importantly, t4 is drastically different from either t2 or t3. However, my understanding is that in principle t4 operation should be able to be performed as fast as t2. That would dramatically speed up my calculations.

Am I missing something here? I am fairly new to Quimb and to the whole world of tensor networks, so I would appreciate if you could provide any input/advice. Thanks!

[jcmgray]

Hi @khidanov, sorry for the slow reply! Yes basically the problem is that the simplification process is data/parameter dependent and so not currently cached. When set_params is called it clears various cached intermediates networks etc, which are then regenerated.

A partial improvement is to use less simplify methods e.g. simplify_sequence ="ADCR", ="R" or ="", though of course for some large circuits these will make the eventual contraction more expensive!

Longer term it is probably possible to trace through the simplifications - as long as the different parameters don't yield different sparsity or low-rank structures - so that the logic doesn't need to be performed each step and even something like jax could be used.

[khidanov]

Hi @jcmgray, thanks for the reply! Changing the simplify_sequence to ="" did slightly speed up my calculation. But yeah, I might need to play with it and see if the speed up would hold for larger circuits.

Could you elaborate what do you mean by "trace through the simplifications" and how it can be done? I do not expect (at least naively) different parameters to yield any difference in the simplification procedure.

[jcmgray]

Tracing, as with jax/torch/autoray, extracts just the array functions into a computational graph that can be called, in the simplest case, on any input arrays with the same shape. However it usually assumes that all one can know about an array is the shape, you can't depend on whats happening inside.

The simplifications in quimb however do depend on whats inside, be it locations of zeros, or whether a gate is low-rank across a particular cut. However, for a random choice of parameters this stuff does not usually change. So one could follow through collecting the operations for an actual set of example arrays including their entries, not just abstract "shapes". Then assume that these operations will be valid for other parameter choices.

[khidanov]

Hi Johnnie, thank you for the reply! I see what you mean here, but do you, by any chance, have examples of how to implement this on practice? Or nudge me in the right direction on how to practically implement this?

Can I do it somehow within Quimb, or would I need to effectively dig into the Quimb code (in particular, the simplification part) and write something new ("in substitute" of the Quimb simplification part) using jax/torch/autoray?