Implement natural gradient optimization for q_diag=True

[Joshuaalbert]

Description

Running an SVGP with NatGrad optimizer fails because meanvarsqrt_to_natural expects a full covariance shape.

MVCE

X = np.random.uniform(size=[N,1])
Y = X
m = gp.models.SVGP(X,Y,gp.kernels.RBF(1),gp.likelihoods.Gaussian(),q_diag=True)
gp.train.NatGradOptimizer(1.).minimize(m, vars_list=[[m.q_mu, m.q_sqrt]])

[hughsalimbeni]

The nat grads update isn't implemented for q_diag=True. I'm not sure how this work, but I'd like to think about it. In any case, the documentation should state that nat grads are only applicable for the full q_sqrt case

[hughsalimbeni]

I've not liked q_diag=True in the past as it has odd properties, but it does save on memory and also saves a big matmul

[Joshuaalbert]

Just noticed this hanging around issue of mine. q_diag=True can be very useful when your prior is close to the posterior, so that the Bayesian update adds little information. In this case, choosing the transformation which decorrelates your prior space y = L.x + m, with prior P(y) = N[m, LL^T], leads to a posterior P(x | d) = P(L.x + m|d) that is "nearly" decorrelated, i.e. P(x|d) \approx N[q_mean, diag(q_sqrt)^2] to very good approximation. This depends on your prior being similar enough to the posterior, which can be verifed e.g. with HMC to get actual samples of P(x|d). In this parametrisation the natural gradients are super simple, since the Fisher Information matrix for q_mean is just

1/diag(q_sqrt^2) and

for q_sqrt.unconstrained is

2 diag( (d q_sqrt / d q_sqrt.unconstrained)^2)/diag(q_sqrt^2)

[hughsalimbeni]

Yes you're right it would be simple to implement the nat grads directly via the fisher. The nat grad code uses autodiff, however, which is probably still fine but would require a new transform.

[st--]

@Joshuaalbert @hughsalimbeni I'm assuming this issue would apply the same to gpflow 2.0?