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.
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:
- '
xis the input batch' - '
Wis the weight matrix' - '
bis the bias vector' - '
zis 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:
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:
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:
- compute the loss gradient at the output
- move backward through the second dense layer
- apply the ReLU derivative
- 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.
- '
GradientTapelets you inspect these values explicitly.' - When debugging dense gradients, watch tensor shapes, activation behavior, and tape scope first.
Related reading
- Tensorflow device CUDA0 not supported by XLA service while setting up XLA_GPU_JIT device number 0
- Tensorflow Diagonal Subtensor for 3D Convolutional NN
- tensorflow difference between multi GPUs and distributed tensorflow
- Tensorflow Documentation
- Tensorflow dense_to_sparse
- Tensorflow Dictionary lookup with String tensor
- Tensorflow Different ways to Export and Run graph in C
- Tensorflow distributed training high bandwidth on Parameter Server
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.