I'd like to show you my rough sketch for scaling up natural gradients to deep neural networks that appears to be easy to automate just like automatic differentiation. I think there must be a flaw somewhere, since natural gradients would otherwise be more commonplace, and I'd like you to point it out.
Background
Neural networks often take the role of parameterizing the distributional parameters of a conditional distribution. Namely, we model
$$p_\theta(y|x) = d(y|\eta_\theta(x)),$$
where $d$ is a simple parametric distribution, whose parameters $\eta_\theta(x)$ (like the mean of a Gaussian distribution, or the probability vector of a categorical distribution) are computed by a neural network that takes $x$ as input and $\theta$ as neural network weights. By using the letter $\eta$, I suggest that we are dealing with natural parameters of an exponential family distribution, but I am more generally thinking of parameters for which natural gradients are easy to compute.
Natural gradients $n_\eta$ are computed as the usual (collumn vector, denominator layout) gradient $g_\eta := \nabla_\eta L(\eta)$ (for some loss function $L(\eta)$) multiplied with the inverted Fisher information matrix $F_\eta^{-1}$:
$$n_\eta := F_\eta^{-1} g_\eta,$$
and analogously for $\theta$. For the low-dimensional and suitably-chosen $\eta$, the computation of $F_\eta^{-1}$, and hence $n_\eta$, works. For the high-dimensional and computationally-entangled $\theta$ – the weights of our deep neural network – it apparently doesn't work well enough to be widely used, but I don't understand why.
Computing the Fisher information matrix for $\theta$
First, denote $G_{\theta,\eta} := \nabla_\theta \eta_\theta(x)$ for the $\dim(\theta)\times\dim(\eta)$ gradient matrix. Note that $g_\theta = G_{\theta,\eta} g_\eta$. Note further that, due to $\dim(\theta) > \dim(\eta)$, $G_{\theta,\eta}^\intercal G_{\theta,\eta}$ is a smaller square matrix than $G_{\theta,\eta}G_{\theta,\eta}^\intercal$.
We have
\begin{align*} F_\theta &= \mathbb{E}_{d(y|\eta_\theta(x))} \nabla_\theta \log d(y|\eta_\theta(x)) (\nabla_\theta \log d(y|\eta_\theta(x)))^\intercal \\ &= \mathbb{E}_{d(y|\eta_\theta(x))} G_{\theta,\eta} \nabla_\eta \log d(y|\eta_\theta(x)) (\nabla_\eta \log d(y|\eta_\theta(x)))^\intercal G_{\theta,\eta}^\intercal \\ &= G_{\theta,\eta} \left(\mathbb{E}_{d(y|\eta_\theta(x))} \nabla_\eta \log d(y|\eta_\theta(x)) (\nabla_\eta \log d(y|\eta_\theta(x)))^\intercal\right) G_{\theta,\eta}^\intercal \\ &= G_{\theta,\eta} F_\eta G_{\theta,\eta}^\intercal. \end{align*}
This result is stated on Wikipedia as well. It provides a straightforward way to compute $F_\theta$ given that we know $F_\eta$. Clearly though, $F_\theta$ is not full-rank, since $\dim(\theta) > \dim(\eta)$. Note that, from here on, I assume that all involved matrices are of their maximal possible rank, so that $F_\theta$ is of rank $\dim(\eta)$.
Justification for natural gradients based on Moore-Penrose inverse
I am not an expert on the topic, and a bit clueless about what it means for the Fisher information to be non-invertible (other than that the model is non-identifiable, which I think is a given in deep learning).
Since we need a substitute for the inverse of the Fisher information matrix, I propose to use the Moore-Penrose inverse, which works well in other contexts. What follows is a justification for that approach.
The usual natural gradient update $\theta_\mathrm{new} = \theta + \delta = \theta + \alpha F_\theta^{-1} g_\theta$, with learning rate $\alpha$, is motivated by the fact that $\delta = \alpha F_\theta^{-1} g_\theta$ minimizes $L(\theta) + \delta^\intercal g_\theta + \frac1\alpha\frac12 \delta^\intercal F_\theta \delta$ where $L(\theta) + \delta^\intercal g_\theta$ is an approximation of $L(\theta+\delta)$ and $\delta^\intercal F_\theta \delta$ is an approximation of $\operatorname{KL}(d(y|\eta_\theta(x))\Vert d(y|\eta_{\theta+\delta}(x)))$.
In our case, $F_\theta^{-1}$ doesn't exist and we want to use $\delta = \alpha F_\theta^+ g_\theta$ instead, where $F_\theta^+$ denotes the Moore-Penrose inverse of $F_\theta$. We want to check whether this is still minimizing $L(\theta) + \delta^\intercal g_\theta + \frac1\alpha\frac12 \delta^\intercal F_\theta \delta$. The first order condition is
$$0 = g_\theta + \frac1\alpha F_\theta \delta.$$
This is a linear system of equations and for singular $F_\theta$, it doesn't necessarily have a solution. Yet the Moore-Penrose pseudoinverse is known to provide a least squares solution.
The usual second order condition for a minimum is that $F_\theta$ is positive semidefinite, which is the case for Fisher information matrices. Without further research I suppose that this is also relevant in the least-squares case. Only if $F_\theta$ was also positive definite, there could be a unique solution. This is the case if and only if it has full rank. We are dealing with exactly the situation where this is not the case. (I suppose this is strongly related to the non-identifiability of $\theta$.) Yet $\delta = \alpha F_\theta^+ g_\theta$ is still not an arbitrary solution since it minimizes the Euclidian norm (which may be irrelevant for natural gradients, though).
Calculating natural gradients with the Moore-Penrose inverse of the Fisher information matrix
I proceed with the calculation of $F_\theta^+$. It relies mostly on properties of the Moore-Penrose inverse which hold for the usual inverse as well. To compute $G_{\theta,\eta}^+$, I use this nice property that holds due to the previously assumed maximal rank of $G_{\theta,\eta}$.
We have
\begin{align*} F_\theta^+ &= \left(G_{\theta,\eta} F_\eta G_{\theta,\eta}^\intercal\right)^+ \\&= \left(G_{\theta,\eta}^+\right)^\intercal F_\eta^+ G_{\theta,\eta}^+ \\&= G_{\theta,\eta} \left(G_{\theta,\eta}^\intercal G_{\theta,\eta}\right)^{-1} F_\eta^{-1}\left(G_{\theta,\eta}^\intercal G_{\theta,\eta}\right)^{-1}G_{\theta,\eta}^\intercal. \end{align*}
Having obtained this result, a "natural gradient" based on the Moore-Penrose inverse can be defined and computed: \begin{align*} n_\theta :&= F_\theta^+ g_\theta \\&= F_\theta^+ G_{\theta,\eta} g_\eta \\&= G_{\theta,\eta} \left(G_{\theta,\eta}^\intercal G_{\theta,\eta}\right)^{-1} F_\eta^{-1}\left(G_{\theta,\eta}^\intercal G_{\theta,\eta}\right)^{-1}G_{\theta,\eta}^\intercal G_{\theta,\eta} g_\eta \\&= G_{\theta,\eta} \left(G_{\theta,\eta}^\intercal G_{\theta,\eta}\right)^{-1} F_\eta^{-1} g_\eta \\&= G_{\theta,\eta} \left(G_{\theta,\eta}^\intercal G_{\theta,\eta}\right)^{-1} n_\eta \end{align*}
This constitutes a chain rule equivalent for natural gradients. For low $\dim(\eta)$, the computation of $(G_{\theta,\eta}^\intercal G_{\theta,\eta})^{-1}$ is no computational concern. Except for this factor, and the factor $F_\eta^{-1}$, the computation is as in usual gradient computations.
Naive implementation (not important to read)
In a practical implementation, one could first compute $G_{\theta,\eta}$ by backpropagation (cumbersome in PyTorch that usually only backpropagates from scalar tensors). With that, $n_\theta$ could be computed and assigned for the whole of $\theta$. This goes a bit against the modular and sequential style of auto-differentiation, that usually assigns gradients whenever leaf-variables of the computation graph are reached. Alternatively, it's possible to only compute just $(G_{\theta,\eta}^\intercal G_{\theta,\eta})^{-1} F_\eta^{-1}$ and multiply it in once the usual gradient descend reaches $g_\eta$. That lets the automatic differentiation system take it's usual way but it requires an additional backward pass.
Modularization
The chain rule we found for $\theta$ and $\eta$ applies between any two variables occurring in a neural network, e.g. between the input&weights and the output both of the same single layer. This allows for the usual modularization of the backpropagation process.
Indeed, if this all would work out, each granular backward step in e.g. Pytorch could be transformed to use the "natural" chain rule instead of the usual one. (A gradient matrix of maximal rank would have to be ensured. If the matrix is of a Small$\times$Large shape instead of the Large$\times$Small shape as before, we have to use the property for linearly independent rows instead, which impedes the cancellation to the right of the Fisher-information-matrix in the above calculation. This is a serious problem for modularization but hopefully rare.) Of course a natural-gradient-mode would be a noticeable (but not prohibitive) more compute intensive and rely on the programmer to make sure that the innermost Fisher information $F_\eta^{-1}$ is computed and multiplied at the right place.
To give an example of the modular "natural gradient" computation: For a scalar multiplication $c=ab$, usually $\pmatrix{g_a\\g_b} = \pmatrix{b\\a} g_c = \pmatrix{bg_c\\ag_c}$ is computed in the backward pass. With the "natural" chain rule, $\pmatrix{n_a\\n_b} = \pmatrix{b\\a}\left(\pmatrix{b\\a}^\intercal\pmatrix{b\\a}\right)^{-1}n_c = \pmatrix{\frac{bn_c}{a^2+b^2}\\\frac{an_c}{a^2+b^2}}$ would be computed instead.
Maybe you understand why the prospect of an automatic "natural" differentiation excites me. (I've recently encountered a training issue that only natural gradients might be able to solve.). So where is my mistake?
