Understanding uncertainty in conditional distributions

[psteinb]

Dear all,

thanks for providing such a great library. If I want to use flows, I resort to zuko these days. Really excellent code and quality of implementation.

I currently am struggling to train an NSF flow on a particular problem. Here is a surrogate code to illustrate the data:

num_dims = 3
hoffset = .25
hscale = .8
hprior = torch.distributions.Uniform(low=hoffset+torch.zeros(num_dims), 
                   high=hoffset + (torch.ones(num_dims)*hscale)
                  )

htheta = hprior.sample((num_training_samples,)).to(device) #our condition later on
htheta *= (1 + .05*torch.randn_like(htheta)) #add more noise, can be removed if need be 

# (optional) set to 0 for an even simpler simulator
# we shift by quite some 
hard_shift = torch.arange(1,num_dims+1).to(device) 

hx = htheta + hard_shift + torch.randn_like(htheta)*torch.sqrt(hard_cov*torch.ones(3))

My goal is to learn $(hx | htheta)$, i.e. during the training loop, I'd have:

hflow = zuko.flows.NSF(features=3, context=3, transforms=6, hidden_features=[256]).to(device)

for step in range(5_000):
    #load batches of htheta -> htheta_, hx -> hx_
    hloss = -hflow(htheta_).log_prob(hx_)  # -log p(hx|htheta)
    hloss = hloss.mean()
    hlosses[hep] = hloss.item()

    hloss.backward()
    hoptimizer.step()
    hoptimizer.zero_grad()

However, the flow doesn't really learn too well. So I am wondering about the cause? Should I normalize the two datasets with min-max scaling or potentially z-scaling (mean, std) before training? Or should I normalize only x because it is clearly outside of $x \in [0,1]$?

What are people's thoughts and experiences?

Cheers,
P

[francois-rozet]

Hi Peter,

Thank you for the kind words!

Could you tell me what hard_cov is? I don't think it is defined in your example. If you could provide summary statistics (mean, std, median, min, max) for (each dimension of) x and theta, that would also be very helpful.

The last thing would be: have you tried with MAF instead of NSF? The former is less powerful but also more forgiving.

[psteinb]

Hi there, thanks for stepping in so swiftly.

So perhaps I'll take the chance and share from where I started:

num_dims = 3
hoffset = 0.
hscale = .8
hprior = torch.distributions.Uniform(low=hoffset+torch.zeros(num_dims), 
                   high=hoffset + (torch.ones(num_dims))
                  )

htheta = hscale*hprior.sample((num_training_samples,)).to(device) #our condition later on
#htheta *= (1 + .05*torch.randn_like(htheta)) #add more noise, can be removed if need be 

# (optional) set to 0 for an even simpler simulator
# we shift by quite some
hard_cov = 0.1 
hard_shift = torch.arange(1,num_dims+1).to(device) 

hx = htheta + hard_shift + torch.randn_like(htheta)*torch.sqrt(hard_cov*torch.ones(3))

I only added the extra code earlier to make the problem more Gaussian like. Here is what the code above generates in each dimension.

If I now look at some validation samples of $p(hx|h\vartheta)$, I get predictions like this:

(magenta line is the validation label, blue line is a histogram of 128 samples around an observed $\vartheta_o$).
As you can see, the support of the samples looks quite broad but at least the label is contained. What strikes me is the somewhat strong bias (see $x_2$) for such a simple problem.

So for this reason, I wanted to know, how users of zuko usually approach normalisation as this is easy to get wrong. On top, I am guessing my data is spilling into the unbounded region of the NSF?

I am surprised that such a simple problem is so tricky for the NSF. Maybe I am overseeing something obvious.

[francois-rozet]

Hi @psteinb, after looking at the data generation code, I can say a few things. First, the distribution of $x$ overlaps greatly for very different values of $\theta$.

This means that given a single reference $x_+$, many $\theta$ are plausible and we should expect $q(\theta | x_+)$ to be very wide (almost the prior distribution). In addition, it also means that even if samples $\bar{\theta} \sim q(\theta | x_+)$ are close to the ground-truth $\theta_+$, $p(x | \bar{\theta})$ will be wide too. The posterior predictive $q(x | x_+)$ combines these two sources of uncertainty, and is therefore even wider.

So the issue here is not with NSF, or MAF, or any model, but the data process itself that defines a very wide (yet exact) posterior $p(\theta | x)$. NSF is uncertain, not because it does not fit the data, but because the posterior is too.

To convince yourself, I would recommend to modify hard_cov to a low value such as 0.01 or 0.001 and you should observe that the inferred posteriors are now much more narrow.

[psteinb]

Thanks so much for these observations and the hint to look at the conditional densities. I continued studying this a bit. Here are some takeaways I made. Feel free to correct me, if I got this wrong:

• in the setting above, we want to learn a conditional density $p(x|\theta)$ using a flow (NSF, MAF, etc)
• if the mapping $\theta \rightarrow x$ is not bijective, this conditional density cannot be learned (i.e. multiple values of $\theta$ map to single value of $x$); this is what your plots illustrate
• normalisation only plays a minor role: we use "global" scalar values in affine transformations for normalisation (e.g. a set of scalars per dimension to shift and/or squeeze distributions); if this is done independently in $x$ and $\theta$ it changes the domain of each, but likely not the characteristics of the mapping itself

My simulator created a non-bijective mapping and hence the flow didn't learn. Once this was fixed (with and without normalisation), the flow was learning well.

Great discussion!

[francois-rozet]

we want to learn a conditional density $p ( x | \theta )$ using a flow (NSF, MAF, etc)

While you can "learn" the likelihood $p(x | \theta)$, we usually want to learn the posterior $p(\theta | x)$, because the likelihood $p(x | \theta)$ is already known (it is defined by your simulator).

if the mapping θ → x is not bijective, this conditional density cannot be learned

I would recommend to not think about simulators as functions/mappings, that can be bijective or not, but as distributions. In this view, uncertainty is not something you want to get rid off. It is a property of the simulator, and a core feature of probabilistic methods. The uncertainty of $p(\theta | x)$ is proportional to the one of $p(x | \theta)$. In fact

$$ p(\theta | x) \propto p(x | \theta) p(\theta) $$

Therefore, the posterior $p(\theta | x)$ can be learned for any likelihood $p(x | \theta)$, deterministic or not.

normalisation only plays a minor role

It can play a bigger role, but in this case most of the uncertainty is inherent to the choice of $p(x | \theta)$, i.e. the simulator.

[psteinb]

On a) yes I know, but learning the posterior was not my goal in the original project where this problem arose.
On b) I fully agree, but it appeared as a helpful metaphor here to help me mentally grasp the problem for conditional density estimation as such.
On c) fair point.