Riemannian approach to batch normalization
Batch Normalization (BN) has proven to be an effective algorithm for deep neural network training by normalizing the input to each neuron and reducing the internal covariate shift. The space of weight vectors in the BN layer can be naturally interpreted as a Riemannian manifold, which is invariant to linear scaling of weights. Following the intrinsic geometry of this manifold provides a new learning rule that is more efficient and easier to analyze. We also propose intuitive and effective gradient clipping and regularization methods for the proposed algorithm by utilizing the geometry of the manifold. The resulting algorithm consistently outperforms the original BN on various types of network architectures and datasets.
Paper Summary
Starting from the observation that batch-normalization induces a particular
form of scale invariance on the weight matrix, the authors propose instead to
directly learn the weights on the unit-sphere. This is motivated from
information geometry as an example of optimization on a Riemannian manifold, in
particular the Stiefel manifold V(1,n) which contains unit-length vectors.  As
the descent direction on the unit sphere is well known (eq 7), the main
contribution of the paper is in extending popular optimization algorithms
(SGD+momentum and Adam) to constrained optimization on the unit-sphere.
Furthermore, the authors propose orthogonality as a (principled) replacement
for L2 regularization, which is no longer meaningful with norm constraints.
The method is shown to be effective across two families of models (VGG, wide
resnet) on CIFAR-10, CIFAR-100 and SVHN.
Review
I like this paper, as it proposes a principled alternative to batch
normalization, which is also computationally lightweight, compared to other
manifold optimization algorithms (e.g. natural gradient). That being said a few
flaws prevent me from giving a clear signal for acceptance.
(1) Clarity. I believe the paper would greatly benefit from increased
simplicity and clarity in the exposition of the idea. First and foremost,
concepts of element-wise scale invariance and optimization under unit-norm
constraints are much simpler to grasp (especially to people not familiar with
information geometry) than Riemannian manifolds, exponential maps and parallel
transport. I recommend the authors preface exposition of their methods with
such an intuitive explanation. Second, I also found there was a general lack of
clarity about what type of invariance is afforded by BN, and which direction
were the weights norms constrained in (column vs row). On line 75, the authors
write that "w^T x is invariant to linear scaling", when really this refers to
element-wise scaling of the columns of the matrix W \in mathbb{R}^{n x p},
where $p$ is the number of hidden units. In Algorithm 4, there is also a discrepancy
between the product "W x" (implies W is p x n) and the "column vectors w_i"
being unit-norm.
(2) Experimental results. As this is inherently an optimization algorithm, I
would expect there to be training curves. Does the algorithm simply converge
faster, does it yield a better minimum, or both ? Also, I would like to see
a controlled experiment showing the impact of the orthogonality regularizer:
what is the test error with and without ?
(3) Baselines and related work. This work ought to be discussed in the more
general context of Riemannian optimization for neural networks, e.g. natural
gradient descent algorithms (TONGA, K-FAC, HF, natural nets, etc.). Path-SGD also
seems closely related and should be featured as a baseline.
Finally, I am somewhat alarmed and puzzled that the authors chose to revert to
standard SGD for over-complete layers, and only use their method in the
under-complete case. Could the authors comment on how many parameters are in this
"under-complete" regime vs total number of parameters ? Are there no appropriate
relaxations of orthogonality ? I would have much rather seen the training curves
using their algorithm for *all* parameters (never mind test error for now) then
this strange hybrid method.