Why is the cost function of neural networks non-convex?

Question

There is a similar thread here (Cost function of neural network is non-convex?) but I was not able to understand the points in the answers there and my reason for asking again hoping this will clarify some issues:

If I am using sum of squared difference cost function, I am ultimately optimizing something of the form $ \Sigma_{i=1}^{N}(y_i - \hat{y_i})^2$ where $y$ is the actual label value during training phase and $\hat{y}$ is the predicted label value. As this has a square form, this should be a convex cost function. So what is it that could make it non-convex in a NN?

Answer 1

$\sum_i (y_i- \hat y_i)^2$ is indeed convex in $\hat y_i$. But if $\hat y_i = f(x_i ; \theta)$ it may not be convex in $\theta$, which is the situation with most non-linear models, and we actually care about convexity in $\theta$ because that's what we're optimizing the cost function over.

For example, let's consider a network with 1 hidden layer of $N$ units and a linear output layer: our cost function is $$ g(\alpha, W) = \sum_i \left(y_i - \alpha_i\sigma(Wx_i)\right)^2 $$ where $x_i \in \mathbb R^p$ and $W \in \mathbb R^{N \times p}$ (and I'm omitting bias terms for simplicity). This is not necessarily convex when viewed as a function of $(\alpha, W)$ (depending on $\sigma$: if a linear activation function is used then this still can be convex). And the deeper our network gets, the less convex things are.

Now define a function $h : \mathbb R \times \mathbb R \to \mathbb R$ by $h(u, v) = g(\alpha, W(u, v))$ where $W(u,v)$ is $W$ with $W_{11}$ set to $u$ and $W_{12}$ set to $v$. This allows us to visualize the cost function as these two weights vary.

The figure below shows this for the sigmoid activation function with $n=50$, $p=3$, and $N=1$ (so an extremely simple architecture). All data (both $x$ and $y$) are iid $\mathcal N(0,1)$, as are any weights not being varied in the plotting function. You can see the lack of convexity here.

Here's the R code that I used to make this figure (although some of the parameters are at slightly different values now than when I made it so they won't be identical):

costfunc <- function(u, v, W, a, x, y, afunc) {
  W[1,1] <- u; W[1,2] <- v
  preds <- t(a) %*% afunc(W %*% t(x))
  sum((y - preds)^2)
}

set.seed(1)
n <- 75  # number of observations
p <- 3   # number of predictors
N <- 1   # number of hidden units


x <- matrix(rnorm(n * p), n, p)
y <- rnorm(n)  # all noise
a <- matrix(rnorm(N), N)
W <- matrix(rnorm(N * p), N, p)

afunc <- function(z) 1 / (1 + exp(-z))  # sigmoid

l = 400  # dim of matrix of cost evaluations
wvals <- seq(-50, 50, length = l)  # where we evaluate costfunc
fmtx <- matrix(0, l, l)
for(i in 1:l) {
  for(j in 1:l) {
    fmtx[i,j] = costfunc(wvals[i], wvals[j], W, a, x, y, afunc)
  }
}

filled.contour(wvals, wvals, fmtx,plot.axes = { contour(wvals, wvals, fmtx, nlevels = 25, 
                                           drawlabels = F, axes = FALSE, 
                                           frame.plot = FALSE, add = TRUE); axis(1); axis(2) },
               main = 'NN loss surface', xlab = expression(paste('W'[11])), ylab = expression(paste('W'[12])))