12
$\begingroup$

Batch normalization has been credited with substantial performance improvements in deep neural nets. Plenty of material on the internet shows how to implement it on an activation-by-activation basis. I've already implemented backprop using matrix algebra, and given that I'm working in high-level languages (while relying on Rcpp (and eventually GPU's) for dense matrix multiplication), ripping everything out and resorting to for-loops would probably slow my code substantially, in addition to being a huge pain.

The batch normalization function is $$ b(x_p) = \gamma \left(x_p - \mu_{x_p}\right) \sigma^{-1}_{x_p} + \beta $$ where

  • $x_p$ is the $p$th node, before it gets activated
  • $\gamma$ and $\beta$ are scalar parameters
  • $\mu_{x_p}$ and $\sigma_{x_p}$ are the mean and SD of $x_p$. (Note that the square root of of the variance plus a fudge factor is normally used -- let's assume nonzero elements for compactness)

In matrix form, batch normalization for a whole layer would be $$ b(\mathbf{X}) = \left(\gamma\otimes\mathbf{1}_p\right)\odot \left(\mathbf{X} - \mu_{\mathbf{X}}\right) \odot\sigma^{-1}_{\mathbf{X}} + \left(\beta\otimes\mathbf{1}_p\right) $$ where

  • $\mathbf{X}$ is $N\times p$
  • $\mathbf{1}_N$ is a column vector of ones
  • $\gamma$ and $\beta$ are now row $p$-vectors of the per-layer normalization parameters
  • $\mu_{\mathbf{X}}$ and $\sigma_{\mathbf{X}}$ are $N \times p$ matrices, where each column is a $N$-vector of columnwise means and standard deviations
  • $\otimes$ is the Kronecker product and $\odot$ is the elementwise (Hadamard) product

A very simple one-layer neural net with no batch normalization and a continuous outcome is $$ y = a\left(\mathbf{X\Gamma}_1\right)\Gamma_2 + \epsilon $$

where

  • $\Gamma_1$ is $p_1 \times p_2$
  • $\Gamma_2$ is $p_2 \times 1$
  • $a(.)$ is the activation function

If the loss is $R = N^{-1}\displaystyle\sum\left(y - \hat{y}\right)^2$, then the gradients are $$ \begin{array}{lr} \frac{\partial R}{\partial \Gamma_1} = -2\mathbf{V}^T \hat\epsilon\\ \frac{\partial R}{\partial \Gamma_2} = \mathbf{X}^T \left(a'(\mathbf{X}\mathbf{\Gamma}_1) \odot -2\hat\epsilon \mathbf{\Gamma}_2^T\right) \\ \end{array} $$

where

  • $\mathbf{V} = a\left(\mathbf{X}\Gamma_1\right)$
  • $\hat{\epsilon} = y-\hat{y}$

Under batch normalization, the net becomes $$ y = a\left(b\left(\mathbf{X}\Gamma_1\right)\right)\Gamma_2 $$ or $$ y = a\Big(\left(\gamma\otimes\mathbf{1}_N\right)\odot \left(\mathbf{X\Gamma_1} - \mu_{\mathbf{X\Gamma_1}}\right) \odot\sigma^{-1}_{\mathbf{X\Gamma_1}} + \left(\beta\otimes\mathbf{1}_N\right)\Big)\mathbf{\Gamma_2} $$ I have no idea how to compute the derivatives of Hadamard and Kronecker products. On the subject of Kronecker products, the literature gets fairly arcane.

Is there a practical way of computing $\partial R/\partial \gamma$, $\partial R/\partial \beta$, and $\partial R/\partial \mathbf{\Gamma_1}$ within the matrix framework? A simple expression, without resorting to node-by-node computation?

Update 1:

I've figured out $\partial R/\partial \beta$ -- sort of. It is: $$ \mathbf{1}_{N}^T \left(a'(\mathbf{X}\mathbf{\Gamma}_1) \odot -2\hat\epsilon \mathbf{\Gamma}_2^T\right) $$ Some R code demonstrates that this is equivalent to the looping way to do it. First set up the fake data:

set.seed(1)
library(dplyr)
library(foreach)

#numbers of obs, variables, and hidden layers
N <- 10
p1 <- 7
p2 <- 4
a <- function (v) {
  v[v < 0] <- 0
  v
}
ap <- function (v) {
  v[v < 0] <- 0
  v[v >= 0] <- 1
  v
}

# parameters
G1 <- matrix(rnorm(p1*p2), nrow = p1)
G2 <- rnorm(p2)
gamma <- 1:p2+1
beta <- (1:p2+1)*-1
# error
u <- rnorm(10)

# matrix batch norm function
b <- function(x, bet = beta, gam = gamma){
  xs <- scale(x)
  gk <- t(matrix(gam)) %x% matrix(rep(1, N))
  bk <- t(matrix(bet)) %x% matrix(rep(1, N))
  gk*xs+bk
}
# activation-wise batch norm function
bi <- function(x, i){
  xs <- scale(x)
  gk <- t(matrix(gamma[i]))
  bk <- t(matrix(beta[i]))
  suppressWarnings(gk*xs[,i]+bk)
}

X <- round(runif(N*p1, -5, 5)) %>% matrix(nrow = N)
# the neural net
y <- a(b(X %*% G1)) %*% G2 + u

Then compute derivatives:

# drdbeta -- the matrix way
drdb <- matrix(rep(1, N*1), nrow = 1) %*% (-2*u %*% t(G2) * ap(b(X%*%G1)))
drdb
           [,1]      [,2]    [,3]        [,4]
[1,] -0.4460901 0.3899186 1.26758 -0.09589582
# the looping way
foreach(i = 1:4, .combine = c) %do%{
  sum(-2*u*matrix(ap(bi(X[,i, drop = FALSE]%*%G1[i,], i)))*G2[i])
}
[1] -0.44609015  0.38991862  1.26758024 -0.09589582

They match. But I'm still confused, because I don't really know why this works. The MatCalc notes referenced by @Mark L. Stone say that the derivative of $\beta \otimes \mathbf{1}_N$ should be

$$ \frac{\partial A \otimes B}{\partial A} = \left(I_{nq} \otimes T_{mp}\right)\left(I_n\otimes vec(B) \otimes I_m\right) $$ where the subscripts $m$, $n$, and $p$, $q$ are the dimensions of $A$ and $B$. $T$ is the commutation matrix, which is just 1 here because both inputs are vectors. I try this and get a result that doesn't seem helpful:

# playing with the kroneker derivative rule
A <- t(matrix(beta)) 
B <- matrix(rep(1, N))
diag(rep(1, ncol(A) *ncol(B))) %*% diag(rep(1, ncol(A))) %x% (B) %x% diag(nrow(A))
     [,1] [,2] [,3] [,4]
 [1,]    1    0    0    0
 [2,]    1    0    0    0
 snip
[13,]    0    1    0    0
[14,]    0    1    0    0
snip
[28,]    0    0    1    0
[29,]    0    0    1    0
[snip
[39,]    0    0    0    1
[40,]    0    0    0    1

This isn't conformable. Clearly I'm not understanding those Kronecker derivative rules. Help with those would be great. I'm still totally stuck on the other derivatives, for $\gamma$ and $\mathbf{\Gamma_1}$ -- those are harder because they don't enter additively like $\beta \otimes \mathbf{1}$ does.

Update 2

Reading textbooks, I'm fairly sure that $\partial R/\partial \Gamma_1$ and $\partial R/\partial \gamma$ will require use of the vec() operator. But I'm apparently unable to follow the derivations sufficiently as to be able to translate them into code. For example, $\partial R/\partial \Gamma_1$ is going to involve taking the derivative of $w\odot\mathbf{X\Gamma_1}$ with respect to $\mathbf{\Gamma_1}$, where $w \equiv (\gamma \otimes \mathbf{1}) \odot \sigma_{\mathbf{X\Gamma_1}}^{-1}$ (which we can treat as a constant matrix for the moment).

My instinct is to simply say "the answer is $w\odot\mathbf{X}$", but that obviously doesn't work because $w$ isn't conformable with $\mathbf{X}$.

I know that $$ \partial(A \odot B) = \partial A \odot B + A \odot \partial B $$

and from this, that

$$ \frac{\partial vec(w \odot \mathbf{X\Gamma_1})}{\partial vec(\mathbf{\Gamma_1})^T} = vec(\mathbf{X\Gamma_1})I\frac{\partial vec(w)}{\partial vec(\mathbf{\Gamma_1})^T} + vec(w)I\frac{\partial vec(\mathbf{X\Gamma_1})}{\partial vec(\mathbf{\Gamma_1})^T} $$ But I'm uncertain how to evaluate this, let alone code it.

Update 3

Making progress here. I woke up at 2AM last night with this idea. Math is not good for sleep.

Here is $\partial R/\partial \mathbf{\Gamma_1}$, after some notational sugar:

  • $w \equiv (\gamma \otimes \mathbf{1}) \odot \sigma_{\mathbf{X\Gamma_1}}^{-1}$
  • $\text{"stub"} \equiv a'(b(\mathbf{X\Gamma}_1)) \odot -2\hat\epsilon \mathbf{\Gamma}_2^T$

Here's what you have after you get to the end of the chain rule: $$ \frac{\partial R}{\partial \Gamma_1} = \frac{\partial w \odot \mathbf{X\Gamma}_1}{\partial \Gamma_1}\left(\text{"stub"}\right) $$ Start by doing this the looping way -- $i$ and $j$ will subscript columns and $\mathbf{I}$ is a conformable identity matrix: $$ \frac{\partial R}{\partial \Gamma_{ij}} = \left(w_i \odot \mathbf{X_i}\right)^T\left(\text{"stub"}_j\right) $$ $$ \frac{\partial R}{\partial \Gamma_{ij}} = \left(\mathbf{I} w_i \mathbf{X_i}\right)^T\left(\text{"stub"}_j\right) $$ $$ \frac{\partial R}{\partial \Gamma_{ij}} = \mathbf{X_i}^T\mathbf{I} w_i\left(\text{"stub"}_j\right) $$ tl;dr you're basically pre-multiplying the stub by the batchnorm scale factors. This should be equivalent to: $$ \frac{\partial R}{\partial \Gamma} = \mathbf{X}^T\left(\text{"stub"}\odot w\right) $$

And, in fact it is:

stub <- (-2*u %*% t(G2) * ap(b(X%*%G1)))
w <- t(matrix(gamma)) %x% matrix(rep(1, N)) * (apply(X%*%G1, 2, sd) %>% t %x% matrix(rep(1, N)))
drdG1 <- t(X) %*% (stub*w)

loop_drdG1 <- drdG1*NA
for (i in 1:7){
  for (j in 1:4){
    loop_drdG1[i,j] <- t(X[,i]) %*% diag(w[,j]) %*% (stub[,j])
  }
}

> loop_drdG1
           [,1]       [,2]       [,3]       [,4]
[1,] -61.531877  122.66157  360.08132 -51.666215
[2,]   7.047767  -14.04947  -41.24316   5.917769
[3,] 124.157678 -247.50384 -726.56422 104.250961
[4,]  44.151682  -88.01478 -258.37333  37.072659
[5,]  22.478082  -44.80924 -131.54056  18.874078
[6,]  22.098857  -44.05327 -129.32135  18.555655
[7,]  79.617345 -158.71430 -465.91653  66.851965
> drdG1
           [,1]       [,2]       [,3]       [,4]
[1,] -61.531877  122.66157  360.08132 -51.666215
[2,]   7.047767  -14.04947  -41.24316   5.917769
[3,] 124.157678 -247.50384 -726.56422 104.250961
[4,]  44.151682  -88.01478 -258.37333  37.072659
[5,]  22.478082  -44.80924 -131.54056  18.874078
[6,]  22.098857  -44.05327 -129.32135  18.555655
[7,]  79.617345 -158.71430 -465.91653  66.851965

Update 4

Here, I think, is $\partial R / \partial \gamma$. First

  • $\widetilde{\mathbf{X\Gamma}} \equiv \left(\mathbf{X\Gamma} - \mu_{\mathbf{X\Gamma}}\right)\odot \sigma^{-1}_\mathbf{X\Gamma}$
  • $\tilde\gamma \equiv \gamma \otimes\mathbf{1}_N$

Similar to before, the chain rule gets you as far as $$ \frac{\partial R}{\partial \tilde\gamma} = \frac{\partial \tilde\gamma \odot \widetilde{\mathbf{X\Gamma}}}{\partial \tilde\gamma}\left(\text{"stub"}\right) $$ Looping gives you $$ \frac{\partial R}{\partial \tilde\gamma_i} = (\widetilde{\mathbf{X\Gamma}})_i^T \mathbf{I}\tilde\gamma_i \left(\text{"stub"}_i\right) $$ Which, like before, is basically pre-multiplying the stub. It should therefore be equivalent to: $$ \frac{\partial R}{\partial \tilde\gamma} = (\widetilde{\mathbf{X\Gamma}})^T \left(\text{"stub"} \odot \tilde\gamma \right) $$

It sort of matches:

drdg <- t(scale(X %*% G1)) %*% (stub * t(matrix(gamma)) %x% matrix(rep(1, N)))

loop_drdg <- foreach(i = 1:4, .combine = c) %do% {
  t(scale(X %*% G1)[,i]) %*% (stub[,i, drop = F] * gamma[i])  
}

> drdg
           [,1]      [,2]       [,3]       [,4]
[1,]  0.8580574 -1.125017  -4.876398  0.4611406
[2,] -4.5463304  5.960787  25.837103 -2.4433071
[3,]  2.0706860 -2.714919 -11.767849  1.1128364
[4,] -8.5641868 11.228681  48.670853 -4.6025996
> loop_drdg
[1]   0.8580574   5.9607870 -11.7678486  -4.6025996

The diagonal on the first is the same as the vector on the second. But really since the derivative is with respect to a matrix -- albeit one with a certain structure, the output should be a similar matrix with the same structure. Should I take the diagonal of the matrix approach and simply take it to be $\gamma$? I'm not sure.

It seems that I have answered my own question but I am unsure whether I am correct. At this point I will accept an answer that rigorously proves (or disproves) what I've sort of hacked together.

while(not_answered){
  print("Bueller?")
  Sys.sleep(1)
}
$\endgroup$
  • 2
    $\begingroup$ Chapter 9 section 14 of "Matrix Differential Calculus with Applications in Statistics and Econometrics" by Magnus and Neudecker, 3rd edition janmagnus.nl/misc/mdc2007-3rdedition covers differentials of Kronecker products and concludes with an exercise on differential of Hadamard product. "Notes on Matrix Calculus" by Paul L. Fackler www4.ncsu.edu/~pfackler/MatCalc.pdf has a lot of material on differentiating Kronceker products $\endgroup$ – Mark L. Stone Feb 20 '18 at 23:36
  • $\begingroup$ Thanks for the references. I've found those MatCalc notes before, but it doesn't cover Hadamard, and anyway I'm never certain whether a rule from non-matrix calculus applies or doesn't apply to to matrix case. Product rules, chain rules, etc. I'll look into the book. I'd accept an answer that points me to all of the ingredients I need to pencil it out myself... $\endgroup$ – generic_user Feb 21 '18 at 0:02
  • $\begingroup$ why are you doing this? why not use framewroks such as Keras/TensorFlow? It's a waste of productive time to implement these low level algorithms, that you could use on solving actual problems $\endgroup$ – Aksakal Feb 22 '18 at 16:17
  • 1
    $\begingroup$ More precisely, I'm fitting networks that exploit known parametric structure -- both in terms of linear-in-parameters representations of input data, as well as longitudinal/panel structure. Established frameworks are so heavily optimized as to be beyond my ability to hack/modify. Plus math is helpful generally. Plenty of codemonkeys have no idea what they're doing. Likewise learning enough Rcpp to implement it efficiently is useful. $\endgroup$ – generic_user Feb 22 '18 at 16:29
  • 1
    $\begingroup$ @MarkL.Stone not only is it theoretically sound, it's practically easy! A more or less mechanical process! &%#$! $\endgroup$ – generic_user Feb 23 '18 at 1:37
1
$\begingroup$

Not a complete answer, but to demonstrate what I suggested in my comment if $$b(X)=(X−e_N\mu_X^T)ΓΣ_X^{-1/2}+e_N\beta^T$$ where $\Gamma=\mathop{\mathrm{diag}}(\gamma)$, $\Sigma_X^{-1/2}=\mathop{\mathrm{diag}}(\sigma_{X_1}^{-1},\sigma_{X_2}^{-1},\dots)$ and $e_N$ is a vector of ones, then by the chain rule $$\nabla_\beta R=[-2\hat{\epsilon}(\Gamma_2^T\otimes I)J_X(a)(I\otimes e_N)]^T$$ Noting that $-2\hat{\epsilon}(\Gamma_2^T\otimes I)=\mathop{\mathrm{vec}}(-2\hat{\epsilon}\Gamma_2^T)^T$ and $J_X(a)=\mathop{\mathrm{diag}}(\mathop{\mathrm{vec}}(a^\prime(b(X\Gamma_1))))$, we see that $$\nabla_\beta R=(I\otimes e_N^T)\mathop{\mathrm{vec}}(a^\prime(b(X\Gamma_1))\odot-2\hat{\epsilon}\Gamma_2^T)=e_N^T(a^\prime(b(X\Gamma_1))\odot-2\hat{\epsilon}\Gamma_2^T)$$ via the identity $\mathop{\mathrm{vec}}(AXB)=(B^T\otimes A)\mathop{\mathrm{vec}}(X)$. Similarly, $$\begin{align}\nabla_\gamma R&=[-2\hat{\epsilon}(\Gamma_2^T\otimes I)J_X(a)(\Sigma_{X\Gamma_1}^{-1/2}\otimes (X\Gamma_1-e_N\mu_{X\Gamma_1}^T))K]^T\\&=K^T\mathop{\mathrm{vec}}((X\Gamma_1-e_N\mu_{X\Gamma_1}^T)^TW\Sigma^{-1/2}_{X\Gamma_1})\\&=\mathop{\mathrm{diag}}((X\Gamma_1-e_N\mu_{X\Gamma_1}^T)^TW\Sigma^{-1/2}_{X\Gamma_1})\end{align}$$ where $W=a^\prime(b(X\Gamma_1))\odot-2\hat{\epsilon}\Gamma_2^T$ (the "stub") and $K$ is an $Np\times p$ binary matrix that selects the columns of the Kronecker product corresponding to the diagonal elements of a square matrix. This follows from the fact that $d\Gamma_{i\neq j}=0$. Unlike the first gradient, this expression is not equivalent to the expression you derived. Considering that $b$ is a linear function w.r.t $\gamma_i$, there should not be a factor of $\gamma_i$ in the gradient. I leave the gradient of $\Gamma_1$ to the OP, but I will say for derivation with fixed $w$ creates the "explosion" the writers of the article seek to avoid. In practice, you will also need to find the Jacobians of $\Sigma_X$ and $\mu_X$ w.r.t $X$ and use product rule.

$\endgroup$

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.