9
$\begingroup$

Let $Z \sim \mathcal N(\mu, \Sigma)$ be in $\mathbb R^d$. What are the mean and covariance matrix of $Z_+ = \max(0, Z)$ (with the max computed elementwise)?

This comes up e.g. because, if we use the ReLU activation function inside a deep network, and assume via the CLT that the inputs to a given layer are approximately normal, then this is the distribution of the outputs.

(I'm sure many people have computed this before, but I couldn't find the result listed anywhere in a reasonably readable way.)

$\endgroup$
4
  • $\begingroup$ It would simplify your answer--perhaps greatly--to observe that you can obtain it by combining the results of two separate questions: (1) what are the moments of a truncated Normal distribution and (2) what are the moments of a mixture? The latter is straightforward and all you need to do is cite results for the former. $\endgroup$
    – whuber
    Feb 1 '18 at 21:45
  • $\begingroup$ @whuber Hmm. Though I didn't say it explicitly, that's essentially what I do in my answer, except that I didn't find results for a truncated bivariate distribution with a general mean and variance and so had to do some scaling and shifting anyway. Is there some way to derive e.g. the covariance without doing the amount of algebra I had to do? I'm certainly not claiming that anything in this answer is novel, just that the algebra was tedious and error-prone, and perhaps someone else will find the solution useful. $\endgroup$
    – Danica
    Feb 1 '18 at 21:53
  • $\begingroup$ Right: I'm sure your algebra is tantamount to what I described, so it looks like we share an appreciation for simplifying the algebra is possible. One easy way to reduce the algebra is to standardize the diagonal elements of $\Sigma$ to unity, because all that does is to establish a unit of measurement for each variable. At that point you can directly plug Rosenbaum's results into the (simple, obvious) expressions for moments of mixtures. Whether that's even worth algebraic simplification might be a matter of taste: without simplification, it leads to a simple, modular computer program. $\endgroup$
    – whuber
    Feb 1 '18 at 21:58
  • 1
    $\begingroup$ I suppose one could write a program that computes moments directly with Rosenbaum's results and mixing appropriately, and then shifts and scales them back into the original space. That probably would have been faster than the way I did it. $\endgroup$
    – Danica
    Feb 1 '18 at 22:01
7
$\begingroup$

We can first reduce this to depend only on certain moments of univariate/bivariate truncated normal distributions: note of course that $ \DeclareMathOperator{\E}{\mathbb E} \DeclareMathOperator{\Var}{Var} \DeclareMathOperator{\Cov}{Cov} \newcommand{\N}{\mathcal N} \newcommand{\T}{\tilde} \newcommand{\v}{\mathcal V} $

\begin{gather} \E[Z_+] = \begin{bmatrix} \E[(Z_i)_+] \end{bmatrix}_i \\ \Cov(Z_+) = \begin{bmatrix} \Cov\left( (Z_i)_+, (Z_j)_+ \right) \end{bmatrix}_{ij} ,\end{gather} and because we're making coordinate-wise transformations of certain dimensions of a normal distribution, we only need to worry about the mean and variance of a 1d censored normal and the covariance of two 1d censored normals.


We'll use some results from

S Rosenbaum (1961). Moments of a Truncated Bivariate Normal Distribution. JRSS B, vol 23 pp 405-408. (jstor)

Rosenbaum considers $$ \begin{bmatrix}\T X \\ \T Y\end{bmatrix} \sim \N\left( \begin{bmatrix}0 \\ 0\end{bmatrix}, \begin{bmatrix}1 & \rho \\ \rho & 1\end{bmatrix} \right) ,$$ and considers truncation to the event $\v = \{ \T X \ge a_X, \T Y \ge a_Y \}$.

Specifically, we'll use the following three results, his (1), (3), and (5). First, define the following: \begin{gather} q_x = \phi( a_x) \qquad q_y = \phi( a_y) \\ Q_x = \Phi(-a_x) \qquad Q_y = \Phi(-a_y) \\ R_{xy} = \Phi\left( \frac{\rho a_x - a_y}{\sqrt{1 - \rho^2}} \right) \qquad R_{yx} = \Phi\left( \frac{\rho a_y - a_x}{\sqrt{1 - \rho^2}} \right) \\ r_{xy} = \frac{\sqrt{1-\rho^2}}{\sqrt{2 \pi}} \phi\left( \sqrt{\frac{h^2 - 2 \rho h k + k^2}{1 - \rho^2}} \right) \end{gather}

Now, Rosenbaum shows that: \begin{align} \Pr(\v) \E[\T X \mid \v] &= q_x R_{xy} + \rho q_y R_{yx} \tag{1} \\ \Pr\left(\v \right) \E\left[\T X^2 \mid \v \right] &= \Pr\left(\v \right) + a_x q_x R_{xy} + \rho^2 a_y q_y R_{yx} + \rho r_{xy} \tag{3} \\ \Pr(\v) \E\left[ \T X \T Y \mid \v \right] &= \rho \Pr(\v) + \rho a_x q_x R_{xy} + \rho a_y q_y R_{yx} + r_{xy} \tag{5} .\end{align}

It will be useful to also consider the special case of (1) and (3) with $a_y = -\infty$, i.e. a 1d truncation: \begin{align} \Pr(\v) \E[\T X \mid \v] &= q_x \tag{*} \\ \Pr\left(\v \right) \E\left[\T X^2 \mid \v \right] &= \Pr\left(\v \right) = Q_x \tag{**} .\end{align}


We now want to consider \begin{align} \begin{bmatrix}X \\ Y\end{bmatrix} &= \begin{bmatrix}\mu_x\\\mu_y\end{bmatrix} + \begin{bmatrix}\sigma_x & 0 \\ 0 & \sigma_y\end{bmatrix}\begin{bmatrix}\T X \\ \T Y\end{bmatrix} \\&\sim \N\left( \begin{bmatrix}\mu_X\\\mu_Y\end{bmatrix}, \begin{bmatrix}\sigma_x^2 & \rho \sigma_x \sigma_y \\ \rho \sigma_x \sigma_y & \sigma_y^2 \end{bmatrix} \right) \\&= \N\left( \mu, \Sigma \right) .\end{align}

We will use $$ a_x = \frac{-\mu_x}{\sigma_x} \qquad a_y = \frac{-\mu_y}{\sigma_y} ,$$ which are the values of $\T X$ and $\T Y$ when $X = 0$, $Y = 0$.

Now, using (*), we obtain \begin{align} \E[ X_+ ] &= \Pr(X_+ > 0) \E[X \mid X > 0] + \Pr(X_+=0) \, 0 \\&= \Pr(X > 0) \left( \mu_x + \sigma_x \E[\T X \mid \T X \ge a_x] \right) \\&= Q_x \mu_x + q_x \sigma_x ,\end{align} and using both (*) and (**) yields \begin{align} \E[ X_+^2 ] &= \Pr(X_+ > 0) \E[X^2 \mid X > 0] + \Pr(X_+=0) 0 \\&= \Pr\left(\T X \ge a_x\right) \E\left[(\mu_x + \sigma_x \T X)^2 \mid \T X \ge a_x\right] \\&= \Pr\left(\T X \ge a_x\right) \E\left[\mu_x^2 + \mu_x \sigma_x \T X + \sigma_x^2 \T X^2 \mid \T X \ge a_x\right] \\&= Q_x \mu_x^2 + q_x \mu_x \sigma_x + Q_x \sigma_x^2 \end{align} so that \begin{align} \Var[X_+] &= \E[X_+^2] - \E[X_+]^2 \\&= Q_x \mu_x^2 + q_x \mu_x \sigma_x + Q_x \sigma_x^2 - Q_x^2 \mu_x^2 - q_x^2 \sigma_x^2 - 2 q_x Q_x \mu_x \sigma_x \\&= Q_x (1 - Q_x) \mu_x^2 + (1 - 2 Q_x) q_x \mu_x \sigma_x + (Q_x - q_x^2) \sigma_x^2 .\end{align}

To find $\Cov(X_+, Y_+)$, we will need \begin{align} \E[X_+ Y_+] &= \Pr(\v) \E[ X Y \mid \v] + Pr(\lnot\v) \, 0 \\&= \Pr(\v) \E\left[ (\mu_x + \sigma_x \T X) (\mu_y + \sigma_y \T Y) \mid \v \right] \\&= \mu_x \mu_y \Pr(\v) + \mu_y \sigma_x \Pr(\v) \E[ \T X \mid \v] + \mu_x \sigma_y \Pr(\v) \E[ \T Y \mid \v] \\&\qquad + \sigma_x \sigma_y \Pr(\v) \E\left[ \T X \T Y \mid \v \right] \\&= \mu_x \mu_y \Pr(\v) + \mu_y \sigma_x (q_x R_{xy} + \rho q_y R_{yx}) + \mu_x \sigma_y (\rho q_x R_{xy} + q_y R_{yx}) \\&\qquad + \sigma_x \sigma_y \left( \rho \Pr\left( \v \right) - \rho \mu_x q_x R_{xy} / \sigma_x - \rho \mu_y q_y R_{yx} / \sigma_y + r_{xy} \right) \\&= (\mu_x \mu_y + \sigma_x \sigma_y \rho) \Pr(\v) + (\mu_y \sigma_x + \mu_x \sigma_y \rho - \rho \mu_x \sigma_y) q_x R_{xy} \\&\qquad + (\mu_y \sigma_x \rho + \mu_x \sigma_y - \rho \mu_y \sigma_x) q_y R_{yx} + \sigma_x \sigma_y r_{xy} \\&= (\mu_x \mu_y + \Sigma_{xy}) \Pr(\v) + \mu_y \sigma_x q_x R_{xy} + \mu_x \sigma_y q_y R_{yx} + \sigma_x \sigma_y r_{xy} ,\end{align} and then subtracting $\E[X_+] \E[Y_+]$ we get \begin{align} \Cov(X_+, Y_+) &= (\mu_x \mu_y + \Sigma_{xy}) \Pr(\v) + \mu_y \sigma_x q_x R_{xy} + \mu_x \sigma_y q_y R_{yx} + \sigma_x \sigma_y r_{xy} \\&\qquad - (Q_x \mu_x + q_x \sigma_x) (Q_y \mu_y + q_y \sigma_y) .\end{align}


Here's some Python code to compute the moments:

import numpy as np
from scipy import stats

def relu_mvn_mean_cov(mu, Sigma):
    mu = np.asarray(mu, dtype=float)
    Sigma = np.asarray(Sigma, dtype=float)
    d, = mu.shape
    assert Sigma.shape == (d, d)

    x = (slice(None), np.newaxis)
    y = (np.newaxis, slice(None))

    sigma2s = np.diagonal(Sigma)
    sigmas = np.sqrt(sigma2s)
    rhos = Sigma / sigmas[x] / sigmas[y]

    prob = np.empty((d, d))  # prob[i, j] = Pr(X_i > 0, X_j > 0)
    zero = np.zeros(d)
    for i in range(d):
        prob[i, i] = np.nan
        for j in range(i + 1, d):
            # Pr(X > 0) = Pr(-X < 0); X ~ N(mu, S) => -X ~ N(-mu, S)
            s = [i, j]
            prob[i, j] = prob[j, i] = stats.multivariate_normal.cdf(
                zero[s], mean=-mu[s], cov=Sigma[np.ix_(s, s)])

    mu_sigs = mu / sigmas

    Q = stats.norm.cdf(mu_sigs)
    q = stats.norm.pdf(mu_sigs)
    mean = Q * mu + q * sigmas

    # rho_cs is sqrt(1 - rhos**2); but don't calculate diagonal, because
    # it'll just be zero and we're dividing by it (but not using result)
    # use inf instead of nan; stats.norm.cdf doesn't like nan inputs
    rho_cs = 1 - rhos**2
    np.fill_diagonal(rho_cs, np.inf)
    np.sqrt(rho_cs, out=rho_cs)

    R = stats.norm.cdf((mu_sigs[y] - rhos * mu_sigs[x]) / rho_cs)

    mu_sigs_sq = mu_sigs ** 2
    r_num = mu_sigs_sq[x] + mu_sigs_sq[y] - 2 * rhos * mu_sigs[x] * mu_sigs[y]
    np.fill_diagonal(r_num, 1)  # don't want slightly negative numerator here
    r = rho_cs / np.sqrt(2 * np.pi) * stats.norm.pdf(np.sqrt(r_num) / rho_cs)

    bit = mu[y] * sigmas[x] * q[x] * R
    cov = (
        (mu[x] * mu[y] + Sigma) * prob
        + bit + bit.T
        + sigmas[x] * sigmas[y] * r
        - mean[x] * mean[y])

    cov[range(d), range(d)] = (
        Q * (1 - Q) * mu**2 + (1 - 2 * Q) * q * mu * sigmas
        + (Q - q**2) * sigma2s)

    return mean, cov

and a Monte Carlo test that it works:

np.random.seed(12)
d = 4
mu = np.random.randn(d)
L = np.random.randn(d, d)
Sigma = L.T.dot(L)
dist = stats.multivariate_normal(mu, Sigma)

mn, cov = relu_mvn_mean_cov(mu, Sigma)

samps = dist.rvs(10**7)
mn_est = samps.mean(axis=0)
cov_est = np.cov(samps, rowvar=False)
print(np.max(np.abs(mn - mn_est)), np.max(np.abs(cov - cov_est)))

which gives 0.000572145310512 0.00298692620286, indicating that the claimed expectation and covariance match Monte Carlo estimates (based on $10,000,000$ samples).

$\endgroup$
2
  • $\begingroup$ can you summarize what those final values are? Are they estimates of the parameters mu and L you generated? Maybe print those target values? $\endgroup$
    – AdamO
    Feb 1 '18 at 21:34
  • $\begingroup$ No, the return values are $\E(Z_+)$ and $\Cov(Z_+)$; what I printed was the $L_\infty$ distance between Monte Carlo estimators of those quantities and the computed value. You could maybe invert these expressions to get a moment-matching estimator for $\mu$ and $\Sigma$ – Rosenbaum actually does that in his section 3 in the truncated case – but that's not what I wanted here. $\endgroup$
    – Danica
    Feb 1 '18 at 21:37

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.