Mahmoud Abas, Full-Stack Software Engineer & AI

Machine Learning & AI - ML Review and Gradient Descent Example

In the print dialog, choose "Save as PDF" as the destination.

ML Review and Gradient Descent Example

Training a linear regression model end to end with TensorFlow

The ML Workflow

Load the dataset

Build the model

multiple linear regression

Evaluate performance

Generating Data from a Known Model

Y = b + X₁w₁ + X₂w₂

True model: b = 1.5, w_1 = 2, w_2 = 5
X_1 and X_2 are drawn uniformly from [0, 10]

python

x1 = tf.random.uniform([150], minval=0, maxval=10)
x2 = tf.random.uniform([150], minval=0, maxval=10)
y = 1.5 + 2.0 * x1 + 5.0 * x2
dataset = tf.stack([y, x1, x2], axis=1)

Goal: recover b, w_1, w_2 purely from the generated (x1, x2, y) data

Splitting Train/Test Sets

Shuffle before splitting to avoid biased splits

python

dataset_shuffled = tf.random.shuffle(dataset, seed=22)
train_data, test_data = dataset_shuffled[:125], dataset_shuffled[125:]
x_train, y_train = train_data[:, 1:], train_data[:, 0]

125 training rows, 25 test rows, about 17% held out

Building the Regression Model

Y = Xw + b

Multiple linear regression stacks every feature into one weight vector
tf.Module plus @tf.function traces the model into a reusable graph

python

class LinearRegression(tf.Module):
    @tf.function
    def __call__(self, x):
        if not self.built:
            self.w = tf.Variable(tf.random.uniform([x.shape[-1], 1]))
            self.b = tf.Variable(tf.random.uniform([]))
            self.built = True
        return tf.squeeze(tf.matmul(x, self.w) + self.b, axis=1)

Defining the Loss Function

Mean Squared Error is the standard loss for regression

MSE = 1m∑mi=1(ŷ_i - y_i)²

python

def mse_loss(y_pred, y):
    return tf.reduce_mean(tf.square(y_pred - y))

Training means adjusting w and b to minimize this loss

Effect of Batch Size

Larger batches: faster training and better GPU utilization, but more memory and can generalize worse
Smaller batches: better generalization, noisier gradients, slower per epoch

Batch Size	When to Use
`1`	Extreme cases, reinforcement learning, online learning
`8` - `16`	Small datasets, memory-constrained GPUs
`32` (default)	Good balance of speed and generalization
`64` - `128`	Large datasets, powerful GPUs
`256+`	Extremely large datasets with TPU/GPU clusters

Effect of Learning Rate

High learning rate: faster convergence, but risks overshooting or diverging
Low learning rate: stable but slow, may get stuck making tiny progress

Model Type	Suggested `learning_rate`
Simple linear regression	`0.01` - `0.1`
Deep neural networks	`0.001` - `0.01`
Convolutional neural networks	`0.0001` - `0.01`
Transformers	`0.00001` - `0.0001`
Reinforcement learning	`0.0001` - `0.1`

Batching the Dataset

tf.data.Dataset handles shuffling and batching efficiently

python

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=125).batch(32)

Produces batches of 32, with a smaller final batch of 29

The Training Loop

For every batch in the epoch, compute the loss
Get the gradients of the loss with respect to w and b
Update the parameters by stepping against the gradient

python

for epoch in range(epochs):
    for x_batch, y_batch in train_dataset:
        with tf.GradientTape() as tape:
            y_pred = lin_reg(x_batch)
            loss = mse_loss(y_pred, y_batch)
        grads = tape.gradient(loss, lin_reg.variables)
        for g, v in zip(grads, lin_reg.variables):
            v.assign_sub(learning_rate * g)

Gradient descent: Stepping in the opposite direction of the gradient, one update at a time, to reduce the loss

Training Results

Epoch	`MSE` Loss
`0`	`299.503`
`10`	`0.078`
`20`	`0.062`
`30`	`0.051`
`40`	`0.041`
`50`	`0.034`
`60`	`0.028`
`70`	`0.024`
`80`	`0.019`
`90`	`0.016`

After 100 epochs: final train loss 0.014, final test loss 0.012
Train and test loss track closely, the model generalizes well, no overfitting

Final Parameters and Loss Curve

Recall the goal: recover b, w_1, w_2 from the generated data alone

Parameter	True Value	Recovered Value
`b`	`1.5`	≈ `1.18`
`w_1`	`2.0`	≈ `2.03`
`w_2`	`5.0`	≈ `5.02`

Plotting train vs. test loss over epochs confirms smooth, stable convergence

Key Takeaways

Gradient descent iteratively updates parameters in the direction that reduces loss
tf.GradientTape plus assign_sub(learning_rate * gradient) is the core training step
Batch size trades off training speed against generalization and memory use
Learning rate is the single most sensitive hyperparameter, too high diverges, too low crawls
Tracking train vs. test loss together is how you catch overfitting early

Open this lesson

Mahmoud AbasML Review and Gradient Descent Example

class LinearRegression(tf.Module): @tf.function def __call__(self, x): if not self.built: self.w = tf.Variable(tf.random.uniform([x.shape[-1], 1])) self.b = tf.Variable(tf.random.uniform([])) self.built = True return tf.squeeze(tf.matmul(x, self.w) + self.b, axis=1)

Batch Size

When to Use

1

Extreme cases, reinforcement learning, online learning

8 - 16

Small datasets, memory-constrained GPUs

32 (default)

Good balance of speed and generalization

64 - 128

Large datasets, powerful GPUs

256+

Extremely large datasets with TPU/GPU clusters

Model Type

Suggested learning_rate

Simple linear regression

0.01 - 0.1

Deep neural networks

0.001 - 0.01

Convolutional neural networks

0.0001 - 0.01

Transformers

0.00001 - 0.0001

Reinforcement learning

0.0001 - 0.1

for epoch in range(epochs): for x_batch, y_batch in train_dataset: with tf.GradientTape() as tape: y_pred = lin_reg(x_batch) loss = mse_loss(y_pred, y_batch) grads = tape.gradient(loss, lin_reg.variables) for g, v in zip(grads, lin_reg.variables): v.assign_sub(learning_rate * g)

Epoch

MSE Loss

0

299.503

10

0.078

20

0.062

30

0.051

40

0.041

50

0.034

60

0.028

70

0.024

80

0.019

90

0.016

Parameter

True Value

Recovered Value

b

1.5

≈ 1.18

w_1

2.0

≈ 2.03

w_2

5.0

≈ 5.02