[Discussion] Adding quantum math features

[MichaelBroughton]

Sometimes, it is useful to be able to analyze circuits in ways that aren't valid on a quantum computer, but might still be interesting for theory reasons. I came across this issue a little while ago when wanting to compute the overlap <psi_1 | psi_2> of a state and had a hard time keeping memory limits under control with larger batch and state sizes using tfq.layers.State. I ultimately ended up simulating the circuits one by one, and doing the inner products in tf manually. It worked, but it wasn't pretty.

Would people be interested in implementing things like:

Examples that can be done right now, but are tricky to get right without blowing up memory:

tfq.math.overlap # => analytically compute state overlap between the two circuits
tfq.math.amplitudes # => analytically compute requested state amplitudes (could use qsimh ?).
tfq.math.fideltiy # => analytically compute fidelity between two circuits.

Other interesting math functions might be worth having

tfq.math.operator_matrix # => ingest a cirq.PauliSum and compute the matrix
tfq.math.operator_commutators # => compute commutators on operator matrices

Does anyone else have any suggestions for useful math function ? @dabacon @zaqqwerty @jaeyoo @we-taper

Final thought: If we do end up making this .math module , maybe we should move calculate_unitary and calculate_state inside of it.

[jaeyoo]

tfq.math.fubini_study_metric # -> used similar to fidelity
tfq.math.fisher_information_matrix # -> used in calculating Hessian
tfq.math.relative_entropy # -> used in measure between two mixed states.

[therooler]

To add to my comments in the TFQ sync meeting about how useful a tfq.math.fubini_study_metric would be:

This is what I had to do in TensorFlow 1.x to calculate the Quantum Natural Gradient:

def QuantumNaturalGradient(state, loss, vrs, optimizer=None, stability_shift=None,
 learning_rate: float = 0.01):
    """
    Implementation of the quantum natural gradient gradient method

    Args:
        *state (Tensor)*:
            Output state of the circuit.

        *loss (Tensor)*:
            Loss function that uses the provided state.

        *vrs (Variables)*:
            Variables to be optimized.

        *optimizer (Tensorflow Optimizer)*:
            Tensorflow optimizer from the tf.train API

    Returns (Operation):
        Train step operation.

    """
    if optimizer == None:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate)
    assign_ops = []
    for variable in vrs:
        variable = [variable]
        jac, nparams = prepGeometricTensor(state, variable)

        grads = tf.gradients(loss, variable)
        sf_metric = []

        # fubini-study metric
        for i in range(nparams):
            for j in range(nparams):
                part_1 = tf.math.conj(tf.reshape(jac[i], (1, -1))) @ tf.reshape(jac[j], (-1, 1))
                part_2 = tf.math.conj(tf.reshape(jac[i], (1, -1))) @ tf.reshape(state, (-1, 1)) + \
                         tf.math.conj(tf.reshape(state, (1, -1))) @ tf.reshape(jac[j], (-1, 1))
                sf_metric.append(part_1 - part_2)
        eta = tf.math.real(tf.reshape(tf.stack(sf_metric), (nparams, nparams)))

        # QNG
        grads = tf.stack(grads)
        if stability_shift is not None:
            eta += tf.eye(*eta.shape.as_list()) * stability_shift
        grads = tf.linalg.solve(eta, tf.reshape(grads, (-1, 1)))
        if len(variable) == 1:
            grads = tf.reshape(grads, (variable[0].shape))
            assign_ops.append(optimizer.apply_gradients(zip([grads], variable)))
        else:
            grads = tf.split(grads, nparams, axis=0)
            assign_ops.append(optimizer.apply_gradients(zip(grads, variable)))
    return assign_ops

It resulted in a huge compute graph that was crazy slow and unstable since inverting eta can be ill-defined. However, for smaller systems this worked and sped up the optimization of a variational circuit significantly (as in https://arxiv.org/abs/1909.02108). If you leave out part_2 of the FS metric, you get Imaginary Time Evolution: https://arxiv.org/abs/1804.03023.

Having an efficient, parallelizable tfq.math.fubini_study_metric instead would be awesome.

[MichaelBroughton]

Since our discussions we have had a number of people also request the fubini study metric to explore the natural gradient. Perhaps once we make the op we should also make a NaturalGradient differentiator ?

[MichaelBroughton]

Was also thinking this might be a good place to add MPS based simulation algorithms with new ops like:

tfq.math.mps_expectation # = > computes the expectation value of a 1d cirq.Circuit.
tfq.math.mps_state # => produce the associated mps tensor
tfq.math.mps_samples #=> produce bitstrings from mps simulator.

[SatyaKuppam]

Hey can I work on this?

I just want to summarize these are the math ops we want to implement ( I have ordered in terms of increasing difficulty, or at least thats what I think).

tfq.math.mps_expectation # = > computes the expectation value of a 1d cirq.Circuit.
tfq.math.mps_state # => produce the associated mps tensor
tfq.math.mps_samples #=> produce bitstrings from mps simulator.
tfq.math.amplitudes # => analytically compute requested state amplitudes (could use qsimh ?).
tfq.math.operator_matrix # => ingest a cirq.PauliSum and compute the matrix
tfq.math.operator_commutators # => compute commutators on operator matrices
tfq.math.fideltiy # => analytically compute fidelity between two circuits.
tfq.math.fubini_study_metric # -> used similar to fidelity
tfq.math.fisher_information_matrix # -> used in calculating Hessian
tfq.math.relative_entropy # -> used in measure between two mixed states.

tfq.math.overlap mentioned above is already implemented as inner product, will use that as reference.

[MichaelBroughton]

Absolutely go for it!

[nishant34]

Since our discussions we have had a number of people also request the fubini study metric to explore the natural gradient. Perhaps once we make the op we should also make a NaturalGradient differentiator ?

Is there any update regarding its status? I also wanted to work on natural and hessian gradients.

[MichaelBroughton]

I think this is something @jaeyoo might be interested in taking on. If he is we should probably break some of these ops off into their own issues and assign people ownership, just so we know who's working on what.

[jaeyoo]

Hi @nishant34 I am already working on natural gradients. Would you mind if you are working on natural gradients optimizer or hessian gradients?

[SatyaKuppam]

Sorry was busy with other stuff, will start working on this now. How do you want me to break this up into multiple issues?

[MichaelBroughton]

I'd say a good approach would be this:

1. Open an issue titled "[Math] Implement tfq.math.whatever op you pick" with the basic steps you envision for implementing it.
2. Assign it to yourself and put a reference back to this issue in the description.
3. Add any needed support code in small PRs before the main op contribution PR.
4. Create the op PR with similar structure to something like #508 or #387 .
5. Once merged we can cross the op off the list on here.

Since op contributions aren't coming in very quickly and are a fairly sizable undertaking, I think anyone looking to contribute a new op should definitely take this "one at a time issue creation followed by PR chain" approach.

[MichaelBroughton]

Similar to MPS simulation we might want to also support clifford circuit simulation. Likely built on top of: https://github.com/quantumlib/Stim. Since expectation values aren't as useful in clifford simulation we'd probably only need:

tfq.math.stabilizer_table      # return stabilizer table/tableu.
tfq.math.stabilizer_sample   # sample from batch of stabilizer circuits

[SatyaKuppam]

We currently have MPS simulator functionality in cirq wondering if I could leverage that for MPS related OPS instead of doing any C++.

[MichaelBroughton]

That's an interesting idea and using python only can in theory work (like we did cirq_ops.py), but it incurs a huge overhead compared to running C++ code. If a PR was opened for implementing these ops in python we wouldn't merge it. We need to make sure that our users can feel confident that whatever ops they are calling into, they can rely on them being very fast.

[SatyaKuppam]

Gotcha, just checking. Will get started with the C++ implementation. Will follow the steps you listed above. Thanks.

[MichaelBroughton]

Looks like we had an implementation of fidelity go in here #554 from @jaeyoo so we can cross it off the list here.

[SatyaKuppam]

Sorry, I was unwell and had to take a couple of months rest, just checking in to say I am still pursuing this.