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:
losswith respect tomodel.trainable_variables - input gradients:
losswith 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:
- Create the input tensor.
- Open a
GradientTape. - Tell the tape to watch the input tensor.
- Run the model forward.
- Compute the loss.
- Ask the tape for
gradient(loss, inputs).
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:
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:
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.GradientTapeto 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 andtape.jacobianfor full output-by-input derivatives. - Be careful about logits versus probabilities and tensor versus NumPy input types.

