Affine layer - gradient shape

Question

In course cs231n, I need to implement backward pass computation for an affine (linear) layer:

def affine_backward(dout, cache):
    """
    Computes the backward pass for an affine layer.

    Inputs:
    - dout: Upstream derivative, of shape (N, M)
    - cache: Tuple of:
      - x: Input data, of shape (N, d_1, ... d_k)
      - w: Weights, of shape (D, M)

    Returns a tuple of:
    - dx: Gradient with respect to x, of shape (N, d1, ..., d_k)
    - dw: Gradient with respect to w, of shape (D, M)
    - db: Gradient with respect to b, of shape (M,)
    """
    x, w, b = cache
    dx, dw, db = None, None, None

I do not understand why the shape of dw is (D, M), as the output of the layer is a matrix (N, M) - N being batch size. This would only make sense if the output is a scalar.

What am I missing?

Thanks in advance.

Answer 1

Your weight tensor is a 2d matrix, in this case (D,M). Therefore the gradient dw of this will also be a (D,M) matrix. The shape of the output of the layer is going to be different because when you multiply a (N,D) input by a (D,M) weight matrix the result will be (N,M). In other words, multiply an input matrix by a weight matrix will (sometimes but not always) give a different shape tensor. Taking the derivative of a matrix won't change its shape.