DenseAutoregresive.init_weights causes unstable learning for 1D target distributions

[feynmanliang]

When I run the example on flowtorch.ai with a 1D distribution:

import torch
import torch.distributions as dist
import flowtorch
import flowtorch.bijectors as bijectors
# Lazily instantiated flow plus base and target distributions
flow = bijectors.AffineAutoregressive(
    flowtorch.params.DenseAutoregressive()
)
base_dist = dist.Normal(torch.zeros(1), torch.ones(1))
target_dist = dist.Normal(torch.zeros(1)+5, torch.ones(1))
# Instantiate transformed distribution and parameters
new_dist, params = flow(base_dist)
# Training loop
opt = torch.optim.Adam(params.parameters(), lr=1e-3)
for idx in range(501):
    opt.zero_grad()
    # Minimize KL(p || q)
    y = target_dist.sample((1000,))
    loss = -new_dist.log_prob(y).mean()
    if idx % 100 == 0:
        print('epoch', idx, 'loss', loss)
        
    loss.backward()
    opt.step()

sns.relplot(
    data=pd.DataFrame(new_dist.sample((100,)).detach().numpy()),
    x=0, y=1
)

The loss goes to NaN unless the learning rate is set extremely low (1e-15 gives sensible results).

Removing the call to self._init_weights in DenseAutoregressiveresolves the issue and allows a more reasonable1e-3` learning rate.

[stefanwebb]

How about we disable my new initialization scheme and treat it more like a research topic? I'll add in a flag and it will be disabled by default

[feynmanliang]

Yeah, I think we should break off an experimental submodule