Project: Fit a Line with Gradient Descent

What you'll build

A flat-rental site wants to tell a landlord "flats like yours rent for about this much". You have 12 flats in Dhaka with their size and monthly rent. You'll fit a straight line, rent = w × size + b, and you'll do it the way every neural network is trained: write the loss, write its gradient, and run gradient descent by hand. No ML library.

Along the way you'll hit the first real training bug most people meet, a learning rate that's either too slow or explodes, and fix it with feature scaling. Then you'll prove your answer is right by checking it against the exact formula.

Prerequisites

• Calculus and Gradients: partial derivatives and the gradient check.
• Optimization: the update rule and learning rates.
• Python 3 or Go. Nothing to install.

Architecture

The training loop has four parts, and they're the same four parts in every model you'll train in this course:

Build it step by step

Step 1: the loss. Mean squared error over all flats: MSE = average of (w·x + b − y)².

Step 2: the gradient. Apply the chain rule to each term. For one flat the error is e = w·x + b − y, and:

∂(e²)/∂w = 2e × x        ∂(e²)/∂b = 2e × 1

Average those over all flats and you have dw and db.

Step 3: the loop. Start at w = b = 0 and repeat w -= lr * dw; b -= lr * db.

Step 4: scale the input. Sizes are around 1,000 while the intercept multiplies 1, so the gradient for w is about a thousand times bigger than for b. One learning rate can't suit both. Fix: rescale sizes to mean 0 and spread 1, (x − mean) / std, train, then convert the weights back to square feet.

Step 5: check. For a straight line there's an exact answer (least squares). Your gradient descent result should match it.

"""Fit rent = w * size + b to 12 Dhaka flats with gradient descent written by hand."""
import math

SIZE = [650, 800, 900, 1000, 1100, 1200, 1250, 1400, 1500, 1650, 1800, 2000]  # square feet
RENT = [12, 15, 16, 19, 20, 22, 24, 25, 28, 30, 33, 36]                       # thousand taka a month
N = len(SIZE)

def mse(w, b, xs):
    return sum((w * x + b - y) ** 2 for x, y in zip(xs, RENT)) / N

def gradients(w, b, xs):
    dw = sum(2 * (w * x + b - y) * x for x, y in zip(xs, RENT)) / N
    db = sum(2 * (w * x + b - y) for x, y in zip(xs, RENT)) / N
    return dw, db

def train(xs, lr, steps, report=()):
    w = b = 0.0
    history = []
    for step in range(1, steps + 1):
        dw, db = gradients(w, b, xs)
        w, b = w - lr * dw, b - lr * db
        if step in report:
            history.append((step, mse(w, b, xs)))
    return w, b, history

# Step 1: raw square feet. Any learning rate is either too slow or explodes.
for lr, runs in ((1e-7, (10, 5000)), (1e-6, (10, 50))):
    for steps in runs:
        w, b, _ = train(SIZE, lr, steps)
        print(f"raw sizes, lr={lr:g}: after {steps} steps loss = {mse(w, b, SIZE):.4g}, b = {b:.4g}")

# Step 2: scale sizes to mean 0, spread 1. Now lr = 0.1 works.
mean = sum(SIZE) / N
std = math.sqrt(sum((x - mean) ** 2 for x in SIZE) / N)
scaled = [(x - mean) / std for x in SIZE]
w, b, history = train(scaled, 0.1, 100, report=(1, 5, 10, 20, 40, 60, 100))
top = history[0][1]
print("\nscaled sizes, lr=0.1")
for step, loss in history:
    print(f"step {step:>3}  loss {loss:8.3f}  {'#' * max(1, round(loss / top * 40))}")

# Step 3: turn the scaled weights back into rent per square foot.
per_sqft = w / std
base = b - w * mean / std
print(f"\nrent = {per_sqft * 1000:.2f} per 1000 sqft x size + {base:.2f}")
print(f"1300 sqft flat -> {per_sqft * 1300 + base:.1f} thousand taka")

# Step 4: check against the exact least-squares answer.
my = sum(RENT) / N
exact_w = sum((x - mean) * (y - my) for x, y in zip(SIZE, RENT)) / sum((x - mean) ** 2 for x in SIZE)
exact_b = my - exact_w * mean
print(f"exact:  {exact_w * 1000:.2f} per 1000 sqft, intercept {exact_b:.2f}")

// Fit rent = w * size + b to 12 Dhaka flats with gradient descent written by hand.
package main

import (
	"fmt"
	"math"
	"strings"
)

var size = []float64{650, 800, 900, 1000, 1100, 1200, 1250, 1400, 1500, 1650, 1800, 2000} // square feet
var rent = []float64{12, 15, 16, 19, 20, 22, 24, 25, 28, 30, 33, 36}                       // thousand taka a month
var n = float64(len(size))

func mse(w, b float64, xs []float64) float64 {
	s := 0.0
	for i, x := range xs {
		e := w*x + b - rent[i]
		s += e * e
	}
	return s / n
}

func gradients(w, b float64, xs []float64) (dw, db float64) {
	for i, x := range xs {
		e := w*x + b - rent[i]
		dw += 2 * e * x
		db += 2 * e
	}
	return dw / n, db / n
}

type point struct {
	step int
	loss float64
}

func train(xs []float64, lr float64, steps int, report map[int]bool) (w, b float64, history []point) {
	for step := 1; step <= steps; step++ {
		dw, db := gradients(w, b, xs)
		w, b = w-lr*dw, b-lr*db
		if report[step] {
			history = append(history, point{step, mse(w, b, xs)})
		}
	}
	return
}

func main() {
	// Step 1: raw square feet. Any learning rate is either too slow or explodes.
	for _, r := range []struct {
		lr   float64
		runs []int
	}{{1e-7, []int{10, 5000}}, {1e-6, []int{10, 50}}} {
		lr := r.lr
		for _, steps := range r.runs {
			w, b, _ := train(size, lr, steps, nil)
			fmt.Printf("raw sizes, lr=%g: after %d steps loss = %.4g, b = %.4g\n", lr, steps, mse(w, b, size), b)
		}
	}

	// Step 2: scale sizes to mean 0, spread 1. Now lr = 0.1 works.
	mean, varSum := 0.0, 0.0
	for _, x := range size {
		mean += x
	}
	mean /= n
	for _, x := range size {
		varSum += (x - mean) * (x - mean)
	}
	std := math.Sqrt(varSum / n)
	scaled := make([]float64, len(size))
	for i, x := range size {
		scaled[i] = (x - mean) / std
	}
	report := map[int]bool{1: true, 5: true, 10: true, 20: true, 40: true, 60: true, 100: true}
	w, b, history := train(scaled, 0.1, 100, report)
	top := history[0].loss
	fmt.Println("\nscaled sizes, lr=0.1")
	for _, p := range history {
		bars := int(math.Max(1, math.RoundToEven(p.loss/top*40)))
		fmt.Printf("step %3d  loss %8.3f  %s\n", p.step, p.loss, strings.Repeat("#", bars))
	}

	// Step 3: turn the scaled weights back into rent per square foot.
	perSqft := w / std
	base := b - w*mean/std
	fmt.Printf("\nrent = %.2f per 1000 sqft x size + %.2f\n", perSqft*1000, base)
	fmt.Printf("1300 sqft flat -> %.1f thousand taka\n", perSqft*1300+base)

	// Step 4: check against the exact least-squares answer.
	my, num := 0.0, 0.0
	for _, y := range rent {
		my += y
	}
	my /= n
	for i, x := range size {
		num += (x - mean) * (rent[i] - my)
	}
	exactW := num / varSum
	fmt.Printf("exact:  %.2f per 1000 sqft, intercept %.2f\n", exactW*1000, my-exactW*mean)
}

Output:

raw sizes, lr=1e-07: after 10 steps loss = 0.3622, b = 1.307e-05
raw sizes, lr=1e-07: after 5000 steps loss = 0.2677, b = 6.107e-05
raw sizes, lr=1e-06: after 10 steps loss = 7.603e+10, b = -0.1486
raw sizes, lr=1e-06: after 50 steps loss = 2.03e+43, b = -2.429e+15

scaled sizes, lr=0.1
step   1  loss  380.887  ########################################
step   5  loss   64.103  #######
step  10  loss    7.099  #
step  20  loss    0.321  #
step  40  loss    0.242  #
step  60  loss    0.242  #
step 100  loss    0.242  #

rent = 17.93 per 1000 sqft x size + 0.54
1300 sqft flat -> 23.9 thousand taka
exact:  17.93 per 1000 sqft, intercept 0.54

Reading it:

• Raw sizes, lr = 1e-7: the loss looks fine (0.27), but after 5,000 steps b is still 0.00006. w learned, b barely moved, because its gradient is a thousand times smaller. The model is stuck with a wrong intercept and you'd never notice from the loss alone.
• Raw sizes, lr = 1e-6: ten times bigger and it explodes, 7.6 × 10¹⁰ after 10 steps and 2 × 10⁴³ after 50. There's no learning rate that works for both knobs.
• Scaled sizes, lr = 0.1: the loss falls from 381 to 0.32 in 20 steps and settles at 0.242. The bars are the loss curve: steep at first, then flat.
• The answer: rent ≈ 17.93 thousand taka per 1,000 sqft, plus 0.54. A 1,300 sqft flat should rent for about 23.9 thousand taka.
• The check: the exact least-squares answer is the same, 17.93 and 0.54. Your loop found the true best line.

Test it

import contextlib
import io

with contextlib.redirect_stdout(io.StringIO()):  # main.py prints its demo on import
    import main
from main import SIZE, gradients, mse, train

def test_gradient_matches_nudging():
    for w, b in [(0.0, 0.0), (0.01, 3.0), (-0.02, 10.0)]:
        dw, db = gradients(w, b, SIZE)
        h = 1e-6
        nw = (mse(w + h, b, SIZE) - mse(w - h, b, SIZE)) / (2 * h)
        nb = (mse(w, b + h, SIZE) - mse(w, b - h, SIZE)) / (2 * h)
        assert abs(dw - nw) / abs(nw) < 1e-4 and abs(db - nb) / abs(nb) < 1e-4

def test_loss_never_goes_up_with_a_good_learning_rate():
    _, _, hist = train(main.scaled, 0.1, 100, report=tuple(range(1, 101)))
    losses = [l for _, l in hist]
    assert all(b <= a + 1e-12 for a, b in zip(losses, losses[1:]))

def test_matches_exact_least_squares():
    assert abs(main.per_sqft - main.exact_w) < 1e-6
    assert abs(main.base - main.exact_b) < 1e-3

if __name__ == "__main__":
    for name, fn in list(globals().items()):
        if name.startswith("test_"):
            fn()
            print("pass", name)

pass test_gradient_matches_nudging
pass test_loss_never_goes_up_with_a_good_learning_rate
pass test_matches_exact_least_squares

• Gradient check: the hand-written gradient must match the nudging estimate at several points. If you get a sign or a factor of 2 wrong, this fails immediately, before you waste time tuning learning rates.
• The loss never rises with a good learning rate on this convex problem. If it does, the learning rate is too high or the gradient is wrong.
• Matches the exact solution. Only possible for simple models, so use it while you can: it proves the whole loop end to end.

Trade-offs

Take it further

1. Add a second feature, the floor number, so rent = w1·size + w2·floor + b. Write the extra partial derivative and check it with the gradient check.
2. Add one outlier (a 1,000 sqft flat at 90 thousand) and watch the line tilt. Then try Huber loss.
3. Train with mini-batches of 4 flats chosen randomly each step, and compare the loss curve with full batch.
4. Plot the loss curve properly with matplotlib, with a log scale on the loss axis.