Normalizing flow with a truncated support

[Bill-Gots]

Running the following lines

import zuko
model = zuko.flows.GF(3)
print(model.base().cdf(0.))

I got a NotImplementedError in torch\distributions\distribution.py.

Plus, I wanted to truncate the base in flows.GF, but I failed to manage it.

import zuko
model = zuko.flows.GF(3)
model.base = zuko.distributions.Truncated(model.base(), -1., 1.)
samples = model().rsample((32,))
print(samples)

Running the code above I got

File "D:\Program Files\Python3.11.4\Lib\site-packages\zuko\distributions.py", line 458, in __init__
    assert not base.event_shape, "'base' has to be univariate"
AssertionError: 'base' has to be univariate

I am new in normalization flows so I don't know how to correct this. Any help is appreciated.

[francois-rozet]

Hello, DiagNormal is a multivariate normal distribution with a diagonal covariance matrix. A multivariate distribution does not have a CDF, or rather there is no unique definition of the CDF of a multivariate distribution. What do you want to achieve exactly?

I guess that you want to define a multivariate truncated normal. zuko.distributions.Truncated works for univariate distributions, as the error message suggests. In your case, you should truncate univariate normal distributions, and then stack them as a single multivariate distribution. Note that the base attribute should be a lazy distribution.

class LazyDiagTruncatedNormal(zuko.flows.LazyDistribution):
    def __init__(self, features: int, lower: float = -1.0, upper: float = 1.0):
        super().__init__()

        self.register_buffer('loc', torch.zeros(features))
        self.register_buffer('scale', torch.ones(features))
        self.register_buffer('lower', torch.as_tensor(lower))
        self.register_buffer('upper', torch.as_tensor(upper))

    def forward(self, c: Tensor = None):
        return Independent(Truncated(Normal(self.loc, self.scale), self.lower, self.upper), 1)

model.base = LazyDiagTruncatedNormal(3)

I don't know why you want a truncated base distribution, but I should warn you that evaluating the log_prob of such flow will likely lead to NaNs / $-\infty$ as the probability outside of the interval is null ($\log 0$).

[Bill-Gots]

Hello, DiagNormal is a multivariate normal distribution with a diagonal covariance matrix. A multivariate distribution does not have a CDF, or rather there is no unique definition of the CDF of a multivariate distribution. What do you want to achieve exactly?

I guess that you want to define a multivariate truncated normal. zuko.distributions.Truncated works for univariate distributions, as the error message suggests. In your case, you should truncate univariate normal distributions, and then stack them as a single multivariate distribution. Note that the base attribute should be a lazy distribution.

class LazyDiagTruncatedNormal(zuko.flows.LazyDistribution):
    def __init__(self, features: int, lower: float = -1.0, upper: float = 1.0):
        super().__init__()

        self.register_buffer('loc', torch.zeros(features))
        self.register_buffer('scale', torch.ones(features))
        self.register_buffer('lower', torch.as_tensor(lower))
        self.register_buffer('upper', torch.as_tensor(upper))

    def forward(self, c: Tensor = None):
        return Independent(Truncated(Normal(self.loc, self.scale), self.lower, self.upper), 1)

model.base = LazyDiagTruncatedNormal(3)

I don't know why you want a truncated base distribution, but I should warn you that evaluating the log_prob of such flow will likely lead to NaNs / −∞ as the probability outside of the interval is null (log⁡0).

Thanks for your immediate reply! After reading the documentation, I'm still not very sure about the usage of many classes because I'm just learning normalization flows. What I want to achieve exactly is a snippet of initializing an NSF (or some other flows in zuko.flows) and then making it "truncated", so I want to add a "truncate transform" in that NSF. I tried to add SoftclipTransform in model.transforms first but I failed because the transform is not a nn.Module. Then I tried Truncated like above and as you could see I failed again.

I'm learning to use zuko but lack some fundamental knowledge, so I'm struggling to understand the usages. Like, why there are "LazyTransforms" (are there some transforms called "hard-working transforms"? just a joke)? So please forgive me for any misunderstanding caused. :)

[francois-rozet]

Like, why there are "LazyTransforms" (are there some transforms called "hard-working transforms"? just a joke)?

Have you checked the Learn the basics tutorial and in particular the parametrization section? It is explained there why LazyDistribution and LazyTransform are necessary, and what they do.

When designing the distributions module, the PyTorch team decided that distributions and transformations should be lightweight objects that are used as part of computations but destroyed afterwards. Consequently, the Distribution and Transform classes are not sub-classes of torch.nn.Module, which means that we cannot retrieve their parameters with .parameters(), send their internal tensor to GPU with .to('cuda') or train them as regular neural networks. In addition, the concepts of conditional distribution and transformation, which are essential for probabilistic inference, are impossible to express with the current interface.

To solve these problems, zuko defines two concepts: the LazyDistribution and the LazyTransform, which are modules whose forward pass returns a distribution or transformation, respectively. These components hold the parameters of the distributions/transformations as well as the recipe to build them, such that the actual distribution/transformation objects are lazily built and destroyed when necessary. Importantly, because the creation of the distribution/transformation object is delayed, an eventual condition can be easily taken into account. This design enables lazy distributions to act like distributions while retaining features inherent to modules, such as trainable parameters.

What I want to achieve exactly is a snippet of initializing an NSF [...] and then making it "truncated"

What do you mean by a "truncated flow"? The support of the distribution defined by the flow should be limited to some interval (hyper-cube)?

[Bill-Gots]

Yes I checked the "Learn the basic" part but I'm not quite sure. I mean LazyDistribution and LazyTransform have inherited nn.Module so we can optimize the parameters by using model.parameters(), which achieves the target mentioned in the "Learn the basic". That's why I can't get what 'lazy' refers to (maybe it means 'creation is delayed'?).

And a "truncated flow", which is what I want to achieve, means the samples of the flow are limited in some intervals given by user. Suppose I have initialized an NSF with features = 3. Now I want to limit the samples of that flow in some intervals like: (-1, 1) for the first and the second features, (-2, 2) for the third feature. Since I think I don't have the need to re-construct the whole architecture, all I need is to add a "truncate transform" to the flow in order to limit the range of the samples, and I just don't know how to add it. Like what can be done in neural networks, if I want to add some layers to the fully connected layers of the networks, I will code like

model.fc = nn.Sequential([model.fc, nn.ReLU(), nn.Linear(...)])

instead of

class New_Network(nn.Module):

    def __init__(self, ...):
        ...
        ...
        ...
        ...
        self.fc = nn.Sequential([nn.Linear(...), nn.ReLU(), nn.Linear(...)])

    def forward(self, x):
        ...
        ...
        ...
        ...
        return self.fc(x)

[francois-rozet]

That's why I can't get what 'lazy' refers to (maybe it means 'creation is delayed'?).

Yes exactly, lazy refers to the fact that the actual distribution/transformation is lazily built when needed.

Now I want to limit the samples of that flow in some intervals.

Ok! This is partially answered in #10, but the answer is a bit outdated. You are right, you can do that by adding a transformation to the flow. However, you should think about the kind of transformation to add and where to add it.

So first, what kind of transformation? Normalizing flows like MAF or NSF work in $\mathbb{R}$. So if you want features that are in $[a, b]$, you need a transformation that maps $\mathbb{R} \mapsto [a, b]$. For example,

$$ f(x) = \mathrm{sigmoid}(x) (b - a) + a $$

where $\mathrm{sigmoid}(x) = \frac{1}{1 + \exp(-x)}$ is the sigmoid function or simply

$$ f(x) = \frac{x + 5}{10} (b - a) + a $$

which is more appropriate for NSF, but might lead to a bit of leaks outside of $[a, b]$. In code, the latter would be

class LazyTruncateTransform(LazyTransform):
    def __init__(self, a: Tensor, b: Tensor):
        super().__init__()

        self.register_buffer('a', a)
        self.register_buffer('b', b)

    def forward(self, c: Tensor = None):
        return AffineTransform((self.b + self.a) / 2, (self.b - self.a) / 10)

Now, where to add this transformation? In a Flow object (result of MAF, NSF, ...), the attribute transform represents the sequence of transformations $f_0, f_1 \circ \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 first transformation $f_0$ should be the inverse of a LazyTruncateTransform instance.

a = torch.tensor([-1.0, -1.0, -2.0])
b = torch.tensor([1.0, 1.0, 2.0])

flow = zuko.flows.NSF(3)
flow.transform.transforms.insert(0, LazyTruncateTransform(a, b).inv)

I think this should work well enough, but as said earlier you might get a bit of leaks outside $[a, b]$ (about one in a million). For another flow (MAF, NAF, ...) the first transformation (with the sigmoid) could be more appropriate, but you would need to ensure that features are strictly between $a$ and $b$ to prevent NaNs.

class LazyTruncateTransform(LazyTransform):
    def __init__(self, a: Tensor, b: Tensor):
        super().__init__()

        self.register_buffer('a', a)
        self.register_buffer('b', b)

    def forward(self, c: Tensor = None):
        return ComposedTransform(
            SigmoidTransform(),
            AffineTransform(self.a, self.b - self.a),
        )

[Bill-Gots]

I'm really appreciated for your informative guide! I'll try that!