Understanding SoftclipTransform

[arnauqb]

Description

I'm trying to limit the output domain of a normalizing flow. I noticed that there is a zuko.transforms.SoftclipTransform transformation that could be useful for my case.

Reproduce

This is what I have tried:

class MyFlow(zuko.flows.FlowModule):
    def __init__(
        self,
        features: int,
        context: int = 0,
        transforms: int = 3,
        randperm: bool = False,
        **kwargs,
    ):
        orders = [
            torch.arange(features),
            torch.flipud(torch.arange(features)),
        ]

        transforms = [
            zuko.flows.MaskedAutoregressiveTransform(
                features=features,
                context=context,
                order=torch.randperm(features) if randperm else orders[i % 2],
                **kwargs,
            )
            for i in range(transforms)
        ]

        base = zuko.flows.Unconditional(
            zuko.distributions.DiagNormal,
            torch.zeros(features),
            torch.ones(features),
            buffer=True,
        )
        
        transforms.append(zuko.flows.Unconditional(zuko.transforms.SoftclipTransform))

        super().__init__(transforms, base)
        

Expected behavior

I would expect the samples to be contrained in the [-5, 5] domain, since that is the default argument of SoftClipTransform.

I have two questions:

1. Why is the domain not limited to [-5, 5]?

flow = MyFlow(1, 1, transforms=5, hidden_features=[50])
samples = flow(torch.rand(1)).sample((10000,))
plt.hist(samples.numpy(), bins=100);

2. How can I change the bound argument of the transformation? I noticed that the transformation class is passed to the flow, instead of its instance, so I'm not sure how one can change the initialization arguments.

[francois-rozet]

Hello @arnauqb 👋

transforms is the sequence of transformations $f_0, f_1, \dots f_n$ from the feature space towards the latent space. In order words, $z = f_n \circ \dots \circ f_1 \circ f_0(x)$ . When sampling, the order of the transformations is reversed, meaning that $x = f_0^{-1} \circ f_1^{-1} \circ \dots \circ f_n^{-1}(z)$.

Therefore, if you want your samples to be bounded, the inverse of the first transformation $f_0^{-1}$ should be the SoftclipTransform, but, in your snippet, you made the last transformation $f_n$ the SoftclipTransform.

The easiest solution is to insert the inverse SoftclipTransform at the front of the transforms list.

class MyFlow(zuko.flows.MAF):
    def __init__(self, *args, bound: float = 5.0, **kwargs):
        super().__init__(*args, **kwargs)

        self.transforms.insert(0, zuko.flows.Unconditional(self.inverse_softclip, bound=bound))

    def inverse_softclip(self, bound: float = 5.0):
        return zuko.transforms.SoftclipTransform(bound).inv

However, I must warn you that this is dangerous. The softclip function maps $\mathbb{R}$ to the interval $[-B, B]$. Therefore, the inverse softclip maps $[-B, B]$ to $\mathbb{R}$, meaning that values close to $B$ are mapped to large values. For instance,

>>> softclip = zuko.transforms.SoftclipTransform(5.0)
>>> softclip.inv(4.9)
... 245.00000000000117
>>> softclip.inv(4.99)
... 2494.9999999999977

This would lead to very high input variances during training, for which neural networks often break. To mitigate this issue, if your inputs lie in $[-B, B]$, you can set a slightly larger bound (e.g. $B + 1$ or $2B$) to SoftclipTransform.

[felixdivo]

This is a very useful recipe for adding custom transforms (e.g. simple data transforms, which may have nonzero log abs det jacobian and thus might need to be accounted for). Could it be useful to provide a more direct API for composing transforms? Much like we currently have torch.distributions.ComposeTransform, just for the transforms in flows.

[francois-rozet]

That is the beauty of Zuko: you can already use torch.distributions.ComposeTransform! For example, if the logarithm of your input $\log x$ is distributed more or less uniformly between $a$ and $b$, you could do

from torch.distributions.transforms import *

preprocess = ComposeTransform([
    ExpTransform().inv,  # maps R+ to R
    AffineTransform((a + b) / (a - b), 2 / (b - a)),  # maps [a, b] to [-1, 1]
])

class MyFlow(zuko.flows.MAF):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.transforms.insert(0, zuko.flows.Unconditional(lambda: preprocess))