Matrix form of partial derivatives of weights in a neural network

Question

I'm implementing a (vanilla) neural network, from scratch in R, following Hastie & Tibshirani. As they suggest on page 397, I'm doing it via a conjugate gradient optimizer (optim in R).

Here's the code:

y <- mtcars$mpg

X <- scale(as.matrix(mtcars[,-1]))

NZ <- 5 #number of hidden units

alpha <- matrix(runif(NZ*(ncol(X)+1), -.7, .7), ncol = NZ)  #Random starting values for alpha
beta <- runif(NZ+1, -.7, .7)                                #Random starting values for beta
Z <- apply(cbind(1,X) %*% alpha, 2, sigma)                  #Calculate values for hidden layer based on starting values
yhat <- cbind(1, Z) %*% beta                                #estimate of yhat based on starting values

#Transfer function
sigma <- function(v, s = 1){
  1/(1+exp(-v*s))
}

#Objective function to be minimized
getweights <- function(par, X, y){
  beta <- par[1:(NZ+1)]
  alpha <- matrix(par[(NZ+2):length(par)], ncol = NZ)
  Z <- apply(cbind(1,X) %*% alpha, 2, sigma)                    
  yhat <- cbind(1, Z) %*% beta
  mean((y-yhat)^2)
}

#Minimizing it
PT <- proc.time()
out <- optim(par = c(beta, alpha), fn = getweights, method = 'CG', X = X, y = y)
proc.time() - PT 

#using the estimated weights to 
beta <- out$par[1:(NZ+1)]
alpha <- matrix(out$par[(NZ+2):length(out$par)], ncol = NZ)
Z <- apply(cbind(1,X) %*% alpha, 2, sigma)                  #Calculate values for hidden layer based on starting values
yhat <- cbind(1, Z) %*% beta
plot(y, yhat)

Now, optim should work faster if I supply gradient functions. Given that I'm working on a regression problem, these are as follows, following the notation in ESLI: $$ \frac{\partial R_i}{\partial \beta_m} = -2(y_i - f(x_i))g'(\beta_m^Tz_i)z_{mi} = -2(y_i - f(x_i))z_{mi} \text{ (given that $g(x) = x$)} $$ and $$ \frac{\partial R_i}{\partial \alpha_{ml}} = -2(y_i - f(x_i))g'(\beta_m^Tz_i)\beta_m\sigma'(\alpha^T_m x_i)x_{il} = -2(y_i - f(x_i))\beta_m\sigma'(\alpha^T_m x_i)x_{il} $$

Where $R$ is the squared error, $i$ is an observation, and $m$ is a hidden unit.

In matrix form, the first one is simply: $$ \frac{\partial R(\theta)}{\partial \beta} = -2(\mathbf{y - \hat{y}})^T\mathbf{z} $$

where $R(\theta) = \displaystyle\sum_i R_i$.

$$ \frac{\partial R(\theta)}{\partial \alpha} = ??? $$

What is the second? It has got to be a $p \times N_Z$ matrix, because that's the dimension of $\alpha$. But a straightforward translation of the second equation gives me matrices that are not conformable.

Here is a brute-force method of computing it with for-loops -- it is painfully slow:

drda <- array(rep(alpha*NA, length(yhat)), dim = c(nrow(alpha), ncol(alpha), length(yhat)))
for (i in 1:length(yhat)){
  for (m in 1:ncol(alpha)){
    for (l in 1:nrow(alpha)){
      drda[l,m,i] <- -2*(y[i]-yhat[i]) * sigma_prime(X[i,] %*% alpha[,m]) * X[i,l]*beta[m]
    }
  }
}
drda <- apply(drda, c(1,2), sum)

Here is a slightly faster version, looping only over $i$. It is still slow:

library(foreach)
drda <- foreach(i=1:nrow(yhat), .combine = '+') %do% {
  t(-2*(y[i]-yhat[i])*beta * sigma_prime(X[i,] %*% alpha)) %*% t(X[i,])
}

What am I missing?

Answer 1

Your second equation defines $$\frac{\partial R_i}{\partial \alpha_{ml}}$$ which has 3 indices, so is a "tensor" rather than a vector or a matrix. That is, for each component $i$ of the $\boldsymbol{R}$ vector, you have a matrix of partial derivatives with respect to the matrix $\boldsymbol{\alpha}$.

To get a matrix result, you would have to "vectorize" the $\boldsymbol{\alpha}$ matrix in the denominator of your derivative.

This would make sense, as most "generic" optimization algorithms accept a vector input. I do not use R, but in Matlab the optimization functions will accept an arbitrary size array as input, but internally they will flatten it into a vector, and any gradient function supplied will be expected to return "flattened" output consistent with this.