bayesiains / nflows

Normalizing flows in PyTorch
MIT License
851 stars 119 forks source link

Forward vs. Inversse #32

Closed nwaftp23 closed 2 years ago

nwaftp23 commented 3 years ago

Hi,

From my understanding to sample from a normalizing flow you sample the base distribution and then push that forward through the transformations, but it seems you have decided to use the inverse. I would like the output of my flow to be bounded by [-1,1] and therefore apply a tanh transformation at the end of my transforms list. Though I am not getting the behavior I expect. I was hoping you could clarify what is meant by forward and inverse in the code.

def build_model(num_layers=2, hids=100, dims=2, context_dims=2,
        batch_norm=False, activation=torch.nn.functional.relu, bins = 15, tail=8.0,
        bounds = 9, device = 'cuda'):
    context_net = nn.Sequential(nn.Linear(context_dims, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 2*dims)
        )
    base_dist = nflows.distributions.ConditionalDiagonalNormal(
        shape=[dims], context_encoder= context_net)

    transforms = []

    def create_net(in_features, out_features):
        return nets.ResidualNet(
            in_features, out_features, context_features=context_dims,
            hidden_features=hids, num_blocks=2,
            use_batch_norm=batch_norm,
            activation=activation)

    for _ in range(num_layers):
        transforms.append(nflows.transforms.ReversePermutation(features=dims))
        if dims > 1:
            mask = nflows.utils.torchutils.create_mid_split_binary_mask(dims)
            transforms.append(
                nflows.transforms.PiecewiseRationalQuadraticCouplingTransform(
                    mask, create_net, tails='linear', num_bins=bins, tail_bound=tail,
                ))
        if dims == 1:
           transforms.append(
                nflows.transforms.MaskedPiecewiseRationalQuadraticAutoregressiveTransform(
                 features=dims,
                 hidden_features=hids,
                 context_features=context_dims,
                 tails='linear',
                 use_batch_norm=batch_norm,
                 num_bins=bins,
                 tail_bound = tail,
                 activation = activation))

    transforms.append(nflows.transforms.Tanh())
    transform = nflows.transforms.CompositeTransform(transforms)

    flow = nflows.flows.Flow(transform, base_dist)
    return flow

Best, Lucas

arturbekasov commented 2 years ago

From my understanding to sample from a normalizing flow you sample the base distribution and then push that forward through the transformations, but it seems you have decided to use the inverse.

What to define as a "forward" pass of a transformation is an arbitrary choice. We can define (as you suggest) that $x = f(z)$, but it makes no less sense to define $z = f(x)$. In fact, the latter definition (one that we use in nflows) is better aligned with the "normalizing" part of the "normalizing flow": the transformation $f$ normalizes the input distribution, i.e. transforms it into a standard normal.

As to your question, if you want the samples from your flow to lie in $[-1, 1]$, you simply need to use an inverse of Tanh (InverseTransform(Tanh())) as the first step of your flow. Alternatively, if it's easier for you to think in terms of $x = f(z)$, you can first define a transformation $f^{-1}$, and then wrap it in InverseTransform before passing to Flow. Note that for some transformations the forward and the inverse pass don't have the same computational cost, so you'd need to be a bit careful with my latter suggestion.

Hope this helps.