Exercise 6¶

What this assignment is about¶

Backpropagation from Scratch:

In this assignment you will compute the gradient of a scalar loss with respect to the weight matrices of a small network, in three different ways, and compare their computational cost. By the end you will understand, at a mechanical level, what PyTorch's autograd is doing when you call .backward().

General Instructions¶

All submitted work must be stored in your GitLab repository. Below it will be explained how to do that. For the repository it is important that:

• one exercise = one GitLab repository → each exercise gets its own GitLab repository
• your repository must be private
• your repository must be named exactly as specified (in the setup below); do not change the name or the upper-case/lower-case letters
• create commits frequently, at least one commit per question
• push your work frequently (this also is your backup if your laptop breaks)
• your repository must be up-to-date before the deadline
• you must double-check your repository in the web browser to make sure all your work is uploaded to the code.vt.edu server
• you must provide screenshots as proof when they are requested; screenshots are also committed to the repository

⚠️ You must run your code yourself! Not following this is an Honor Code violation.

⚠️ You must create complete log files of all AI conversations! Not following this is an Honor Code violation.

0. Setup: Git and GitLab¶

1. Create a new repository on GitLab at code.vt.edu:

   • in the top-left corner, click the "+" icon
   • in the new menu that pops up, click "New project/repository"; a new page loads
   • in the new page, click the "Create black project" button; a new page loads:
     • enter the "Project name":

       cmda3634-ex6
       

     • in the drop-down field "Pick a group or namespace", select your username
     • the "Project slug" field should populate automatically after entering a project name; do not change the slug
     • for the "Visibility level", select "private" (very important!)
     • for the "Project configuration", select only "Initialize repository with README"

2. Add the instructor and TA to your repository as a member with Developer access. In the web browser with your new repository open:

   • in the left sidebar, hover the mouse cursor over "Manage"
   • in the menu that pops up, click "Members"; a new page loads
   • on the new page, click button "Invite members"
     • enter the jrudi
     • select maximum role: "Developer"
     • click the button "Invite"
   • repeat to invite a new member:
     • enter the jraind
     • select maximum role: "Developer"
     • click the button "Invite"

Background: Why differentiate with respect to the weights?¶

In a neural network, the input \(x\) and the target \(y\) are fixed data — we cannot change them. The only things we can change are the weight matrices \(W_1, W_2, W_3\). Training means repeatedly adjusting the weights to reduce the loss \(\ell\), and the gradient \(\partial \ell / \partial W_k\) tells us in which direction to adjust \(W_k\).

Concretely, one step of gradient descent updates each weight matrix as

where \(\eta > 0\) is called the learning rate. Computing these gradients efficiently is exactly what backpropagation does, and exactly what this assignment asks you to implement.

PyTorch syntax reference¶

The table below covers every operation you will need.

1-D vs 2-D vectors. PyTorch distinguishes a 1-D tensor of shape (n,) from a 2-D row vector of shape (1, n). With @: (n,) @ (n, m) → (m,), and (1, n) @ (n, m) → (1, m). Working with 1-D tensors throughout is fine and keeps the code cleaner.

Setup¶

The network has two hidden layers and a scalar output.

import torch
import time

torch.manual_seed(0)

n, d = 8, 6          # input dim, hidden dim  (kept small for Part 1)

x  = torch.randn(n)  # input,  shape (n,)
y  = torch.randn(1)  # target, shape (1,)

# weight matrices — no autograd yet
W1 = torch.randn(d, n)   # W_1 : ℝⁿ → ℝᵈ,  shape (d, n)
W2 = torch.randn(d, d)   # W_2 : ℝᵈ → ℝᵈ,  shape (d, d)
W3 = torch.randn(1, d)   # W_3 : ℝᵈ → ℝ¹,  shape (1, d)

The network computes:

followed by the squared loss:

The objects you want to compute are the three weight gradients:

Each gradient is a matrix of the same shape as the corresponding weight.

1. Part — Forward pass and the full Jacobian (30 points)¶

1a — Implement the forward pass¶

Using only PyTorch tensor operations (no nn.Module, no autograd), implement the forward pass and store \(z_1, z_2, z_3\). Print the loss \(\ell\).

Hint. Each step is torch.mv(Wk, z_prev). Check .shape at each step.

Background: The Jacobian¶

The gradient \(\partial \ell / \partial W_k\) is a matrix, so it cannot directly be obtained by a vector-matrix product. One systematic approach is to vectorise the weight matrices and compute the Jacobian of \(\ell\) with respect to \(\text{vec}(W_k)\), the entries of \(W_k\) laid out as a long vector.

Using the identity \(\text{vec}(AXB) = (B^T \otimes A)\,\text{vec}(X)\), we can derive that:

Then we multiply each Jacobian by the scalar \(\delta = \partial \ell / \partial z_3 = z_3 - y\) to get the gradient with respect to each vectorised weight.

Think about it. How does the cost of forming the Kronecker products above scale with \(d\)?

1b — Implement and time the full Jacobian¶

Implement the computation of the Kronecker products above. For each weight matrix, form the explicit Jacobian using torch.kron, multiply by \(\delta\), and reshape the result back to the shape of the weight matrix.

delta = (z3 - y).item()   # scalar  ∂ℓ/∂z₃

# Jacobian w.r.t. W3  (simplest case — no Kronecker product needed)
J3 = delta * z2           # shape (d,)  →  reshape to (1, d)

# Jacobian w.r.t. W2
Id = torch.eye(d)
J2 = delta * (W3 @ torch.kron(z1.unsqueeze(0), Id))  # shape (1, d²) → reshape to (d, d)

# Jacobian w.r.t. W1  — fill this in yourself

Hint. After computing a vectorised gradient of shape (1, d*k), recover the matrix gradient with .reshape(d, k). Check that the result has the same shape as the corresponding weight.

Time the full computation and report the wall time.

Think about it. The Kronecker product \(x^T \otimes I_d\) has shape \((d, dn)\). Forming it requires storing \(d^2 n\) numbers. For \(d = 64, n = 20\) that is already 81 920 floats — and the matrix multiply on top costs \(O(d^3 n)\). How does this compare to simply storing the weight matrix itself?

2. Part — Reverse mode by hand (30 points)¶

Background: The cheaper formula¶

The Kronecker-product approach in Part 1 is expensive because it builds large intermediate matrices. But notice that the result of multiplying those matrices by \(\delta\) can always be written as a rank-1 outer product.

Starting from the scalar \(\delta = z_3 - y\), define the upstream gradient vectors by propagating \(\delta\) backward through each layer:

Here \(g_k = \partial \ell / \partial z_k\) is the gradient of the loss with respect to the output of layer \(k\).

The weight gradients are then simply outer products:

Think about it. Each outer product \(g_k z_{k-1}^T\) costs \(O(d^2)\). How does the total cost compare to Part 1?

The pattern to remember: The gradient of the loss with respect to a weight matrix is always:

The upstream gradient \(g_k\) tells you how the loss changes as the output of layer \(k\) changes; the input \(z_{k-1}\) tells you how the output changes as \(W_k\) changes.

2a — Implement and time the reverse pass¶

Implement the reverse pass:

1. Forward pass: compute and store \(z_1, z_2, z_3\) and the loss.
2. Backward pass: compute \(g_3, g_2, g_1\) by propagating backward with matrix–vector products.
3. Weight gradients: compute \(\partial\ell/\partial W_k = \texttt{torch.outer}(g_k, z_{k-1})\) for each layer (taking care of the shapes for \(W_3\)).

Hint 1. \(g_3\) is a length-1 tensor containing \(\delta = z_3 - y\). The update \(g_2 = W_3^T g_3\) is W3.T @ g3 or torch.mv(W3.T, g3), giving shape (d,). Similarly for \(g_1\). These are matrix–vector products, costing \(O(d^2)\) each — not matrix–matrix products.

Hint 2. For \(W_3\) (shape (1, d)), the gradient is torch.outer(g3, z2) which gives shape (1, d). For \(W_2\) and \(W_1\) it is torch.outer(g2, z1) → (d, d) and torch.outer(g1, x) → (d, n).

Time the reverse pass and verify against Part 1:

torch.allclose(dW1_reverse, dW1_jacobian, atol=1e-5)
torch.allclose(dW2_reverse, dW2_jacobian, atol=1e-5)
torch.allclose(dW3_reverse, dW3_jacobian, atol=1e-5)

2b — Scaling experiment¶

Vary \(d \in \{8, 16, 32, 64, 128\}\) (keep \(n = d / 2\) for simplicity). For each value of \(d\), time both methods and record the ratio

Report in a table or plot \(r(d)\). At what \(d\) does the speedup become substantial?

Hint. Keep weight generation outside the timed block. Repeat each measurement several times and take the minimum to reduce noise:

start = time.perf_counter()
# ... your code ...
elapsed = time.perf_counter() - start

3. Part — PyTorch autograd (30 points)¶

3a — Verify with autograd¶

Re-run the forward pass with requires_grad=True on each weight matrix:

W1_a = W1.clone().requires_grad_(True)
W2_a = W2.clone().requires_grad_(True)
W3_a = W3.clone().requires_grad_(True)

Recompute the forward pass using these tensors, call loss.squeeze().backward(), and verify that W1_a.grad, W2_a.grad, W3_a.grad match your Part 2 results.

Hint. After .backward(), Wk_a.grad has the same shape as Wk_a — exactly the shape of the weight matrix. This matches the outer-product result from Part 2.

Common mistake. Gradients accumulate in PyTorch. If you call .backward() more than once, gradients are summed rather than replaced. Recreate the W_a tensors from scratch between runs, or call .grad.zero_() on each before the next backward pass.

3b — Time autograd¶

Time the combined forward + backward pass and compare to Parts 1 and 2.

3c — Reflection¶

In one short paragraph, explain in your own words what PyTorch's autograd engine is doing during the forward pass and what happens when .backward() is called. Use the vocabulary from Parts 1 and 2: computation graph, Jacobian, upstream gradient, outer product, chain rule.

Bonus — Add a nonlinearity (10 points)¶

Introduce ReLU activations:

The Jacobian of \(z_k\) with respect to the pre-activation \(a_k = W_k z_{k-1}\) is \(\text{diag}(\mathbf{1}[a_k > 0])\). This inserts an extra factor into each backward step:

Modify your hand-written reverse pass to include this masking step and re-verify against autograd.

Hint. Store the pre-activation \(a_k = W_k z_{k-1}\) alongside \(z_k\) during the forward pass — you need it to compute the mask (a_k > 0).float(). Apply the mask to the upstream gradient before computing the outer product and before passing the gradient further left: g_k_masked = (ak > 0).float() * g_k, then use g_k_masked in place of g_k for both the weight gradient and the next vector–matrix product.

Final README and reflection¶

Update the README.md file in your submission repository. (If no file named README.md exists in your repository, then create one.)

Add complete log files of ALL your AI conversations – just like you did in the previous exercises. Store them in the files named ai_log.md.

Write "Reflection" and "AI use" statements into the README.

Submission¶

Submit this exercise, as done with all previous exercises, via Canvas and GitLab.

Do not forget to submit the link in Canvas on time and do not edit the GitLab after the link has been submitted to Canvas.

Congrats! You're done! 🍪