Introduction

The Hessian matrix captures second-order curvature information for a scalar function, but many applications only need its diagonal entries. In TensorFlow, that usually means computing second derivatives of each gradient component with respect to the matching input variable, without spending memory on the entire dense Hessian if you do not need it.

Understand What the Diagonal Means

For a scalar function f(x) where x is a vector, the Hessian diagonal contains:

• 'd2f / dx0^2'
• 'd2f / dx1^2'
• and so on

Those terms tell you how curvature behaves along each coordinate independently. They are useful in optimization diagnostics, uncertainty approximations, and diagonal preconditioning methods.

TensorFlow's automatic differentiation makes this possible with nested GradientTape objects.

Compute the Diagonal with Nested Gradient Tapes

The cleanest pattern is:

1. watch the input vector
2. compute the scalar output
3. compute the first gradient
4. differentiate each gradient component again and keep the matching coordinate

1import tensorflow as tf
2
3
4def f(x):
5    return tf.reduce_sum(x**3 + 2.0 * x**2)
6
7
8def hessian_diagonal(x):
9    x = tf.convert_to_tensor(x, dtype=tf.float32)
10
11    with tf.GradientTape(persistent=True) as outer_tape:
12        outer_tape.watch(x)
13        with tf.GradientTape() as inner_tape:
14            inner_tape.watch(x)
15            y = f(x)
16        grad = inner_tape.gradient(y, x)
17
18    diagonal = []
19    for i in range(x.shape[0]):
20        second_grad = outer_tape.gradient(grad[i], x)
21        diagonal.append(second_grad[i])
22
23    del outer_tape
24    return tf.stack(diagonal)
25
26
27x = tf.constant([1.0, 2.0, 3.0])
28print(hessian_diagonal(x).numpy())

For the function above, the exact second derivative per coordinate is 6x + 4, so the output should be [10, 16, 22].

This approach avoids manually writing derivatives and is easy to verify against analytic formulas.

When jacobian Is Simpler but More Expensive

TensorFlow can also build the full Hessian by taking the Jacobian of the gradient and then extracting the diagonal.

1import tensorflow as tf
2
3
4def full_hessian_diagonal(x):
5    with tf.GradientTape() as outer_tape:
6        outer_tape.watch(x)
7        with tf.GradientTape() as inner_tape:
8            inner_tape.watch(x)
9            y = tf.reduce_sum(x**3 + 2.0 * x**2)
10        grad = inner_tape.gradient(y, x)
11
12    hessian = outer_tape.jacobian(grad, x)
13    return tf.linalg.diag_part(hessian)
14
15
16x = tf.constant([1.0, 2.0, 3.0])
17print(full_hessian_diagonal(x).numpy())

This is often shorter, but it materializes the full Hessian first. For large parameter vectors, that can be far more memory-intensive than needed.

So the tradeoff is:

• use jacobian when the vector is small or code clarity matters more
• use coordinate-wise second derivatives when you truly want only the diagonal

Batched Inputs Need Extra Care

If x has a batch dimension, decide what the scalar function actually is. GradientTape expects a scalar target for gradient, so you often reduce across the batch first or compute per-example diagonals deliberately.

A common pattern is to compute loss per example, then loop or vectorize over examples. The important point is to be explicit about whether you want:

• Hessian diagonal of the total loss
• Hessian diagonal per example
• Hessian diagonal with respect to model weights instead of input features

Those are different objects and can produce very different shapes.

Common Pitfalls

• Calling gradient on a non-scalar target without understanding how TensorFlow reduces it.
• Using outer_tape.jacobian on large vectors and accidentally building a huge dense Hessian.
• Forgetting persistent=True when taking multiple second-derivative queries from the same outer tape.
• Mixing batch dimensions and feature dimensions without deciding which Hessian you actually want.

Summary

• The Hessian diagonal contains the second derivative of each coordinate with respect to itself.
• In TensorFlow, nested GradientTape objects are the standard way to compute it.
• 'jacobian plus diag_part is simple but can allocate the full Hessian.'
• A coordinate-wise second-derivative loop is often better when only the diagonal is needed.
• Be explicit about whether you are differentiating with respect to inputs, per-example losses, or model parameters.