tensorflow
deep learning
neural networks
machine learning
gradients

Tensorflow dense gradient explanation?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

A dense layer computes a matrix multiply, adds a bias, and optionally applies an activation. The gradient question is really about how the loss changes with respect to the layer's weights, biases, and inputs during backpropagation.

Start with the Dense-Layer Formula

Ignoring activation for a moment, a dense layer is:

z = xW + b

where:

  • 'x is the input batch'
  • 'W is the weight matrix'
  • 'b is the bias vector'
  • 'z is the pre-activation output'

If an activation exists, the final output is y = activation(z).

For a batch of examples, the shapes usually look like:

  • 'x: (batch_size, input_dim)'
  • 'W: (input_dim, units)'
  • 'b: (units,)'
  • 'z: (batch_size, units)'

Understanding these shapes makes the gradients much easier to interpret.

What Gradients TensorFlow Computes

Suppose your loss is L. Backpropagation for the linear part of the layer gives:

  • gradient with respect to weights: dL/dW = x^T * dL/dz
  • gradient with respect to bias: batchwise sum of dL/dz
  • gradient with respect to input: dL/dx = dL/dz * W^T

Intuitively:

  • the weight gradient tells you how to change each input-to-output connection
  • the bias gradient tells you how to shift each unit
  • the input gradient is what gets passed to the previous layer

If the layer has an activation, TensorFlow first applies the chain rule through the activation, then through the matrix multiply.

Inspecting Dense Gradients with GradientTape

Here is a minimal example:

python
1import tensorflow as tf
2
3layer = tf.keras.layers.Dense(
4    units=2,
5    use_bias=True,
6    kernel_initializer=tf.keras.initializers.Constant([[1.0, 2.0], [3.0, 4.0]]),
7    bias_initializer=tf.keras.initializers.Constant([0.5, -0.5]),
8)
9
10x = tf.constant([[1.0, 2.0]], dtype=tf.float32)
11
12with tf.GradientTape() as tape:
13    y = layer(x)
14    loss = tf.reduce_sum(y)
15
16grads = tape.gradient(loss, layer.trainable_variables)
17
18for variable, grad in zip(layer.trainable_variables, grads):
19    print(variable.name)
20    print(grad.numpy())

Because the loss is just the sum of the outputs, the math stays easy to reason about.

For this example:

  • the gradient for the kernel depends on the input values
  • the bias gradient becomes ones because each output contributes directly to the sum

That is why dense-layer gradients often look simple in toy examples but become less intuitive once the loss and activation are more complex.

Why the Weight Gradient Depends on the Input

Each weight connects one input feature to one neuron. If a particular input feature is large, changing its corresponding weight changes the output more strongly. That is why the input batch appears directly in the kernel gradient.

A useful mental model is:

  • weight gradient = "how much this connection mattered"
  • bias gradient = "how much this neuron's constant shift mattered"

For a batch, TensorFlow accumulates contributions across all examples.

Add an Activation and the Chain Rule

Now consider a ReLU activation:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential(
4    [
5        tf.keras.layers.Dense(3, activation="relu", input_shape=(2,)),
6        tf.keras.layers.Dense(1),
7    ]
8)
9
10x = tf.constant([[1.0, -1.0], [2.0, 3.0]], dtype=tf.float32)
11target = tf.constant([[0.5], [1.0]], dtype=tf.float32)
12
13with tf.GradientTape() as tape:
14    prediction = model(x)
15    loss = tf.reduce_mean(tf.square(prediction - target))
16
17grads = tape.gradient(loss, model.trainable_variables)
18
19for variable, grad in zip(model.trainable_variables, grads):
20    print(variable.name, grad.shape)

TensorFlow automatically handles the chain rule. You do not manually derive every intermediate term in normal Keras training code.

Still, it helps to know what happens conceptually:

  1. compute the loss gradient at the output
  2. move backward through the second dense layer
  3. apply the ReLU derivative
  4. continue backward through the first dense layer

Why Some Gradients Are Zero

Dense-layer gradients can contain zeros for several reasons:

  • ReLU blocked the gradient because the pre-activation was negative
  • the loss did not depend on a variable in the current tape scope
  • the variable was not trainable

This is not automatically a bug. A zero gradient can be completely correct.

Common Pitfalls

Confusing dL/dy with dL/dW is common. The loss gradient at the layer output is only one step in the backprop chain, not the final weight gradient.

Ignoring tensor shapes makes dense-layer gradients look mysterious. Most confusion disappears once you track the dimensions carefully.

Reading a zero gradient as a failure without checking the activation function or tape scope often leads to incorrect debugging.

Assuming TensorFlow computes gradients for non-trainable variables or operations outside the tape scope will produce None and surprise many beginners.

Summary

  • A dense layer is matrix multiplication plus bias, with an optional activation afterward.
  • TensorFlow computes gradients for kernel, bias, and inputs using the chain rule.
  • The kernel gradient depends directly on the input batch and the output-side gradient.
  • 'GradientTape lets you inspect these values explicitly.'
  • When debugging dense gradients, watch tensor shapes, activation behavior, and tape scope first.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.