Gradient Descent
TensorFlow
Machine Learning
Neural Networks
Model Training

Update of the weights after Gradient Descent in TensorFlow

Master System Design with Codemia

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

Introduction

In TensorFlow, weights are updated after gradient descent when an optimizer applies the computed gradients to trainable variables. The important idea is that TensorFlow does not change weights just because you computed a loss; the update happens only when gradients are explicitly applied.

The Core Update Rule

The basic gradient descent rule is:

w = w - learning_rate * gradient

TensorFlow follows that same idea, but usually through an optimizer object rather than a handwritten assignment.

The training loop has three conceptual steps:

  1. run the model and compute the loss
  2. compute gradients of the loss with respect to the weights
  3. apply those gradients to the weights

A Minimal Manual Example

Here is a small TensorFlow example showing the full process.

python
1import tensorflow as tf
2
3w = tf.Variable(5.0)
4learning_rate = 0.1
5
6for step in range(5):
7    with tf.GradientTape() as tape:
8        loss = (w - 1.0) ** 2
9
10    grad = tape.gradient(loss, w)
11    w.assign_sub(learning_rate * grad)
12
13    print(f"step={step} w={w.numpy():.4f} loss={loss.numpy():.4f} grad={grad.numpy():.4f}")

This example uses assign_sub to perform the actual update. That is the TensorFlow equivalent of writing w = w - ..., but it updates the tf.Variable in place.

The Usual Approach: Use An Optimizer

In real model code, you usually let an optimizer handle the variable update logic.

python
1import tensorflow as tf
2
3w = tf.Variable(5.0)
4optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
5
6for step in range(5):
7    with tf.GradientTape() as tape:
8        loss = (w - 1.0) ** 2
9
10    grad = tape.gradient(loss, w)
11    optimizer.apply_gradients([(grad, w)])
12
13    print(f"step={step} w={w.numpy():.4f} loss={loss.numpy():.4f}")

The important line is optimizer.apply_gradients(...). That is where the weight update happens.

What Happens In model.fit

When you use model.fit, Keras is doing the same logic for you internally.

Conceptually, the framework:

  • runs the forward pass
  • computes the loss
  • computes gradients with GradientTape
  • applies gradients through the optimizer

So if you ever wonder "when are the weights updated," the answer is: after each training step when the optimizer applies gradients to the trainable variables.

A Small Keras Model Example

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(1, input_shape=(1,))
5])
6
7optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
8loss_fn = tf.keras.losses.MeanSquaredError()
9
10x = tf.constant([[1.0], [2.0], [3.0]])
11y = tf.constant([[2.0], [4.0], [6.0]])
12
13with tf.GradientTape() as tape:
14    predictions = model(x)
15    loss = loss_fn(y, predictions)
16
17grads = tape.gradient(loss, model.trainable_variables)
18optimizer.apply_gradients(zip(grads, model.trainable_variables))

After apply_gradients, the layer weights have changed.

Why Computing Gradients Is Not Enough

A very common misunderstanding is thinking that tape.gradient(...) already updates the weights. It does not. That call only computes the gradients.

The update happens only when you:

  • call assign_sub manually
  • or call optimizer.apply_gradients(...)
  • or use a higher-level training API that does it internally

That distinction is the heart of the question.

Different Optimizers Mean Different Update Details

Although the basic intuition is gradient descent, optimizers such as Adam, RMSprop, and momentum SGD use more than just the raw gradient. They keep internal state and modify the update rule.

But the structure is still the same:

  • compute gradients
  • apply optimizer-specific update to weights

So even when the math becomes more sophisticated, the update point in TensorFlow is still the optimizer application step.

Common Pitfalls

The most common mistake is computing the loss and gradients but never calling apply_gradients or a manual variable update. In that case, training appears to run but the weights never change.

Another mistake is forgetting that only tf.Variable objects or model trainable variables can be updated this way. Plain tensors are not mutable weights.

Developers also sometimes inspect weights before the optimizer step and conclude that TensorFlow did nothing. The order of operations matters.

Finally, remember that model.predict does not update weights. Only training steps do.

Summary

  • In TensorFlow, weights are updated when gradients are applied to trainable variables.
  • 'GradientTape computes gradients but does not update weights by itself.'
  • Manual updates use methods like assign_sub.
  • Optimizer-based updates use optimizer.apply_gradients(...).
  • 'model.fit performs the same pattern internally during each training step.'

Course illustration
Course illustration

All Rights Reserved.