Avoiding tensors when differentiating with respect to weight matrices in backpropagation

Question

Consider a neural network consisting of only a single affine transformation with no non-linearity. Use the following notation:

• $\textbf{Inputs}: x \in \mathbb{R}^n$
• $\textbf{Weights}: W \in \mathbb{R}^{h \times n}$
• $\textbf{Biases}: b \in \mathbb{R}^h$
• $\textbf{Output}: z = W x + b \in \mathbb{R}^h$
• $\textbf{Loss function}: \ell(z) \in \mathbb{R}$

We seek to calculate the gradient $\frac{\partial \ell}{\partial W}$ of the loss with respect to the weights. In doing backpropogation, we break this gradient into two factors: $$ \frac{\partial \ell}{\partial W} = \frac{\partial \ell}{\partial z} \cdot \frac{\partial z}{\partial W} $$ However, $\frac{\partial z}{\partial W}$ is the derivative of a vector with respect to matrix, while will result in a 3 dimensional tensor. How can we keep compute $\frac{\partial \ell}{\partial W}$ while keeping things in two dimensions?

Answer 1

Because this is purely a question of notation, it helps to examine it abstractly.

Multivariate derivatives

Every finite-dimensional real vector space $\mathbb R^d$ can be considered a set of functions $\mathbf v:S\to \mathbb R$ where $S$ is a set with $d$ elements (the "index set"), endowed with its usual pointwise addition and scalar multiplication. It is canonically equipped with $d$ coordinate functions defined by restricting $\mathbf v$ to a single element of $S.$ That is, for $s\in S,$

$$\pi_s(\mathbf v) = \mathbf v(s).$$

Your situation concerns the composition of two differentiable maps

$$\mathbb R^{nh} \overset{g}\to \mathbb R^h \overset{l}\to \mathbb R \approx \mathbb R^1$$

where

$$g(W) = Wx+b.$$

The Chain Rule says the derivative of $l\circ g$ at some point $W\in\mathbb R^{nh}$ is the composition of their derivatives evaluated at relevant points,

$$D(l\circ g)_W = D(l)_{g(W)}\circ D(g)_W.$$

Calculating derivatives

The final preliminary result about differentiation is the basic theorem that these derivatives, in terms of the coordinate functions, are the Jacobian matrices of partial derivatives--whence the composition on the right hand side is computed via matrix multiplication. I don't know that this theorem even has a name--and I suspect most people forget it's even a theorem, it is so basic.

You ask how to calculate these derivatives. It is aided by the Kronecker delta. This can be considered a notational convenience given by

$$\delta_{t,s} = \delta_{s,t} = \mathcal I(s=t) = \left\{\matrix{1, & s=t\\0, & s\ne t}\right.$$

for any two indices $s$ and $t.$ Delta simplifies sums wherever it involves a summation index:

$$\sum_{i} a_i \delta_{i,j} = a_j$$

is a model of how it works. You will see this applied twice below.

Comment. One could maintain that $\delta$ is a tensor (and it indeed can be formalized as such), but as you will see, it is employed in the solution below merely as a notational convenience for indicating whether or not two indices are equal. This keeps the conceptual focus squarely on the fact that we are computing the derivative of the composition of two maps of vector spaces -- no tensors need get involved.

The solution

We may now apply the Chain Rule mindlessly. But for those to whom this approach might be unfamiliar, I will spell out all the details.

Let $i$ be an index for $\mathbb R^h$ and $j$ an index for $\mathbb R^n,$ so that we may use the ordered pairs $(i,j)$ to index the components of the matrix $W.$ Let $k$ be an index for $\mathbb R^h$ to index the components of the vector $z.$ Finally, let $s$ be an index for $\mathbb R^n$ used for the components of the vector $x.$ Here is the calculation:

$$\begin{aligned} \frac{\partial (l \circ z)}{\partial W_{ij}} &= \sum_{k=1}^h \frac{\partial l}{\partial z_k} \frac{\partial z_k}{\partial W_{ij}}&\text{Chain Rule}\\ &=\sum_{k=1}^h \frac{\partial l}{\partial z_k}\ \frac{\partial}{\partial W_{ij}}\left(\sum_{s=1}^n W_{ks}x_s + b_k\right)&\text{Expand }z = Wx + b\\ &=\sum_{k=1}^h \frac{\partial l}{\partial z_k}\ \sum_{s=1}^n \left(x_s\, \frac{\partial W_{ks}}{\partial W_{ij}} + \frac{\partial b_k}{\partial W_{ij}}\right)&\text{Linearity of differentiation}\\ &=\sum_{k=1}^h \frac{\partial l}{\partial z_k} \ \sum_{s=1}^n x_s\,(\delta_{ks, \,ij} + 0)&\text{Derivatives expressed as } \delta\\ &=\sum_{k=1}^h \frac{\partial l}{\partial z_k} \ \sum_{s=1}^n x_s\,\delta_{k,i}\,\delta_{j,s}&\text{Separability}\\ &=\left(\sum_{k=1}^h \frac{\partial l}{\partial z_k} \delta_{k,i}\right)\left(\sum_{s=1}^n x_s\,\delta_{j,s}\right)&\text{Arithmetic}\\ &=\frac{\partial l}{\partial z_i} x_j.&\delta\text{ simplification of sums} \end{aligned}$$

Answer 2

Update: here is a better solution with slightly improved notation:

Let $f(W) = \ell(z(W)) = \ell(Wx+b)$, treating $x$ and $b$ as constants. Then proceed using differentials. We see immediately that $dz = dW x$. Further, $$df = \frac{\partial \ell}{\partial z} dz = \frac{\partial \ell}{\partial z} dW x$$ Since $df$ is a scalar we may take the trace of both sides and rearrange using the cyclic property: $$df = \text{tr}(df) = \text{tr} \left(\frac{\partial \ell}{\partial z} dW x \right) = \text{tr}\left(x \frac{\partial \ell}{\partial z} dW \right)$$ The expression $x \frac{\partial \ell}{\partial z}$ is $n \times h$ while $dW$ is $h \times n$ just as $W$ was. Therefore this is a Frobenius inner product $$ df = \left\langle \frac{\partial \ell}{\partial z}^\top x^\top, dW \right\rangle_F $$ and by the relation of inner products of differentials to "derivatives," using numerator layout, $$ \frac{\partial f}{\partial W} = x \frac{\partial \ell}{\partial z} $$

---

Longer solution using components

We adopt the numerator layout (Jacobian formulation), where the gradient of a scalar with respect to a vector is a row vector by convention.

Convert to component notation

First, let us write $ \dfrac{\partial \ell}{\partial z} $ in component notation as follows: $$ \frac{\partial \ell}{\partial z} = \begin{bmatrix} \dfrac{\partial \ell}{\partial z_1} & \dfrac{\partial \ell}{\partial z_2} & \dots & \dfrac{\partial \ell}{\partial z_h} \end{bmatrix} \in \mathbb{R}^{1 \times h} $$ and write $z$ in component form as: $$ z_i = \sum_{j=1}^n W_{ij} x_j + b_i $$ Taking the partial derivative of $ z_k $ with respect to $ W_{ij} $ we find that $$ \frac{\partial z_k}{\partial W_{ij}} = \begin{cases} x_j & \text{if } k = i \\ 0 & \text{if } k \neq i \end{cases} $$

Calculate $ \dfrac{\partial \ell}{\partial W} $

Using the chain rule: $$ \left( \frac{\partial \ell}{\partial W} \right)_{ij} = \sum_{k=1}^h \frac{\partial \ell}{\partial z_k} \frac{\partial z_k}{\partial W_{ij}} = \frac{\partial \ell}{\partial z_i} x_j $$ In numerator layout, $$ \left(\frac{\partial l}{\partial W}\right)_{ij} = \frac{\partial l}{\partial W_{ji}}. $$ where the $i$ and $j$ appear in swapped positions between the two sides of the equality. Substituting the expression from above (but swapping the indices $i$ and $j$ on the right-hand side as required) yields $$ \left(\frac{\partial l}{\partial W}\right)_{ij} = \frac{\partial l}{\partial z_j} x_i $$ To recognize this as an outer product, note that:

• $\frac{\partial l}{\partial z}$ is a $1 \times h$ row vector in numerator layout
• $x$ is an $n \times 1$ column vector

The $(i,j)$ component of the outer product $x \left(\frac{\partial l}{\partial z} \right)$ is precisely $x_i \frac{\partial l}{\partial z_j}$ Therefore, in matrix form: $$ \frac{\partial l}{\partial W} = x\frac{\partial l}{\partial z} $$ This expression is an $n \times h$ matrix, equal in dimension to $W^T$, as required by numerator layout.