Neural Networks
Keras
TensorFlow
Optimization
Machine Learning

Find input that maximises output of a neural network using Keras and TensorFlow

Master System Design with Codemia

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

Introduction

Finding an input that maximizes a model output is an optimization problem over the input space instead of over the model weights. This is useful for feature visualization, adversarial analysis, and understanding what kind of signal a trained network responds to most strongly.

Treat the Input as a Variable

The basic idea is simple:

  1. freeze the trained model
  2. create an input tensor as a trainable variable
  3. compute the target output
  4. use gradient ascent on the input

With TensorFlow 2 and Keras, tf.GradientTape makes this straightforward.

python
1import tensorflow as tf
2from tensorflow import keras
3
4model = keras.Sequential([
5    keras.layers.Input(shape=(2,)),
6    keras.layers.Dense(8, activation="tanh"),
7    keras.layers.Dense(1)
8])
9
10model.compile(optimizer="adam", loss="mse")
11
12x_train = tf.random.normal((256, 2))
13y_train = tf.reduce_sum(x_train, axis=1, keepdims=True)
14model.fit(x_train, y_train, epochs=5, verbose=0)

At this point the model is trained. Now the optimization target shifts from weights to input.

Gradient Ascent on the Input

python
1import tensorflow as tf
2
3candidate = tf.Variable([[0.0, 0.0]], dtype=tf.float32)
4optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
5
6for step in range(100):
7    with tf.GradientTape() as tape:
8        score = model(candidate, training=False)
9        loss = -tf.reduce_mean(score)  # negative because optimizer minimizes
10
11    grads = tape.gradient(loss, candidate)
12    optimizer.apply_gradients([(grads, candidate)])
13
14print("candidate:", candidate.numpy())
15print("score:", model(candidate, training=False).numpy())

This performs gradient ascent by minimizing the negative score.

Pick the Right Output Target

For a scalar regression model, the target is obvious. For classification or multi-output models, choose a specific output unit.

python
1classifier = keras.Sequential([
2    keras.layers.Input(shape=(4,)),
3    keras.layers.Dense(16, activation="relu"),
4    keras.layers.Dense(3, activation="softmax")
5])
6
7classifier.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
8
9candidate = tf.Variable([[0.1, 0.1, 0.1, 0.1]], dtype=tf.float32)
10target_index = 1
11optimizer = tf.keras.optimizers.Adam(0.05)
12
13for _ in range(80):
14    with tf.GradientTape() as tape:
15        probs = classifier(candidate, training=False)
16        target_score = probs[:, target_index]
17        loss = -tf.reduce_mean(target_score)
18
19    grads = tape.gradient(loss, candidate)
20    optimizer.apply_gradients([(grads, candidate)])
21
22print(candidate.numpy())
23print(classifier(candidate, training=False).numpy())

Here the optimization asks, "what input most increases class one probability for this model."

Add Constraints So the Result Stays Meaningful

Unconstrained optimization often finds strange, extreme inputs that maximize the output but are not interpretable. That is why regularization matters.

Two common constraints are:

  • clipping the input to a valid range
  • adding a penalty for large input magnitude
python
1candidate = tf.Variable([[0.0, 0.0]], dtype=tf.float32)
2optimizer = tf.keras.optimizers.Adam(0.05)
3
4for _ in range(100):
5    with tf.GradientTape() as tape:
6        score = model(candidate, training=False)
7        penalty = 0.1 * tf.reduce_sum(tf.square(candidate))
8        loss = -(tf.reduce_mean(score) - penalty)
9
10    grads = tape.gradient(loss, candidate)
11    optimizer.apply_gradients([(grads, candidate)])
12    candidate.assign(tf.clip_by_value(candidate, -2.0, 2.0))

Without constraints, the optimizer may exploit unrealistic directions in the learned function.

For Images, Start from Noise and Normalize

In image-based models, this same technique is used for feature visualization. The input starts as noise, then gradient ascent pushes it toward patterns that activate a chosen filter or class.

In that context, you usually also need:

  • image-space clipping
  • smoothness regularization
  • smaller learning rates
  • periodic normalization

The math is the same. Only the input shape and regularization become more important.

Watch for Saturation and Local Maxima

This is still non-convex optimization. The result depends on:

  • initialization
  • learning rate
  • number of steps
  • regularization strength

If the optimization stalls, try several random starting points instead of assuming the first result is the global optimum.

Common Pitfalls

The biggest mistake is accidentally updating model weights instead of the input tensor. The model should stay fixed during this process.

Another issue is optimizing an unconstrained input and then treating the result as meaningful. High-scoring garbage is easy to generate if you do not regularize.

A third problem is maximizing the wrong output tensor, especially in multi-class models where you intended to target one class but optimized the full vector implicitly.

Summary

  • Maximize model output by treating the input as the optimization variable.
  • Use GradientTape to compute gradients with respect to the input.
  • For classification, optimize a specific output unit or class score.
  • Add clipping or penalties so the result stays interpretable.
  • Expect multiple runs and tuning because the optimization landscape is non-convex.

Course illustration
Course illustration

All Rights Reserved.