Expectation of the absolute value of the product of correlated jointly gaussians?

Question

I am reading the Performer paper https://arxiv.org/abs/2009.14794. To understand their ReLU kernel used to approximate softmax attention, I need to evaluate $\mathbb{E}[ReLU(x^T w) \cdot ReLU(y^T w)]$ where $w \sim \mathcal{N}(0, I)$. I was able to reduce the problem to the more general one of computing $\mathbb{E}[|X \cdot Y|]$ where $X,Y$ are correlated jointly Gaussians. However, now I am stuck. Is there an approach of computing this expression or bound it from above reasonably without numerical simulation? Cauchy-Schwarz seems to loose for the upper bound.

Answer 1

Probably not the best solution, but a direction that may be useful.

Suppose \begin{align} \begin{bmatrix} X \\ Y \end{bmatrix} \sim N_2\left(\begin{bmatrix} \mu_1 \\ \mu_2 \end{bmatrix}, \begin{bmatrix} \sigma_1^2 & \rho\sigma_1\sigma_2 \\ \rho\sigma_1\sigma_2 & \sigma_2^2 \end{bmatrix}\right). \end{align}

It is well known that the conditional distribution of $X$ given $Y = y$ is univariate normal: \begin{align} X | Y = y \sim N(\mu_1 + \rho\sigma_1\sigma_2^{-1}(y - \mu_2), (1 - \rho^2)\sigma_1^2). \tag{1} \end{align}

This suggests that, to bound $E[|XY|] = E[E[|X||Y||Y]] = E[|Y|E[|X||Y]]$, we can evaluate $E[|X||Y]$ first. In this answer (or use the property of folded normal distribution), it is shown that if $Z \sim N(\mu, \sigma^2)$, then \begin{align} E[|Z|] = \sigma\sqrt{\frac{2}{\pi}}e^{-\frac{\mu^2}{2\sigma^2}} + \mu(1 - 2\Phi(-\mu\sigma^{-1})), \tag{2} \end{align} where $\Phi$ is CDF of $N(0, 1)$. $(1)$ and $(2)$ together then imply \begin{align} & E[|X||Y] = \sqrt{1 - \rho^2}\sigma_1\sqrt{\frac{2}{\pi}}e^{-\frac{[\mu_1 + \rho\sigma_1\sigma_2^{-1}(Y - \mu_2)]^2}{2(1 - \rho^2)\sigma_1^2}} \\ +& (\mu_1 + \rho\sigma_1\sigma_2^{-1}(Y - \mu_2))\left(1 - 2\Phi\left(-\frac{\mu_1 + \rho\sigma_1\sigma_2^{-1}(Y - \mu_2)}{\sqrt{1 - \rho^2}\sigma_1}\right)\right). \tag{3} \end{align}

A trivial upper bound for $(3)$ is \begin{align} \sqrt{1 - \rho^2}\sigma_1\sqrt{\frac{2}{\pi}} + |\mu_1| + |\rho|\sigma_1\sigma_2^{-1}(|Y| + |\mu_2|), \end{align} which gives an upper bound for $E[|XY|]$ as \begin{align} & CE[|Y|] + |\rho|\sigma_1\sigma_2^{-1}E[Y^2] \\ =& C\left[\sigma_2\sqrt{\frac{2}{\pi}}e^{-\frac{\mu_2^2}{2\sigma_2^2}} + \mu_2(1 - 2\Phi(-\mu_2\sigma_2^{-1}))\right] + |\rho|\sigma_1\sigma_2^{-1}(\sigma_2^2 + \mu_2^2), \tag{4} \end{align} where $C = \sqrt{1 - \rho^2}\sigma_1\sqrt{\frac{2}{\pi}} + |\mu_1| + |\rho|\sigma_1\sigma_2^{-1}|\mu_2|$.

---

Updates

Added R code of comparing bound $(4)$ and the Cauchy-Schwarz bound $\sqrt{(\sigma_1^2 + \mu_1^2)(\sigma_2^2 + \mu_2^2)}$. As expected, bound $(4)$ is sharper when $X$ and $Y$ are weakly correlated while Cauchy-Schwarz bound is sharper when $X$ and $Y$ are highly correlated. Different setups of $\mu_i$ and $\sigma_i$ also result in different winners. Those who are interested can play with it.

bounds <- function(mu1, mu2, sigma1, sigma2, rho) {
  C1 <- sqrt(1 - rho^2) * sigma1 * sqrt(2/pi) + abs(mu1) + abs(rho * mu2) * sigma1/sigma2
  C2 <- sigma2 * sqrt(2/pi) * exp(-mu2^2/(2 * sigma2^2)) + 
    mu2 * (1 - 2 * pnorm(-mu2/sigma2))
  C3 <- abs(rho) * sigma1 * (sigma2^2 + mu2^2)/sigma2
  UB1 <- C1 * C2 + C3
  UB2 <- sqrt((sigma1^2 + mu1^2) * (sigma2^2 + mu2^2))
  c(UB1, UB2)
}
para_mat <- cbind(rep(0, 10), rep(0, 10), rep(1, 10), rep(1, 10), seq(0.1, 0.9, len = 10))
colnames(para_mat) <- c("mu1", "mu2", "sigma1", "sigma2", "rho")
My_bound <- apply(para_mat, 1, function(t) bounds(t[1], t[2], t[3], t[4], t[5])[1])
CS_bound <- apply(para_mat, 1, function(t) bounds(t[1], t[2], t[3], t[4], t[5])[2])
cbind(para_mat, My_bound, CS_bound)
Result:
    mu1 mu2 sigma1 sigma2       rho  My_bound CS_bound

[1,]   0   0      1      1 0.1000000 0.7334287        1
   [2,]   0   0      1      1 0.1888889 0.8140485        1
   [3,]   0   0      1      1 0.2777778 0.8893436        1
   [4,]   0   0      1      1 0.3666667 0.9589474        1
   [5,]   0   0      1      1 0.4555556 1.0222792        1
   [6,]   0   0      1      1 0.5444444 1.0784391        1
   [7,]   0   0      1      1 0.6333333 1.1260001        1
   [8,]   0   0      1      1 0.7222222 1.1625473        1
   [9,]   0   0      1      1 0.8111111 1.1834650        1
  [10,]   0   0      1      1 0.9000000 1.1774961        1