[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?