Hessian calculation fails using tf.jacobian

[therooler]

Python: 3.8
TFQ: 0.4.0
Hi,

I am trying to get some Hessians from the one of my parameterized quantum circuits. This trainstep works as intended:

def build_train_step(circuit: cirq.Circuit, symbols: List,
                     paulisum: List[cirq.PauliSum], learning_rate: float) -> \
        Tuple[Any, tf.Variable, tfq.layers.Expectation]:
    model_params = tf.Variable(tf.random.uniform([1, len(symbols)]) * 2,
                               constraint=lambda x: tf.clip_by_value(x, 0, 4))
    expectation_layer = tfq.layers.Expectation()
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

    @tf.function
    def train_step():
        with tf.GradientTape() as tape:
            expectation_batch = expectation_layer(circuit,
                                                  symbol_names=symbols,
                                                  symbol_values=model_params,
                                                  operators=paulisum)
            energy = tf.reduce_sum(expectation_batch)
        gradients = tape.gradient(energy, model_params)
        optimizer.apply_gradients(zip([gradients], [model_params]))
        return energy

    return train_step, model_params, expectation_layer

Following this example in the TensorFlow 2 docs I was hoping I could get the Hessian with the following code:

def build_train_step_hessians(circuit: cirq.Circuit, symbols: List,
                     paulisum: List[cirq.PauliSum], learning_rate: float) -> \
        Tuple[Any, tf.Variable, tfq.layers.Expectation]:
    model_params = tf.Variable(tf.random.uniform([1, len(symbols)]) * 2,
                               constraint=lambda x: tf.clip_by_value(x, 0, 4))
    expectation_layer = tfq.layers.Expectation()
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

    @tf.function
    def train_step():
        with tf.GradientTape() as t2:
            with tf.GradientTape() as t1:
                expectation_batch = expectation_layer(circuit,
                                                      symbol_names=symbols,
                                                      symbol_values=model_params,
                                                      operators=paulisum)
                energy = tf.reduce_sum(expectation_batch)
            gradients = t1.gradient(energy, model_params)
        hess = t2.jacobian(gradients, model_params)
        optimizer.apply_gradients(zip([gradients], [model_params]))

        return energy, hess

    return train_step, model_params, expectation_layer

But this throws the error:

...
    LookupError: No gradient defined for operation 'TfqAdjointGradient' (op type: TfqAdjointGradient)

From which I conclude that calculating gradients of gradients is not supported yet ( I tried the other differentiators as well). Am I out of luck here? Or is there a hack I can use to get the Hessians from the circuit? Thanks! If you need an example where I use this train step I can throw one together.

P.S.
I am in the process of rewriting all my research code to TFQ and so far everything has worked like a charm. No more super slow graph building times and worrying about how to extract stuff the graph with my own TF1 simulator. And the adjoint differentiator in TFQ is amazing as well; I ran a VQE optimization with like 500 parameters the other day without any issues. Great stuff!

[gillverd]

Hi Roeland, it's great to hear that your research is benefitting from TFQ!

As for the Hessians of circuits, it is likely that we will add a function to implement the parameter-shift-based Hessian (as in e.g. this work by the Coles group ). Really, this is just applying recursively the parameter shift gradient update rule twice. @zaqqwerty is working on such implementations for second and higher-order gradients and may have a temporary workaround leveraging some of the recent methods for exposing the innards of TFQ differentiators.

[phyjoon]

Hi, I also want to compute the Hessian of parameterized circuits using TFQ. Does anyone have a workaround solution? Thanks!

[mikeevmm]

Hello all,

I'm hitting this error but only when running in graph mode.

What's strange is that I'm only asking for gradients (not a Jacobian), and calling the circuit (part of a model) with tf.py_function (so as to run the model in eager mode), so I don't understand where the error is coming from. Is there a way to side-step this?

def eager_model_eval(inputs, params):
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(params)
        # `model` and `tan_loss` are defined out of scope
        overlap = model([inputs,params])
        loss = tan_loss(overlap, [1.])
    gradients = tape.gradient(loss, [params])
    return loss, gradients

# ... Later in graph mode ...

eager_model_result = tf.py_function(
                func=eager_model_eval,
                inp=[inputs, params],
                Tout=(tf.float32, tf.float32))

(My guess is that a differentiation of the gradient is being attempted, but only in graph mode for some reason?)

Thanks in advance.

[born-2learn]

Hi @therooler @zaqqwerty @QuantumVerd

May I know if you were able to find a solution to calculate Jacobians/Hessians from Parameterized Quantum circuits? I'm facing a similar problem and it would be appreciated if you could provide some insights! Thanks.