Keras
loss gradient
backpropagation
deep learning
model inputs

How to compute loss gradient w.r.t to model inputs in a Keras model?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Computing the gradient of a loss with respect to a model's input is useful for saliency maps, adversarial examples, sensitivity analysis, and debugging. In modern Keras and TensorFlow, the standard tool for this is tf.GradientTape.

Why Input Gradients Are Different from Weight Gradients

During training, Keras normally computes gradients with respect to trainable variables such as kernels and biases. For interpretability or adversarial workflows, you want the derivative of the loss with respect to the input tensor instead.

The core idea is the same backpropagation graph, but the source of the gradient changes:

  • parameter gradients: loss with respect to model.trainable_variables
  • input gradients: loss with respect to the input batch

If your input is a plain tensor and not a tf.Variable, TensorFlow will not watch it automatically. You must call tape.watch(inputs).

Basic Pattern with tf.GradientTape

The pattern is:

  1. Create the input tensor.
  2. Open a GradientTape.
  3. Tell the tape to watch the input tensor.
  4. Run the model forward.
  5. Compute the loss.
  6. Ask the tape for gradient(loss, inputs).
python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(8, activation="relu", input_shape=(4,)),
5    tf.keras.layers.Dense(3)
6])
7
8loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
9
10x = tf.constant([[0.2, -0.1, 0.4, 0.7]], dtype=tf.float32)
11y_true = tf.constant([2], dtype=tf.int32)
12
13with tf.GradientTape() as tape:
14    tape.watch(x)
15    logits = model(x, training=False)
16    loss = loss_fn(y_true, logits)
17
18grad_x = tape.gradient(loss, x)
19
20print("loss:", float(loss.numpy()))
21print("gradient shape:", grad_x.shape)
22print(grad_x.numpy())

That gives the gradient of the scalar loss with respect to each input feature in the batch.

Choosing the Target Correctly

The tape computes gradients of a target tensor with respect to a source tensor. Most of the time, the target should be a scalar loss. If you pass a non-scalar tensor, TensorFlow will usually behave as if you asked for the gradient of the sum of that tensor.

That matters because these are different questions:

  • gradient of classification loss with respect to the input
  • gradient of one output logit with respect to the input
  • full Jacobian of all outputs with respect to the input

If you want the gradient of one class score instead of the training loss, select that score explicitly:

python
1with tf.GradientTape() as tape:
2    tape.watch(x)
3    logits = model(x, training=False)
4    target_score = logits[:, 1]
5
6score_grad = tape.gradient(target_score, x)
7print(score_grad.numpy())

For a full output-by-input derivative structure, use tape.jacobian, not gradient.

Working with a Compiled Keras Model

You do not need to call model.fit or modify the model definition. A compiled model is fine, but the gradient calculation itself still needs to happen inside a tape context.

If you already have a compiled loss object, you can reuse it:

python
1loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
2
3x = tf.constant([[1.0, 0.5, -0.5, 0.3]], dtype=tf.float32)
4y_true = tf.constant([[0.0, 1.0, 0.0]], dtype=tf.float32)
5
6with tf.GradientTape() as tape:
7    tape.watch(x)
8    logits = model(x, training=False)
9    loss = loss_fn(y_true, logits)
10
11grad_x = tape.gradient(loss, x)

This works the same way for subclassed models, functional models, and sequential models.

Practical Uses

Input gradients are often used to answer one of these questions:

  • Which input dimensions most affect the current prediction?
  • How should I perturb the input to increase or decrease the loss?
  • Is the model sensitive to noise in a specific region?

In image models, you usually reshape or visualize the gradient as a heatmap. In tabular models, you often inspect the per-feature magnitudes directly.

Common Pitfalls

The most common mistake is forgetting tape.watch(x). TensorFlow automatically watches trainable variables, but not arbitrary constant tensors.

Another common mistake is passing a NumPy array directly and expecting a gradient back. Convert it to a tensor first with tf.constant or tf.convert_to_tensor.

People also mix up logits and probabilities. If your loss expects logits, pass raw outputs and set from_logits=True. A mismatch there changes both the loss and the gradient.

Finally, if you want per-output derivatives, do not use gradient on a matrix and assume it is a full Jacobian. Use tape.jacobian when the shape structure matters.

Summary

  • Use tf.GradientTape to compute input gradients in Keras.
  • Call tape.watch(inputs) when the input is not a trainable variable.
  • Compute a scalar target such as the loss or a selected class score inside the tape.
  • Use tape.gradient(target, inputs) for ordinary gradients and tape.jacobian for full output-by-input derivatives.
  • Be careful about logits versus probabilities and tensor versus NumPy input types.

Course illustration
Course illustration

All Rights Reserved.