Introduction

The diagonal of the Hessian contains the second derivative of a scalar objective with respect to each variable individually. In TensorFlow 2, the cleanest way to compute it is usually with nested GradientTape objects: first compute the gradient, then differentiate that gradient again to recover the Hessian or at least its diagonal.

What the Hessian Diagonal Means

For a scalar function of a vector input, the Hessian is the matrix of all second partial derivatives. The diagonal entries are the terms that look like "second derivative with respect to the same variable twice."

Those values are often useful when you care about curvature per parameter but do not need the full dense Hessian.

The main reason people ask for only the diagonal is cost. A full Hessian is much more expensive than a gradient.

Nested GradientTape for a Small Example

For a small vector-valued input, you can compute the full Hessian and then extract the diagonal.

1import tensorflow as tf
2
3x = tf.Variable([1.0, 2.0, 3.0])
4
5with tf.GradientTape() as outer_tape:
6    with tf.GradientTape() as inner_tape:
7        y = tf.reduce_sum(x ** 3)
8    grad = inner_tape.gradient(y, x)
9
10hessian = outer_tape.jacobian(grad, x)
11diag = tf.linalg.diag_part(hessian)
12
13print("gradient:", grad.numpy())
14print("hessian diagonal:", diag.numpy())

For the function sum(x^3), the second derivative with respect to each component is 6x, so the diagonal should come out as [6, 12, 18] for the example input.

Why This Works

The inner tape computes the first derivative of the scalar loss with respect to x. The outer tape then takes the Jacobian of that gradient vector with respect to x again.

That Jacobian of the gradient is the Hessian.

Once you have the full Hessian tensor, tf.linalg.diag_part extracts the diagonal efficiently.

When You Only Need the Diagonal

For higher-dimensional parameter vectors, building the full Hessian can become expensive in both memory and time. If you only need the diagonal, you should think carefully about whether computing the full Hessian first is acceptable.

For small examples and teaching code, it is fine. For larger models, it can become impractical quickly.

That is why people often use approximations, Hessian-vector products, or more specialized curvature methods in real training systems.

Practical Considerations

To make this work smoothly:

• the objective should be scalar,
• the watched variables should be differentiable tensors,
• and the operations in the objective must support second-order differentiation.

If some operation in the graph is not differentiable twice, the second derivative path can return None or fail to produce the expected result.

It is also worth remembering that TensorFlow 2 runs eagerly by default, which makes these small autodiff experiments much easier to inspect than old graph-mode code.

Common Pitfalls

• Forgetting that the Hessian is defined for a scalar objective, not an arbitrary vector output.
• Trying to compute a full Hessian for a very large parameter vector and running into memory problems.
• Expecting second derivatives to work through operations that are not twice differentiable.
• Mixing up the gradient vector with the Hessian diagonal itself.
• Using TensorFlow 2 but thinking in old graph-mode patterns instead of using GradientTape directly.

Summary

• In TensorFlow 2, the Hessian diagonal is commonly computed with nested GradientTape objects.
• Compute the gradient first, then take its Jacobian with respect to the original variables.
• Extract the diagonal with tf.linalg.diag_part.
• This is practical for small inputs and demonstrations, but full Hessian construction gets expensive quickly.
• If you only need curvature information at scale, think carefully before materializing the full Hessian.