Dense Layer Backward Pass

A dense layer computes Y = X @ W + b, mapping a batch of n inputs of width d_in to n outputs of width d_out. During training you are handed dOut, the gradient of the scalar loss with respect to Y, and you must push it back to the layer’s parameters and to its input. This is the single most reused backprop step in deep learning, and interviewers use it to check whether you can derive matmul gradients from the chain rule and get every shape and transpose right.

Definition

For the forward map Y = X @ W + b with X shape (n, d_in), W shape (d_in, d_out), and bias b shape (d_out,) broadcast across the n rows, the chain rule gives three gradients:

dX = dOut @ W.T        # (n, d_out) @ (d_out, d_in) -> (n, d_in)
dW = X.T @ dOut        # (d_in, n)  @ (n, d_out)    -> (d_in, d_out)
db = dOut.sum(axis=0)  # sum over the batch -> (d_out,)

Each rule is forced by shape matching. dX must look like X, so the only legal product is dOut @ W.T. dW must look like W, so it must be X.T @ dOut. The bias was added once per example by broadcasting, so its gradient collapses that broadcast by summing dOut down the batch axis.

Task

Complete dense_backward(dOut, X, W) so it returns the tuple (dX, dW, db) using the three rules above, fully vectorized with NO Python loops over examples or features. Use only NumPy matrix products and a sum; do not call any autodiff or framework helper. The returned shapes must be exactly dX: (n, d_in), dW: (d_in, d_out), db: (d_out,).

Example

X = np.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0]])          # (2, 3) -> n=2, d_in=3
W = np.array([[1.0, 0.0],
              [0.0, 1.0],
              [1.0, 1.0]])               # (3, 2) -> d_out=2
dOut = np.array([[1.0, 1.0],
                 [1.0, 1.0]])            # (2, 2)

dX, dW, db = dense_backward(dOut, X, W)
# dX -> array([[1., 1., 2.],
#              [1., 1., 2.]])
# dW -> array([[5., 5.],
#              [7., 7.],
#              [9., 9.]])
# db -> array([2., 2.])

dX reads each row of W.T summed by the ones in dOut; dW accumulates each input feature against the upstream gradient over both examples; db sums the two rows of dOut.

Constraints

• 1 <= n <= 10^4, 1 <= d_in, d_out <= 10^3; use vectorized NumPy matmuls, no Python loop over rows or columns.
• Inputs are real-valued float64; no NaN or inf. Return the three gradients in the order (dX, dW, db).
• db is a 1-D array of shape (d_out,), not (1, d_out).

Approach

Backprop through Y = X @ W + b is three chain-rule steps, each a matmul or a sum whose form is pinned by shape matching. Since dX must match X’s shape (n, d_in) and only dOut @ W.T produces it, that is the input gradient. Since dW must match W’s shape (d_in, d_out), it must be X.T @ dOut, which accumulates every example’s contribution. The bias was broadcast identically across the n rows, so its gradient sums dOut down the batch axis to shape (d_out,).

Solution

import numpy as np
from typing import Tuple

def dense_backward(
    dOut: np.ndarray, X: np.ndarray, W: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    dOut = np.asarray(dOut, dtype=np.float64)
    X = np.asarray(X, dtype=np.float64)
    W = np.asarray(W, dtype=np.float64)

    dX = dOut @ W.T          # (n, d_out) @ (d_out, d_in) -> (n, d_in)
    dW = X.T @ dOut          # (d_in, n)  @ (n, d_out)    -> (d_in, d_out)
    db = dOut.sum(axis=0)    # sum over the batch axis    -> (d_out,)

    return dX, dW, db

Walkthrough

On the example, X is (2, 3), W is (3, 2), and dOut is all ones with shape (2, 2).

• dX = dOut @ W.T. W.T is [[1,0,1],[0,1,1]], shape (2, 3). Each row of dOut is [1,1], so each output row is [1,1] @ W.T = [1+0, 0+1, 1+1] = [1, 1, 2], giving dX = [[1,1,2],[1,1,2]], shape (2, 3) matching X.
• dW = X.T @ dOut. X.T is (3, 2) with rows [1,4], [2,5], [3,6]. Multiplying by dOut (columns of ones) sums each X.T row across examples: 1+4=5, 2+5=7, 3+6=9, broadcast into both output columns, so dW = [[5,5],[7,7],[9,9]], shape (3, 2) matching W.
• db = dOut.sum(axis=0) = [1+1, 1+1] = [2, 2], shape (2,) matching the bias width.

All three shapes line up with their forward counterparts, the check that confirms the transposes and sum axis are correct.

Complexity & notes

• Time O(n * d_in * d_out) for each of the two matmuls (BLAS runs them in optimized C); the bias sum is O(n * d_out). Space O(n * d_in + d_in * d_out + d_out) for the three gradients, matching the input, weight, and bias shapes.
• The transposes are not a choice: dX must be shaped like X and dW like W, and only one arrangement of dOut, W, and X yields each. Memorizing “gradient w.r.t. the left factor uses the right factor transposed on the right; gradient w.r.t. the right factor uses the left factor transposed on the left” reproduces both rules instantly.
• db sums over axis=0 because the same bias vector was added to all n rows in the forward pass; differentiating a broadcast (a copy) sends gradients back through every copy and adds them. Returning shape (d_out,), not (1, d_out), keeps it aligned with b.
• If the forward pass averaged the loss over the batch rather than summing, the upstream dOut already carries the 1/n factor, so no extra scaling belongs here; keep the layer’s backward pure and let the loss own the averaging.
• With L2 weight decay you would add lambda * W to dW outside this function; the layer’s job is only the data-dependent gradient.