Repository: probabilists/zuko

Issue:
Add Extrapolation to Bernstein Polynomial Flow
### Description

As discussed in #32, we now implemented linear extrapolation outside the bounds of the Bernstein Polynomial.
The feature becomes active if `linear=True`.

Here is a simple Python script to visualize the resulting effect:

```python
# %% Imports
import torch

from matplotlib import pyplot as plt
from zuko.transforms import BernsteinTransform

# %% Globals
M = 10
batch_size = 10
torch.manual_seed(1)
theta = torch.rand(size=(M,)) * 500  # creates a random parameter vector

# %% Sigmoid
bpoly = BernsteinTransform(theta=theta, linear=False)

x = torch.linspace(-15, 15, 2000)
y = bpoly(x)

adj = bpoly.log_abs_det_jacobian(x, y).detach()
J = torch.diag(torch.autograd.functional.jacobian(bpoly, x)).abs().log()

# %% Plot
fig, axs = plt.subplots(2, sharex=True)
fig.suptitle("Bernstein polynomial with Sigmoid")
axs[0].plot(x, y, label="Bernstein polynomial")
axs[0].scatter(
    torch.linspace(-10, 10, bpoly.order + 1),
    bpoly.theta.numpy().flatten(),
    label="Bernstein coefficients",
)
axs[0].legend()
axs[1].plot(x, adj, label="ladj")
# axs[1].scatter(
#     torch.linspace(-10, 10, bpoly.order),
#     bpoly.dtheta.numpy().flatten(),
#     label="dtheta",
# )
axs[1].plot(x, J, label="ladj (autograd)")
axs[1].legend()
fig.tight_layout()
fig.savefig("sigmoid.png")
```

![sigmoid](https://github.com/probabilists/zuko/assets/38564291/8c3817f8-79a6-441b-8e56-12ead6b17b6d)

```python
# %% Extrapolataion
bpoly = BernsteinTransform(theta=theta, linear=True)

x = torch.linspace(-15, 15, 2000)
y = bpoly(x)

adj = bpoly.log_abs_det_jacobian(x, y).detach()
J = torch.diag(torch.autograd.functional.jacobian(bpoly, x)).abs().log()

# %% Plot

fig, axs = plt.subplots(2, sharex=True)
fig.suptitle("Bernstein polynomial with linear extrapolation")
axs[0].plot(x, y, label="Bernstein polynomial")
axs[0].scatter(
    torch.linspace(-10, 10, bpoly.order + 1),
    bpoly.theta.numpy().flatten(),
    label="Bernstein coefficients",
)
axs[0].legend()
axs[1].plot(x, adj, label="ladj")
# axs[1].scatter(
#     torch.linspace(-10, 10, bpoly.order),
#     bpoly.dtheta.numpy().flatten(),
#     label="dtheta",
# )
axs[1].plot(x, J, label="ladj (autograd)")
axs[1].legend()
fig.tight_layout()
fig.savefig("linear.png")
```

![linear](https://github.com/probabilists/zuko/assets/38564291/fe55a813-8d70-4b23-8431-0b3685b1f7f4)

This makes the `BPF` Implementation more robust to data laying outside the domain.of the Bernstein Polynomial,without the need for thenon linear sigmoid function.

@oduerr Do you have anything else to add?

### Implementation

The implementation can be found in the [bpf_extrapolation branch of my fork](https://github.com/MArpogaus/zuko/tree/bpf_extrapolation).

My changes specifically include

1. Optional linear extrapolation in the [call](https://github.com/MArpogaus/zuko/blob/1c48b9b748b0377d47d5b6a1d8806c8b561c71cc/zuko/transforms.py#L695) method
2. Custom Implementation of [log_abs_det_jacobian](https://github.com/MArpogaus/zuko/blob/1c48b9b748b0377d47d5b6a1d8806c8b561c71cc/zuko/transforms.py#L703) since the gradient dos not seam to pass through the `torch.where` statement and to [improve numerical stability ind the sigmoidal case](https://github.com/MArpogaus/zuko/blob/1c48b9b748b0377d47d5b6a1d8806c8b561c71cc/zuko/transforms.py#L714).

Hints from issue comments:
@francois-rozet: Hello @MArpogaus, thanks for the feature request and PR! I am curious about the issues you had with gradients. Normally `torch.where` should work fine with autograd as it computes more or less `cond * left + (1 - cond) * right`, except if `left` or `right` is NaN. So my guess is that evaluating the polynomial outside of $[0, 1]$ returns a NaN, which breaks autograd.

A solution could be to cast $x$ to 0.5 when $x \not \in [0, 1]$ before evaluating the polynomial.

Also how did you handle the inverse transformation with the extrapolation? I guess you would also need to look whether $y$ is larger/smaller than the values at the boundaries.@MArpogaus: > I am curious about the issues you had with gradients. Normally `torch.where` should work fine with autograd as it computes more or less `cond * left + (1 - cond) * right`, except if `left` or `right` is NaN. So my guess is that evaluating the polynomial outside of [0,1] returns a NaN, which breaks autograd.
> 
> A solution could be to cast x to 0.5 when x∉[0,1] before evaluating the polynomial.

I had a similar problem in my TF implementation, and casting to 0.5 helped.
But I did not succeed with it here. But I might give it another try.

Is it preferred to use autograd over analytical Jacobians?

Additionally, in the sigmoidal case, we discovered numerical issues since the derivative of sigma (sigma * (a-sigma)) converges towards zero for +/- Inf:

![sigma](https://github.com/probabilists/zuko/assets/38564291/db63b683-c109-406e-8abb-b3c528b5a89a)

Hence, we decided to [add `finfo(dtype).tiny`](https://github.com/MArpogaus/zuko/blob/924fb6d940b8540a5bbf3a72ca217a8f584fc327/zuko/transforms.py#L718) to avoid Infs when calculating the logarithm.

> Also how did you handle the inverse transformation with the extrapolation? I guess you would also need to look whether y is larger/smaller than the values at the boundaries.

Good, point! I did not think of the inverse yet, but of course, the bisection is limited to the bounds.
Thanks for pointing this out, I'll look into this.@francois-rozet: > Is it preferred to use autograd over analytical Jacobians?

An analytical Jacobian is fine (if it is exact), but if the output $y$ of the transformation is not differentiable, the transformation itself cannot be trained through gradient descent.

Can you provide a small example demonstrating the incorrect behavior?

> Additionally, in the sigmoidal case, we discovered numerical issues since the derivative of sigma (sigma * (a-sigma)) converges towards zero for +/- Inf

I see, but is it problematic in practice? Neural network layers expect standardized features (zero mean, unit variance), so the features are never too large.

Anyway if you have a working implementation with the extrapolation, we can drop the sigmoid mapping.@MArpogaus: > A solution could be to cast x to 0.5 when x∉[0,1] before evaluating the polynomial.

This actually solved the problem, my binary mask was wrong in the first place.
I update the code accordingly and drop the analytical Jacobian in [a recent commit](https://github.com/probabilists/zuko/blob/f4d9e088c386b624e3ee1532525f16a524943fba/zuko/transforms.py#L697).

> Anyway if you have a working implementation with the extrapolation, we can drop the sigmoid mapping.

[dropped it](https://github.com/probabilists/zuko/pull/37/commits/f4d9e088c386b624e3ee1532525f16a524943fba).

> Also how did you handle the inverse transformation with the extrapolation? I guess you would also need to look whether is larger/smaller than the values at the boundaries.

I now added a [`_inverse`](https://github.com/probabilists/zuko/blob/a8260ef59038badc07c770585af6b352146f2de8/zuko/transforms.py#L719) function with inverse extrapolation.

![extra_with_invers](https://github.com/probabilists/zuko/assets/38564291/c79f2aa6-5e2d-4351-ba5f-9c51574a9c66)@MArpogaus: I have added an additional option to get a smooth transition into the extrapolation by enforcing the second order derivative to be zero on the bounds ([b2869d4](https://github.com/probabilists/zuko/pull/37/commits/b2869d42e9c1beb82b1c737a49fccd1ef1c21e65)).

![without_smooth_bounds](https://github.com/probabilists/zuko/assets/38564291/c7d44db5-c1eb-48ee-9503-bfdb43551fbb)

![with_smooth_bounds](https://github.com/probabilists/zuko/assets/38564291/4571c2ee-ef63-4edd-a4c3-edd9656a41a8)

@oduerr what do you think? Do we need such an option?

I was thinking about adding an other optional augment to specify the codomain of the polynomial.

first minimal working example:

```python
low = -3
high = 33

positive_fn = torch.exp
theta = torch.rand(size=(M,))  # creates a random parameter vector
theta_low = low # - theta[:1].exp() optionally allow flexible bounds
theta_high = high # + theta[-1:].exp()
diff = torch.nn.functional.softmax(theta, dim=-1) # use theta[1:-1] if flexible_bounds=True
diff *= (theta_high - theta_low) - 2 * bpoly.eps

torch.cat([torch.cumsum(torch.cat([theta_low, diff], dim=-1), dim=-1), theta_high], dim=-1)
```

This is especially usefully when chaining several Transforms or using a base distribution with limited support.@francois-rozet: > I have added an additional option to get a smooth transition into the extrapolation by enforcing the second order derivative to be zero on the bounds.

Nice! I would actually enable this by default if it is not too expensive.

> I was thinking about adding an other optional augment to specify the codomain of the polynomial.

Why not! However, instead of specifying the co-domain, it is probably better to fix the co-domain to `[-B, B]` such that stacking several Bernstein transformations is safe and the co-domain covers well the standard Gaussian distribution. Note that this would not trivially enable base distributions with a limited support (e.g. uniform) as the extrapolation could lead to out-of-support values.

Note that implementing the `[-B, B]` co-domain would allow to drop the `SoftclipTransform` between transformations that are currently necessary in `BPF`.@MArpogaus: > Nice! I would actually enable this by default if it is not too expensive.

Ok, i enabled it by default in a [recent commit](https://github.com/probabilists/zuko/pull/37/commits/803aeb7893d90fa4420de26922b0858b0720fc27).

> Why not! However, instead of specifying the co-domain, it is probably better to fix the co-domain to `[-B, B]` such that stacking several Bernstein transformations is safe and the co-domain covers well the standard Gaussian distribution. Note that this would not trivially enable base distributions with a limited support (e.g. uniform) as the extrapolation could lead to out-of-support values.

I can add it like this for now, but I certainly need asymmetric bounds for my own models.
Could you maybe roughly scat what has to be adjusted in order for Zuko to support this?
Then can work on this in a separate PR.

> Note that implementing the `[-B, B]` co-domain would allow to drop the `SoftclipTransform` between transformations that are currently necessary in `BPF`.

Thanks for the hint. I'll make it conditional, depending on a new `keep_in_bounds` attribute inside the Transformation.
@MArpogaus: > Thanks for the hint. I'll make it conditional, depending on a new keep_in_bounds attribute inside the Transformation.

This is not possible, as the transform is not initialized when entering the BPF constructor.
We can either add an this as an argument to BPF, or just keep it in the Transform for now, for use in custom flows.

@francois-rozet  What is your opinion?@MArpogaus: Just pushed a first draft: [13ead7d](https://github.com/probabilists/zuko/pull/37/commits/13ead7d156a66be307928e32ea313978153f4e0b)@francois-rozet: What I had in mind is to make the Bernstein transform such that it **always** (not as an option) maps $\pm B$ to $\pm B$. Like this you can stack many transformations without any risk of collapse. This is important because the co-domain should always cover the base distribution support.

This is what the rational quadratic spline transform does. It maps B to B and extrapolates linearly outside. In addition, the derivatives at the bounds are 1 so that $y = x$ when $x$ is outside of the bounds.

Note that if you need another (asymmetric) domain or co-domain you can always combine a transformation with a (monotonic) affine transformation.

What do you think?@MArpogaus: Ok, that is basically what i implemented, but i did not yet enable it by default.

What is your opinion on chaining flexible transformations?
I normally consider this an unnecessary increase in complexity and prefer single layer models with just one very flexible polynomial.
But its certainly nice to have the ability of chaining them.

Currently the [`_constrain_theta`](https://github.com/probabilists/zuko/blob/13ead7d156a66be307928e32ea313978153f4e0b/zuko/transforms.py#L664) takes the argument `bound` and transforms the coefficients accordingly.

If we want we could also force the derivative to 1 on the bound to ensure we get the identity function when extrapolating.

Could look sth like this:

```python
    def _constrain_theta(unconstrained_theta: Tensor, bound: float) -> Tensor:
        r"""Processes the unconstrained output of the hyper-network to be increasing."""

        if bound:
            theta_min = -bound * torch.ones_like(unconstrained_theta[..., :1])

            def fn(x):
                return torch.cat(
                     (
                        torch.ones_like(unconstrained_theta[..., :2]), #  f'(0) = 1, f''(0)=0
                        torch.nn.functional.softmax(x, dim=-1) * (2 * bound - 4),
                        torch.ones_like(unconstrained_theta[..., :2])  #  f'(1) = 1, f''(1)=0
                     ),
                    dim = -1
                 )
        else:
            shift = math.log(2.0) * unconstrained_theta.shape[-1] / 2
            theta_min = unconstrained_theta[..., :1] - shift
            unconstrained_theta = unconstrained_theta[..., 1:]
            fn = torch.nn.functional.softplus

        widths = fn(unconstrained_theta)

        widths = torch.cat((theta_min, widths), dim=-1)
        theta = torch.cumsum(widths, dim=-1)

        return theta
```

However, i would personally prefer to keep an option for a simple "ordered" constrain function in the library, as we had in the beginning for less restrictive setups..
@oduerr What is your opinion?@francois-rozet: > What is your opinion on chaining flexible transformations? I normally consider this an unnecessary increase in complexity and prefer single layer models with just one very flexible polynomial.

In multivariate flows, it is generally necessary to chain several multi-variate (think autoregressive or coupling) transformations, even if the univariate transformations are very expressive. And even if a single multi-variate transformation is used, you should ensure that its co-domain covers the support of the base distribution.

> However, i would personally prefer to keep an option for a simple "ordered" constrain function in the library, as we had in the beginning for less restrictive setups..

You mean the one currently in the lib or one with extrapolation + smooth bounds?

Anyway, instead of adding more and more options to the class you can also sub-class it (e.g. `BoundedBernsteinTransform`).@oduerr: Sorry for the late reply. 

We (@MArpogaus and I) just had a discussion and are a bit puzzled by:
> In multivariate flows, it is generally necessary to chain several multi-variate (think autoregressive or coupling) transformations, even if the univariate transformations are very expressive. 

We thought that if the 1-D transformation function is flexible enough, utilizing the AR flows allows us to express even complex distributions due to the chain rule of prob. See also [1]. While we still believe this is true in theory, we observed much better performance on the UCI Benchmarks when we did chaining. Do you have intuition why one gets a better performance when chaining, or do you even know a paper?

@MArpogaus will implement two classes as you suggested soon (in the following week): one unbounded and one Bounded in which the coefficients are scaled. Sorry for the delay.

--- 
[1] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan, “Normalizing Flows for Probabilistic Modeling and Inference,” Journal of Machine Learning Research, vol. 22, no. 57, pp. 1–64, 2021.


@francois-rozet: Hello @oduerr, no problem, everyone is busy :smile:

> We thought that if the 1-D transformation function is flexible enough, utilizing the AR flows allows us to express even complex distributions due to the chain rule of prob. See also [1]. While we still believe this is true in theory, we observed much better performance on the UCI Benchmarks when we did chaining. Do you have intuition why one gets a better performance when chaining, or do you even know a paper?

You are right, a single auto-regressive transformation should be enough if the uni-variate transformation is an universal (monotonic) function approximator. However, the hyper network (MADE/MaskedMLP) conditioning the transformation must also be a universal function approximator. In practice, the capacity of the hyper network is finite and the function it should approximate might be complex depending on the data, the uni-variate function (and its parametrization) and the auto-regressive order. The latter, in particular, can have a huge impact: some orders can lead to very simple functions while others can lead to almost unlearnable functions. Stacking several multi-variate transformations with different (sometimes randomized) orders help to alleviate this issue.

For example if $p(x_1, x_2, ..., x_n) = p(x_1) \prod p(x_i | x_1)$ (star-shaped dependency graph), the order $1 \to n$ leads to a simple function, while the order $n \to 1$ does not. This is somewhat related to finding the Bayesian network (a directed acyclic graph) with the least number of edges that explains the data. In fact, the "simplest" Bayesian network is sometimes confused with the causal graph of the data.@oduerr: Hello @francois-rozet,

thanks for your patience. And many thanks for the intuitive example that makes it very clear! @MArpogaus: Hello everybody,

I have finally found the time to continue working on this.
Here is what i got so far:

There are now two alternative versions of the Bernstein Transform:

The class [`BernsteinTransform`](https://github.com/probabilists/zuko/blob/83727f516c6133eea6d786c086c3ea8c99a43906/zuko/transforms.py#L603) is similar to the previous implementation, but linearly extrapolates outside the bounds and applies the "smooth bound constrain" to the coefficients, discussed above.
This is not configurable anymore, as we agreed to implement different behavior via sub classing.
However, i don't have a strong opinion here, if you want we could easily add a conditional for that.

![bpoly_extra_with_invers](https://github.com/probabilists/zuko/assets/38564291/a64f4f03-1424-435c-b003-e999468f4414)

Additionally [`BoundedBernsteinTransform`](https://github.com/probabilists/zuko/blob/83727f516c6133eea6d786c086c3ea8c99a43906/zuko/transforms.py#L734) is an alternative implementation optimized for chaining.
There are two major differences:

- The coefficients are scaled into the interval $[-B, B]$
- Additionally $Be'(0,1)$ are forced to $1$ to get the identity function outside the bounds.

![bounded_bpoly_extra_with_invers](https://github.com/probabilists/zuko/assets/38564291/7a4a5bc8-2c4a-4139-81a8-b39e820bac40)

For now [`BPF`](https://github.com/probabilists/zuko/blob/83727f516c6133eea6d786c086c3ea8c99a43906/zuko/flows/polynomial.py#L93C32-L93C57) uses the bounded Version.

What are your opinions?

I would like to discuss two remaining questions:

- Should we provide an alternative `BPF` Implementation based on the previous attempt with soft clip and fewer constraints on the parameters?
- Should we Provide a univariate version using the unbounded polynomial?@francois-rozet: This looks really nice! I am curious to see how it compares to the other flows, especially NSF.

> Should we provide an alternative BPF implementation based on the previous attempt with soft clip and fewer constraints on the parameters?

Unless you think it is necessary, I am fine with (and actually prefer) only proposing BPF with `BoundedBernsteinTransform`. Softclip usually deteriorates the performances quite a lot. 

> Should we provide a univariate version using the unbounded polynomial?

I don't think so. Users can construct their own flow with `BernsteinTransform` if necessary.

```python
Flow(
    transform=ElementWiseTransform(
        features=1,
        ...
        univariate=BernsteinTransform,
        shape=[(16,)],
    ),
    base=...,
)
```@MArpogaus: Hey @francois-rozet, thanks for the fast response!

I am fine with it too.
How about you @oduerr , dos this also meet your requirements?@oduerr: Hello @Marcel,

Thanks a lot for the changes; they look fine. The idea of fixing the coefficients for a bounded transformation is quite clever!
There is no need from my side for an unbounded flow as I can roll my own easily, as described in   https://github.com/probabilists/zuko/issues/36#issuecomment-2022804132 @MArpogaus: > This looks really nice! I am curious to see how it compares to the other flows, especially NSF.

some early results on `"checkerboard"` data, with default `BPF` params:

# Dataset

![dataset](https://github.com/probabilists/zuko/assets/38564291/88b02bdf-5fce-43c3-bca6-95c70b8b076e)

# Samples

![samples](https://github.com/probabilists/zuko/assets/38564291/da4ff810-2c4b-4ae8-a8aa-abcaa9f7c53c)

# Evolution through training

![ani](https://github.com/probabilists/zuko/assets/38564291/118da8b1-6a69-46d9-96e4-adf800fab957)

# Density

![density](https://github.com/probabilists/zuko/assets/38564291/8f2eef22-8984-4534-ba8e-17502e55aa37)
![density3d](https://github.com/probabilists/zuko/assets/38564291/fad5a96a-ce2c-4856-a013-7d86aba31870)

@francois-rozet do you have some results from other flows that you can share?@francois-rozet: Hello @MArpogaus, here are the results for `NSF` with basic settings.

<img src="https://github.com/probabilists/zuko/assets/37352336/de1885b4-e324-4812-b2f5-28d5f18d524c" width="50%">

You can check the run at https://wandb.ai/francois-rozet/zuko-benchmark-2d/runs/5xllo4bz .
@francois-rozet: Closing as #37 was merged :partying_face:

Solution patch:
diff --git a/README.md b/README.md
index 1262915..260d26c 100644
--- a/README.md
+++ b/README.md
@@ -90,7 +90,7 @@ For more information, check out the documentation and tutorials at [zuko.readthe
 | `UNAF`  | 2019 | [Unconstrained Monotonic Neural Networks](https://arxiv.org/abs/1908.05164) |
 | `CNF`   | 2018 | [Neural Ordinary Differential Equations](https://arxiv.org/abs/1806.07366) |
 | `GF`    | 2020 | [Gaussianization Flows](https://arxiv.org/abs/2003.01941) |
-| `BPF`   | 2020 | [Bernstein-Polynomial Normalizing Flows](https://arxiv.org/abs/2204.13939) |
+| `BPF`   | 2020 | [Bernstein-Polynomial Normalizing Flows](https://arxiv.org/abs/2004.00464) |
 
 ## Contributing
 
diff --git a/docs/conf.py b/docs/conf.py
index a94dd77..10ccd70 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -12,7 +12,7 @@ import zuko
 package = 'zuko'
 project = 'Zuko'
 version = zuko.__version__
-copyright = '2022-2023'
+copyright = '2022-2024'
 repository = 'https://github.com/probabilists/zuko'
 commit = subprocess.check_output(['git', 'rev-parse', 'HEAD'], text=True).strip()
 
diff --git a/zuko/flows/polynomial.py b/zuko/flows/polynomial.py
index 5b544fa..59504ce 100644
--- a/zuko/flows/polynomial.py
+++ b/zuko/flows/polynomial.py
@@ -5,12 +5,11 @@ __all__ = [
     'BPF',
 ]
 
-from functools import partial
 
 # isort: local
 from .autoregressive import MAF
 from .core import Unconditional
-from ..transforms import BernsteinTransform, SoftclipTransform, SOSPolynomialTransform
+from ..transforms import BoundedBernsteinTransform, SoftclipTransform, SOSPolynomialTransform
 
 
 class SOSPF(MAF):
@@ -62,14 +61,17 @@ class BPF(MAF):
     r"""Creates a Bernstein polynomial flow (BPF).
 
     Warning:
-        Invertibility is only guaranteed for features within the interval :math:`[-10,
-        10]`. It is recommended to standardize features (zero mean, unit variance)
-        before training.
+        The Bernstein polynomial is bounded to the interval :math:`[-5, 5]`. Any feature
+        outside of this domain is not transformed. It is recommended to standardize
+        features (zero mean, unit variance) before training.
 
     See also:
-        :class:`zuko.transforms.BernsteinTransform`
+        :class:`zuko.transforms.BoundedBernsteinTransform`
 
     References:
+        | Deep transformation models: Tackling complex regression problems with neural network based transformation models (Sick et al., 2020)
+        | https://arxiv.org/abs/2004.00464
+
         | Short-Term Density Forecasting of Low-Voltage Load using Bernstein-Polynomial Normalizing Flows (Arpogaus et al., 2022)
         | https://arxiv.org/abs/2204.13939
 
@@ -77,7 +79,6 @@ class BPF(MAF):
         features: The number of features.
         context: The number of context features.
         degree: The degree :math:`M` of the Bernstein polynomial.
-        linear: Whether to use a linear or sigmoid mapping to :math:`[0, 1]`.
         kwargs: Keyword arguments passed to :class:`zuko.flows.autoregressive.MAF`.
     """
 
@@ -86,18 +87,12 @@ class BPF(MAF):
         features: int,
         context: int = 0,
         degree: int = 16,
-        linear: bool = False,
         **kwargs,
     ):
         super().__init__(
             features=features,
             context=context,
-            univariate=partial(BernsteinTransform, linear=linear),
+            univariate=BoundedBernsteinTransform,
             shapes=[(degree + 1,)],
             **kwargs,
         )
-
-        transforms = self.transform.transforms
-
-        for i in reversed(range(1, len(transforms))):
-            transforms.insert(i, Unconditional(SoftclipTransform, bound=11.0))
diff --git a/zuko/transforms.py b/zuko/transforms.py
index da61ddc..7751b76 100644
--- a/zuko/transforms.py
+++ b/zuko/transforms.py
@@ -13,6 +13,7 @@ __all__ = [
     'MonotonicRQSTransform',
     'MonotonicTransform',
     'BernsteinTransform',
+    'BoundedBernsteinTransform',
     'GaussianizationTransform',
     'UnconstrainedMonotonicTransform',
     'SOSPolynomialTransform',
@@ -30,7 +31,7 @@ import torch.nn.functional as F
 
 from textwrap import indent
 from torch import BoolTensor, LongTensor, Size, Tensor
-from torch.distributions import Transform, constraints
+from torch.distributions import Distribution, Transform, constraints
 from torch.distributions.transforms import *  # noqa: F403
 from torch.distributions.utils import _sum_rightmost
 from typing import Any, Callable, Iterable, Tuple, Union
@@ -603,12 +604,18 @@ class MonotonicTransform(Transform):
 class BernsteinTransform(MonotonicTransform):
     r"""Creates a monotonic Bernstein polynomial transformation.
 
-    .. math:: f(x) = \frac{1}{M + 1} \sum_{i=0}^{M} b_{i+1,M-i+1}(\sigma(x)) \, \theta_i
+    .. math:: f(x) = \frac{1}{M + 1} \sum_{i=0}^{M} b_{i+1,M-i+1} \left( \frac{x + B}{2B} \right) \, \theta_i
 
-    where :math:`b_{i,j}` are the Bernstein basis polynomials and :math:`\sigma(x)` is
-    the sigmoid function.
+    where :math:`b_{i,j}` are the Bernstein basis polynomials.
+
+    As the polynomial :math:`f(x)` is only defined for :math:`x \in [-B, B]`, the
+    transformation linearly extrapolates it outside this domain. The second derivative
+    at the bounds is enforced to be zero for smooth extrapolation.
 
     References:
+        | Deep transformation models: Tackling complex regression problems with neural network based transformation models (Sick et al., 2020)
+        | https://arxiv.org/abs/2004.00464
+
         | Short-Term Density Forecasting of Low-Voltage Load using Bernstein-Polynomial Normalizing Flows (Arpogaus et al., 2022)
         | https://arxiv.org/abs/2204.13939
 
@@ -617,10 +624,8 @@ class BernsteinTransform(MonotonicTransform):
 
     Arguments:
         theta: The unconstrained polynomial coefficients :math:`\theta`,
-            with shape :math:`(*, M + 1)`.
-        linear: Whether to replace the sigmoid function with a linear mapping
-            :math:`\frac{x + B}{2B}`. If :py:`True`, input features are assumed to be
-            in :math:`[-B, B]`. Failing to satisfy this constraint will result in NaNs.
+            with shape :math:`(*, M - 2)`.
+        bound: The polynomial's domain bound :math:`B`.
         kwargs: Keyword arguments passed to :class:`MonotonicTransform`.
     """
 
@@ -629,42 +634,166 @@ class BernsteinTransform(MonotonicTransform):
     bijective = True
     sign = +1
 
-    def __init__(self, theta: Tensor, linear: bool = False, **kwargs):
-        super().__init__(None, phi=(theta,), **kwargs)
+    def __init__(self, theta: Tensor, bound: float = 5.0, **kwargs):
+        super().__init__(None, phi=(theta,), bound=bound, **kwargs)
 
-        self.theta = self._increasing(theta)
-        self.linear = linear
+        self.theta = self._constrain_theta(theta)
+        self.basis = self._bernstein_basis(self.order, device=theta.device, dtype=theta.dtype)
 
-        degree = theta.shape[-1]
-        alpha = torch.arange(1, degree + 1, device=theta.device, dtype=theta.dtype)
-        beta = torch.arange(degree, 0, -1, device=theta.device, dtype=theta.dtype)
+        self.offset, self.slope = self._offset_and_slope()
+
+    @property
+    def order(self) -> int:
+        return self.theta.shape[-1] - 1
+
+    def _offset_and_slope(self) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]:
+        r"""Calculates the offsets and slopes at the domain bounds for extrapolation."""
+
+        dtheta = self.order * (self.theta[..., 1:] - self.theta[..., :-1])
+        dbasis = self._bernstein_basis(
+            self.order - 1, device=self.theta.device, dtype=self.theta.dtype
+        )
 
-        self.basis = torch.distributions.Beta(alpha, beta)
+        bounds = [
+            self.theta.new_tensor(self.eps),
+            self.theta.new_tensor(1 - self.eps),
+        ]
+
+        offset = [self._bernstein_poly(x, self.theta, self.basis) for x in bounds]
+        slope = [self._bernstein_poly(x, dtheta, dbasis) for x in bounds]
+
+        return tuple(offset), tuple(slope)
 
     @staticmethod
-    def _increasing(theta: Tensor) -> Tensor:
-        r"""Processes the unconstrained output of the hyper-network to be increasing."""
+    def _constrain_theta(unconstrained_theta: Tensor) -> Tensor:
+        """Processes the unconstrained output of the hyper-network to be increasing."""
+
+        shift = math.log(2.0) * unconstrained_theta.shape[-1] / 2
+
+        theta_min = unconstrained_theta[..., :1]
+        unconstrained_theta = unconstrained_theta[..., 1:]
+
+        # ensure smooth bounds
+        unconstrained_theta = torch.cat(
+            (
+                unconstrained_theta[..., :1],
+                unconstrained_theta,
+                unconstrained_theta[..., -1:],
+            ),
+            dim=-1,
+        )
+
+        diffs = torch.nn.functional.softplus(unconstrained_theta)
+        diffs = torch.cat((theta_min, diffs), dim=-1)
 
-        shift = math.log(2.0) * theta.shape[-1] / 2
+        theta = torch.cumsum(diffs, dim=-1) - shift
 
-        widths = torch.nn.functional.softplus(theta[..., 1:])
-        widths = torch.cat((theta[..., :1], widths), dim=-1)
+        return theta
 
-        return torch.cumsum(widths, dim=-1) - shift
+    @staticmethod
+    def _bernstein_basis(order: int, **kwargs) -> Distribution:
+        alpha = torch.arange(1, order + 2, **kwargs)
+        beta = torch.arange(order + 1, 0, -1, **kwargs)
+        basis = torch.distributions.Beta(alpha, beta)
+        return basis
+
+    @staticmethod
+    def _bernstein_poly(x: Tensor, theta: Tensor, basis: Distribution) -> Tensor:
+        b = basis.log_prob(x.unsqueeze(-1)).exp()
+        y = torch.mean(b * theta, dim=-1)
+        return y
 
     def f(self, x: Tensor) -> Tensor:
-        if self.linear:
-            x = (x + self.bound) / (2 * self.bound)  # map [-B, B] to [0, 1]
-        else:
-            x = torch.nn.functional.sigmoid(x)  # map [-inf, inf] to [0, 1]
+        x = (x + self.bound) / (2 * self.bound)  # map [-B, B] to [0, 1]
 
-        x = x * (1 - 2e-6) + 1e-6
-        x = x.unsqueeze(-1)
-        b = self.basis.log_prob(x).exp()
-        y = torch.mean(b * self.theta, dim=-1)
+        lower_bound = x <= self.eps
+        upper_bound = x >= 1 - self.eps
+        x_safe = torch.where(lower_bound | upper_bound, 0.5 * torch.ones_like(x), x)
+
+        y = self._bernstein_poly(x_safe, self.theta, self.basis)
+
+        # f'(eps) * (x - eps) + f(eps)
+        y0 = self.slope[0] * (x - self.eps) + self.offset[0]
+
+        # f'(1-eps) * (x - 1 - eps) + f(1-eps)
+        y1 = self.slope[1] * (x - 1 + self.eps) + self.offset[1]
+
+        y = torch.where(lower_bound, y0, y)
+        y = torch.where(upper_bound, y1, y)
 
         return y
 
+    def _inverse(self, y: Tensor) -> Tensor:
+        left_bound = y <= self.offset[0]
+        right_bound = y >= self.offset[1]
+
+        x = super()._inverse(y)
+        x0 = (y - self.offset[0]) / self.slope[0] + self.eps
+        x1 = (y - self.offset[1]) / self.slope[1] - self.eps + 1
+
+        # map [0, 1] to [-B, B]
+        x0 = x0 * 2 * self.bound - self.bound
+        x1 = x1 * 2 * self.bound - self.bound
+
+        x = torch.where(left_bound, x0, x)
+        x = torch.where(right_bound, x1, x)
+
+        return x
+
+
+class BoundedBernsteinTransform(BernsteinTransform):
+    r"""Creates a bounded Bernstein polynomial transformation.
+
+    This subclass of :class:`BernsteinTransform` scales the Bernstein coefficients so
+    that the polynomial's domain and codomain match the interval :math:`[-B, B]`. It
+    also enforces that the first derivative at the bounds is one and that the second
+    derivative is zero, ensuring a smooth transition to the identity function outside
+    the bounds.
+
+    These constraints make the transformation suitable for chaining in flows.
+
+    Arguments:
+        theta: The unconstrained polynomial coefficients :math:`\theta`,
+            with shape :math:`(*, M - 5)`.
+        kwargs: Keyword arguments passed to :class:`BernsteinTransform`.
+    """
+
+    def _constrain_theta(self, unconstrained_theta: Tensor) -> Tensor:
+        theta_min = -self.bound * torch.ones_like(unconstrained_theta[..., :1])
+
+        diff_on_bounds = (2 * self.bound) / (unconstrained_theta.shape[-1] + 4)
+        diffs = torch.nn.functional.softmax(unconstrained_theta, dim=-1) * (
+            2 * self.bound - 4 * diff_on_bounds
+        )
+
+        # ensure identity on bounds by enforcing Be'(0,1) = 1 and Be''(0,1) = 0
+        # Be'(0) = order * theta_1 - theta_0 = order * diff_0 -> diff_0 = 1 / order
+        # Be''(0) = (order - 1) * (diff_1 - diff_0) -> diff_0 = diff_1
+        diffs = torch.cat(
+            (
+                theta_min,
+                diff_on_bounds * torch.ones_like(diffs[..., :2]),
+                diffs,
+                diff_on_bounds * torch.ones_like(diffs[..., :2]),
+            ),
+            dim=-1,
+        )
+
+        return torch.cumsum(diffs, dim=-1)
+
+    def _offset_and_slope(self) -> Tuple[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor]]:
+        offset = (
+            self.theta.new_tensor(-self.bound),
+            self.theta.new_tensor(self.bound),
+        )
+
+        slope = (
+            self.theta.new_tensor(2 * self.bound),
+            self.theta.new_tensor(2 * self.bound),
+        )
+
+        return offset, slope
+
 
 class GaussianizationTransform(MonotonicTransform):
     r"""Creates a gaussianization transformation.


Test patch:
diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index 474346c..53e777f 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -25,7 +25,7 @@ def test_univariate_transforms(batched: bool):
         MonotonicRQSTransform(randn(*batch, 8), randn(*batch, 8), randn(*batch, 7)),
         MonotonicTransform(lambda x: x**3),
         BernsteinTransform(randn(*batch, 16)),
-        BernsteinTransform(randn(*batch, 16), linear=True),
+        BoundedBernsteinTransform(randn(*batch, 16)),
         GaussianizationTransform(randn(*batch, 8), randn(*batch, 8)),
         UnconstrainedMonotonicTransform(lambda x: torch.exp(-(x**2)) + 1e-2, randn(batch)),
         SOSPolynomialTransform(randn(*batch, 3, 5), randn(batch)),