[FR] add an alternative parameterization for lower triangular cholesky factors

[martinjankowiak]

add an alternative parameterization for lower triangular cholesky factors and consider using this in AutoMultivariateNormal:

instead of parameterizing a lower cholesky factor as an unconstrained strictly lower triangular piece and a positive diagonal we instead parameterize as

L = unit_scale_tril @ scale_diag

where unit_scale_tril is lower triangular with ones along the diagonal and scale_diag is a positive diagonal matrix.

for more details see the corresponding NumPyro PR

[fritzo]

I agree we should improve the parametrization of AutoNormal. We did something similar to AutoLowRankMultivariateNormal #2127, and indeed AutoLowRankMultivariateNormal now works much better than AutoMultivariateNormal.

[fritzo]

Note I've been using a similar overparametrization in the AutoStructured guide, since the raw MVN via scale_tril performed so poorly:

pyro/pyro/infer/autoguide/guides.py

Lines 1564 to 1573 in c8dc40a

	# This overparametrizes by learning (scale,scale_tril),
	# enabling faster learning of the more-global scale parameter.
	aux_value = pyro.sample(
	name + "_aux",
	dist.Normal(zero, 1).to_event(1),
	infer={"is_auxiliary": True},
	)
	scale = _deep_getattr(self.scales, name)
	scale_tril = _deep_getattr(self.scale_trils, name)
	aux_value = aux_value @ scale_tril.T * scale

[martinjankowiak]

interesting. note this isn't an overparametrization though: the diagonal on unit_scale_tril is all ones