Probabilistic bounds on neuron death in deep rectifier networks

Despite the explosion of interest in deep learning over the last decade, network design remains largely experimental. Engineers are faced with a wide variety of design choices, including preprocessing of the input data, the number of layers in the network, the use of residual and skip connections between the layers, the width or number of neurons in each layer, and which types of layers are used. Even with modern parallel processing hardware, training a deep neural network can take days or even weeks depending on the size and nature of the input data. The problem is particularly acute in applications like biomedical image analysis [13]. This imposes practical limits on how much experimentation can be done, even with the use of automatic model searching tools [11]. In these cases, theoretical analysis can shed light on which models are likely to work a priori, reducing the search space and accelerating innovation.

This paper explores the model design space through the lens of neuron death. We focus on the rectified linear unit (ReLU), which is the basic building block of the majority of current deep learning models [6, 7, 15]. The ReLU neuron has the interesting property that it maps some data points to a constant function, at which point they do not contribute to the training gradient. The property that certain neurons focus on certain inputs may be key to the success of ReLU neurons over the sigmoid type. However, a ReLU neuron sometimes maps all data to a constant function, in which case we say that it is dead. If all the neurons in a given layer die, then the whole network is rendered untrainable. In this work, we show that neuron death can help explain why certain model architectures work better than others. While it is difficult to compute the probability of neuron death directly, we derive upper and lower bounds by a symmetry argument, showing that these bounds rule out certain model architectures as intractable.

Neuron death is a well-known phenomenon which has inspired a wide variety of research directions. It is related to the unique geometry of ReLU networks, which have piecewise affine or multi-convex properties depending on whether the inference or training loss function is considered [2, 18, 20]. Many activation functions have been designed to preclude neuron death, yet the humble ReLU continues to be preferred in practice, suggesting these effects are not completely understood [3, 4, 8, 9, 19]. Despite widespread knowledge of the phenomenon, there are surprisingly few theoretical works on the topic. Several works analyzed the fraction of living neurons in various training scenarios. Wu et al. studied the empirical effects of various reinforcement leaning schemes on neuron death, deriving some theoretical bounds for the single-layer case [27]. Ankevist et al. studied the effect of different training algorithms on neuron death, both empirically as well as theoretically using a differential equation as a model for network behavior [1]. To our knowledge, a pair of recent works from Lu and Shin et al. were the first to discuss the problem of an entire network dying at initialization, and the first to analyze neuron death from a purely theoretical perspective [17, 24]. They argued that, as a ReLU network architecture grows deeper, the probability that it is initialized dead goes to one. If a network is initialized dead, the partial derivatives of its output are all zero, and thus it is not amenable to differential training methods such as stochastic gradient descent [16]. This means that, for a bounded width, very deep ReLU networks are nearly impossible to train. Both works propose upper bounds on the probability of living initialization, suggesting new initialization schemes to improve this probability. However, there is much left to be said on this topic, as bounds are complex and the derivations include some non-trivial assumptions.

In this work, we derive simple upper and lower bounds on the probability that a random ReLU network is alive. Our proofs are based on a rigorous symmetry argument requiring no special assumptions. Our upper bound rigorously establishes the result of Lu et al., while our lower bound establishes a new positive result, that a network can grow infinitely deep so long as it grows wider as well [17]. We show that the true probability agrees with our bounds in the extreme cases of a single-layer network with the largest possible input set, or a deep network with the smallest possible input set. We also show that the true probability converges to our lower bound along any path through hyperparameter space such that neither the width nor depth is bounded. Our proof of the latter claim requires an assumption about the output variance of the network. We also note that our lower bound is exactly the probability that a single data point is alive. All of these results are confirmed by numerical simulations.

Finally, we discuss how information loss by neuron death furnishes a compelling interpretation of various network architectures, such as residual layers (ResNets), batch normalization, and skip connections [10, 12, 23]. This analysis provides a priori means of evaluating various model architectures, and could inform future designs of very deep networks, as well as their biological plausibility. Based on this information, we propose a simple sign-flipping initialization scheme guaranteeing with probability one that the ratio of living training data points is at least $2^{-k}$, where $k$ is the number of layers in the network. Our scheme preserves the marginal distribution of each parameter, while modifying the joint distribution based on the training data. We confirm our results with numerical simulations, suggesting the actual improvement far exceeds the theoretical minimum. We also compare this scheme to batch normalization, which offers similar, but not identical, guarantees.

The contributions of this paper are summarized as follows:

1. 1.

   New upper and lower bounds on the probability of ReLU network death,

2. 2.

   proofs of the optimality of these bounds,

3. 3.

   interpretation of various neural network architectures in light of the former, and

4. 4.

   a tractable initialization scheme preventing neuron death.

Given an input dimension $d$, with weights $a\in\mathbb{R}^{d}$ and $b\in\mathbb{R}$, and input data $x\in\mathbb{R}^{n}$, a ReLU neuron is defined as

A ReLU layer of with $n$ is just the vector concatenation of $n$ neurons, which can be written $g(x)=\max\{Ax+b,0\}$ with parameters $A\in\mathbb{R}^{n\times n}$ and $b\in\mathbb{R}^{n}$, where the maximum is taken element-wise. A ReLU network with $k$ layers is $F_{k}(x)=g_{k}\circ\dots\circ g_{1}(x)$, the composition of the $k$ layers. The parameters of a network are denoted $\theta=(A_{1},b_{1},\dots,A_{n},b_{n})$, a point in $\mathbb{R}^{N(n,k)}$ where $N(n,k)=k(n^{2}+n)$ is the total number of parameters in the network. To simplify the proofs and notation, we assume that the width of each layer is the same throughout a network, always denoted $n$.

In practice, neural networks parameters are often initialized from some random probability distribution at the start of training. This distribution is important to our results. In this work, as in practice, we always assume that $\theta$ follows a symmetric, zero-mean probability density function (PDF). That is, the density of a parameter vector $\theta$ is not altered by flipping the sign of any component. Furthermore, all components of $\theta$ are assumed to be statistically independent and identically distributed (IID), except where explicitly stated otherwise. We sometimes follow the practical convention that $b=0$ for all layers, which sharpens the upper bound slightly, but most of our results hold either way.

We sometimes refer to the response of a network layer before the rectifying nonlinearity. The pre-ReLU response of the $k^{th}$ layer is denoted $\tilde{F}_{k}(x)$, and consists of the first $k-1$ layers composed with the affine part of the $k^{th}$ layer, without the final maximum. For short-hand, the response of the $l^{th}$ neuron in the $k^{th}$ layer is denoted $F_{k.l}$, while the pre-ReLU response is $\tilde{F}_{k.l}$.

Let $S_{0}\subseteq\mathbb{R}^{n}$ be the input domain of the network. Let $S_{k}=f_{k}\circ\dots\circ f_{1}(S_{0})$, the image of $S_{0}$ under the $k^{th}$ layer, which gives the domain of $F_{k+1}$. Let the random variable $x\in S_{0}$ denote the input datum, the precise distribution of which is not relevant for this paper. We say that $x$ is dead at layer $k$ if the Jacobian $\mathcal{D}F_{k}(x)$ is the zero matrix,¹¹1For convenience, assume the convention that a partial derivative is zero whenever $\tilde{F}_{k,l}(x)=0$, in which case it would normally be undefined. taken with respect to the parameters $(A_{k},b_{k})$. By the chain rule, $x$ is then dead in any layer after $k$ as well. Note that if $b=0$, this is equivalent to the statement that $F_{k,l}$ is the zero vector, while for $b\neq 0$ it can be any constant.

The dead set of a ReLU neuron is

which is a half-space in $\mathbb{R}^{n}$. The dead set of a layer is just the intersection of the half-space of each neuron, a convex polytope, which is possibly empty or even unbounded. For the sake of simplicity, this paper follows the convention that the number of features $n$ is identical in each layer. In this case, the dead set is actually an affine cone, as there is exactly one vertex which is the intersection of $n$ linear equations in $\mathbb{R}^{n}$, except for degenerate cases with probability zero. Furthermore, if we follow the usual practice that $b=0$ at initialization, then the dead set is a convex cone with vertex at the origin, as seen in figure 1.

For convenience, we denote the dead set of the $l^{th}$ neuron in the $k^{th}$ layer as $D_{k,l}$, while the dead set of the whole $k^{th}$ layer is $D_{k}=\cap_{l=1}^{n}D_{k,l}$. If $\tilde{F}_{k}(x)\in D_{k}$, and $k$ is the least layer in which this occurs, we say that $x$ is killed by that layer. A dead neuron is one for which all of $S_{0}$ is dead. In turn, a dead layer is one in which all neurons are dead. Finally, a dead network is one for which the final layer is dead, which is guaranteed if any intermediate layer dies.

We use $P(n,k)$ to denote the probability that a network of width $n$ features and depth $k$ layers is alive, i.e. not dead, the estimation of which is the central study of this work. For convenience in some proofs, we use $A_{k}\subset\mathbb{R}^{N(n,k)}$ denote the event that the $k^{th}$ layer of the network is alive, a function of the parameters $\theta$. Under this definition, $P(n,k)=P(A_{k})$.

First, we derive a conservative upper bound on $P(n,k)$. We know that $S_{k}\subseteq\mathbb{R}_{+}^{n}$ for all $k>1$, due to the ReLU nonlinearity. Thus, if layer $k>1$ kills $\mathbb{R}_{+}^{n}$, then it also kills $S_{k}$ . Now, $R_{+}^{n}\subseteq D_{k}$ if and only if none of the $n^{2}+n$ parameters are positive. Since the parameters follow symmetric independent distributions, we have $P(\mathbb{R}_{+}^{n}\subseteq D_{k})=2^{-n^{2}-n}$.

Next, note that layer $k$ is alive only if $k-1$ is, or more formally, $A_{k}\subset A_{k-1}$. The law of total probability yields the recursive relationship

Then, since $\mathbb{R}_{+}^{n}\subseteq D_{k}$ is independent of $A_{k-1}$, we have the upper bound

Note that the bound can be sharpened to an exponent of $-n^{2}$ if the bias terms are initialized to zero. For fixed $n$, this is a geometric sequence in $k$. Note the limit is zero, verifying the claim of Lu et al [17].

The previous section showed that deep networks die in probability. Luckily, we can bring our networks back to life again if they grow deeper and wider simultaneously. Our insight is to show that, while a deeper network has a higher chance of dead initialization, a wider network has a lower chance. In what follows, we derive a lower bound for $P(n,k)$ from which we compute the minimal width for a given network depth. The basic idea is this: in a living network, $S_{k}\neq\emptyset$ for all $k$. Recall from equation 2 that the dead set of a neuron is a sub-level set of a continuous function, thus it is a closed set [14]. Given any point $x\in\mathbb{R}^{n}$, $x\notin D_{k}$ implies that there exists some neighborhood around $x$ which is also not in $D_{k}$. This means that a layer is alive so long as it contains a single living point, so a lower bound for the probability of a living layer is the probability of a single point being alive.

It remains to compute the probability of a single point being alive, the compliment of which we denote by $\gamma=\inf_{x\in\mathbb{R}^{n}}P(x\in D)$, where $D$ is the dead set of some neuron. Surprisingly, $P(x\in D)$ does not depend on the value of $x$. Given some symmetric distribution $\rho(a,b)$ of the parameters, we have

Now, the surface $a\cdot x+b=0$ can be rewritten as $(a\,b)\cdot(x,1)=0$, a hyperplane in $\mathbb{R}^{2(n+1)}$. Since this surface has Lebesgue measure zero, its contribution to the integral is negligible. Combining this with the definition of a PDF, we have

Then, change variables to $(\tilde{a},\tilde{b})=(-a,-b)$, with Jacobian determinant $(-1)^{2(n+1)}=1$, and apply the symmetry of $f$ to yield

Thus $\gamma=1-\gamma$, so $\gamma=1/2$. The previous calculation can be understood in a simple way devoid of any formulas: for any half-space containing $x$, the closure of its compliment is equally likely and with probability one, exactly one of these sets contains $x$. Thus, the probability of drawing a half-space containing $x$ is the same as the probability of drawing one that does not, so each must have probability $1/2$.

Finally, if any of the $n$ neurons in layer $k$ is alive at $x$, then the whole layer is alive. Since these are independent events with probability $1/2$, $P(n,k|A_{k-1})\geq 1-2^{-n}$ . It follows from equation 3 that

From this inequality, we can compute the width required to achieve $P(n,k)=p$ as

See figure 2 for plots of these curves for various values of $p$.

Combining the results of the previous sections, we have the bounds

Let $\ell(n,k)$ denote the lower bound, and $\upsilon(n,k)$ the upper. We now show that these bounds are asymptotically tight, as is later verified by our experiments. More precisely, $\upsilon(n,1)=P(n,1)$ for any $n$, with $S_{0}=\mathbb{R}^{n}$ . Furthermore, $P(n,k)\rightarrow\ell(n,k)$ along any path $n(k)$ such that $\lim_{k\rightarrow\infty}n(k)=\infty$. This means that these bounds are optimal for the extreme cases of a single-layer network, or a very deep network.

We now confirm our previously established bounds on neuron death by numerical simulation. We followed the straightforward methodology of generating a number of random networks, applying these to random data, and computing various statistics on the output. All experiments were conducted using SciPy and required several hours’ compute time on a standard desktop CPU [26].

First we generated the random data and networks. Following the typical convention of using white training data, we randomly generated $1024$ data points from standard normal distributions in $\mathbb{R}^{n}$, for each feature dimensionality $n\in\{1,\dots,15\}$ [25]. Then, we randomly initialized $1024$ fully-connected ReLU networks having $1\leq k\leq 256$ layers. Following the popular He initialization scheme, we initialized the multiplicative network parameters with normal distributions having variance $\sigma_{\theta}^{2}=2/n$, while the additive parameters were initialized to zero [9]. Then, we computed the number of living data points for each network, concluding that a network has died if no living data remain at the output. The results are shown in figure 4 alongside the upper and lower bounds from inequality 10.

Next, to show asymptotic tightness of the lower bound, we conducted a similar experiment using the hyperparameters defined by equation 9, rounding $n(k)$ upwards to the next integer, for various values of $p$. These plots ensure that the lower bound is approximately constant at $p$, while the upper bound is approximately equal to one, that $P(n,k)$ converges to a nontrivial value. The results are shown in figure 5.

We draw two main conclusions from these graphs. First, the simulations agree with our theoretical bounds, as $\ell(n,k)\leq P(n,k)\leq\nu(n,k)$ in all cases. Second, the experiments support asymptotic tightness of the lower bound, as $P(n,k)\rightarrow\ell(n,k)$ with increasing $n$ and $k$. Interestingly, the variance decreases exponentially in figure 5, despite the fact that He initialization was designed to stabilize it [9]. This could be explained by the extra term in lemma 5.2. However, we do not notice exponential decay after normalization by the mean $|\lambda|^{2}$, as would be required by assumption 5.6. Recall that assumption 5.6 is a strong claim meant to furnish a sufficient, yet unnecessary condition that $P(n,k)\rightarrow\nu(n,k)$. That is, the lower bound might still be optimal even if the assumption is violated, as suggested by the graphs.

Our previous experiments and theoretical results demonstrated some of the issues with IID ReLU networks. Now, we propose a slight deviation from the usual initialization strategy which partially circumvents the issue of dead neurons. As in the experiments, we switch to a discrete point of view, using a finite dataset $x_{1},\dots,x_{M}\in\mathbb{R}^{n}$. Our key observation is that, with probability one, negating the parameters of a layer revives any data points killed by it. Thus, starting from the first layer, we can count how many data points are killed, and if this exceeds half the number which were previously alive, we negate the parameters of that layer. This alters the joint parameter distribution based on the available data, while preserving the marginal distribution of each parameter. Furthermore, this scheme is practical, costing essentially the same as a forward pass over the dataset.

For a $k$-layer network, this scheme guarantees that at least $\left\lfloor M/2^{k}\right\rfloor$ data points live. The caveat is that the pre-ReLU output cannot be all zeros, or else negating the parameters would not change the output. However, this edge case has probability zero since $\tilde{F}_{k}$ depends on a polynomial in the parameters of layer $k$, the roots of which have measure zero. Batch normalization provides a similar guarantee: if $\mathbb{E}[\tilde{F}_{k}(x)|\theta]=0$ for all $k$, then with probability one there is at least one living data point [12]. Similar to our sign-flipping scheme, this prevents total network death while still permitting individual data points to die.

These two schemes are simulated and compared in figure 6. Note that for IID data $P(\alpha)=\ell(n,k)$, as that bound comes from the probability that any random data point is alive. Sign-flipping significantly increased the proportion of living data, far exceeding the $2^{-k}$ fraction that is theoretically guaranteed. In contrast, batch normalization ensured that the network was alive, but did nothing to increase the proportion of living data over our IID baseline.

Neuron death is not only relevant for understanding how to initialize a very deep neural network, but also helps to explain various aspects of the model architecture itself. Equation 9 in section 4 tells us how wide a $k$-layer network needs to be; if a network is deep and skinny, most of the data will die. We have also seen in section 7 that batch normalization prevents neuron death, supplementing the existing explanations of its efficacy [22]. Surprisingly, many other innovations in network design can also be viewed through the lens of neuron death.

Up to this point we focused on fully-connected ReLU networks. Here we briefly discuss how our results generalize to other network types. Interestingly, many of the existing network architectures bear relevance in preventing neuron death, even if they were not originally designed for this purpose. What follows is meant to highlight the major categories of feed-forward models, but our list is by no means exhaustive.

We have established several theoretical results concerning ReLU neuron death, confirmed these results experimentally, and discussed their implications for network design. The main results were an upper and lower bound for the probability of network death $P(n,k)$, in terms of the network width $n$ and depth $k$. We provided several arguments for the asymptotic tightness of these bounds. From the lower bound, we also derived a formula relating network width to depth in order to guarantee a desired probability of living initialization. Finally, we showed that our lower bound coincides with the probability of data death, and developed a sign-flipping initialization scheme to reduce this probability.

We have seen that neuron death is a deep topic covering many aspects of neural network design. This article only discusses what happens at random initialization, and much more could be said about the complex dynamics of training a neural network [1, 27]. We do not claim that neuron death offers a complete explanation of any of these topics. On the contrary, it adds to the list of possible explanations, one of which may someday rise to preeminence. Even if neuron death proves unimportant, it may still correlate to some other quantity with a more direct relationship to model efficacy. In this way, studying neuron death could lead to the discovery of other important factors, such as information loss through a deep network. At any rate, the theory of deep learning undoubtedly lags behind the recent advances in practice, and until this trend reverses, every avenue is deserving of exploration.