A Robust Normalizing Flow using Bernstein-type Polynomials

The degree $n$ Bernstein polynomials, $\binom{n}{k}x^{k}(1-x)^{n-k},\,0\leq k\leq n,\,n\in\mathbb{N}$, were first introduced by Bernstein in his constructive proof of the Weierstrass theorem in Bernstein (1912). In fact, given a continuous function $f:[0,1]\to\mathbb{R}$, its degree $n$ Bernstein approximation, $B_{n}(f):[0,1]\to\mathbb{R}$, given by

$$B_{n}(f)(x)=\sum_{k=0}^{n}f\left(\frac{k}{n}\right)\binom{n}{k}x^{k}(1-x)^{n-k},\vspace{-5pt}$$

(1)

is such that $B_{n}(f)\to f$ uniformly in $[0,1]$ as $n\to\infty$. Moreover, Bernstein polynomials form a basis for degree $n$ polynomials on $[0,1]$. More generally, polynomials of Bernstein-type can be defined as follows.

As we shall see below, one can control various properties of Bernstein-type polynomials like strict monotonicity, range and universality by specifying conditions on the coefficients, and the error of approximation depends on the degree of the polynomials used.

The supports of distributions of samples used when training and applying NFs are not fixed. So, it is important to be able to easily control the range of the coupling functions. In the case of Bernstein-type polynomials, $B_{n}$s, this is very straightforward. Note that if $B_{n}$ is defined on $[a,b]$, then $B_{n}(a)=\alpha_{0}$ and $B_{n}(b)=\alpha_{n}$. Therefore, one can fix the values of a Bernstein-type polynomial at the end points of $[a,b]$ by fixing $\alpha_{0}$ and $\alpha_{n}$. So, if $B_{n}$ is increasing (which will be the case in our model; see Section 2.3), then its range is $[\alpha_{0},\alpha_{n}]$. This translates to a significant advantage when training for compactly supported targets because we can achieve any desired range $[c,d]$ (the support of the target) by fixing $\alpha_{0}=c$ and $\alpha_{n}=d$ and letting only $\alpha_{k},\,0<k<n$ vary. So, $B_{n}$s are ideal for modeling compactly supported targets. In fact, in most other methods except splines in Durkan et al. (2019a, b), either there is no obvious way to control the range or the range is infinite. We present the following theorem to establish that for the purpose of training, we can assume the target has compact support (up to a known diffeomorphism).

In the previous literature that uses transformations with compact range, this fact was implicitly assumed without proper justification. As a consequence of the above theorem, our coupling functions having finite ranges is not a restriction, and in any NF model, even if the target density is not compactly supported, the learning procedure can be implemented by first converting the target density to a density with a suitable compact support via a diffeomorphism, and then training on the transformed data. Since we deal with compactly supported targets, in practice, we do not need to construct deep architectures (with a higher number of layers), as we can increase the degree of the polynomials to get a better approximation. In other polynomial based methods, a practical problem arises because the higher order polynomials could predict extremely high values initially leading to unstable gradients (e.g., Jaini et al. (2019)). In contrast, we can avoid that problem as the range of our transformations can be explicitly controlled from the beginning by fixing $\alpha_{0}$ and $\alpha_{n}$ .

In triangular flows, the coupling maps are expected to be invertible. Since strict monotonocity implies invertibility, it is sufficient that the $B_{n}$s we use are strictly monotone.

This result was mentioned as folklore without proof and used in Farouki (2000). On the other hand, in Lindvall (2002), the conclusion is stated with monotonicity and not strict monotonicity (which is absolutely necessary for invertibility). Hence, we added a complete proof of the statement in Appendix 2. According to this result, the strict monotonicity of $B_{n}$s depends entirely on the strict monotonicity of the coefficients $\alpha_{k}$s. It is easy to see that the assumption of strict monotonicity of the coefficients is not a further restriction on the optimization problem. For example, if the required range is $[c,d]$, we can take $\alpha_{n-k}=c+(d-c)(1+v^{2}_{0}+\dots+v^{2}_{k})^{-1}$ where $v_{k}$s are real valued. This converts the constrained problem of finding $\alpha_{k}$s to an unconstrained one of finding $v_{k}$s. Alternatively, we can take $\alpha_{0}=c$ and $\alpha_{k}=|v_{1}|+\dots+|v_{k}|$, and after each iteration, linearly scale $\alpha_{k}$s in such a way that $\alpha_{n}=d$. After guaranteeing invertibility, we focus on computing the inverse, i.e., at each iteration, given $x$ we solve for $z\in[0,1]$,

$$B_{n}(z)=\sum_{k=0}^{n}\alpha_{k}\binom{n}{k}z^{k}(1-z)^{n-k}=x\iff\sum_{k=0}^{n}(\alpha_{k}-x)\binom{n}{k}z^{k}(1-z)^{n-k}=0\vspace{-5pt}$$

(4)

because Bernstein polynomials form a partition of unity on $[0,1]$. So, finding inverse images, i.e., solving the former is equivalent to finding solutions to the latter. Due to our assumption of increasing $\alpha_{k}$s, $B_{n}$ is increasing, and has at most one root on $[0,1]$. The condition $(\alpha_{0}-x)(\alpha_{n}-x)<0$ (which can be easily checked) guarantees the existence of a unique solution, and hence, the invertibility of the original transformation.

Due to the extensive use of Bernstein-type polynomials in computer-aided geometric design, there are several well-established efficient root finding algorithms at our disposal  (Spencer, 1994). For example, the parabolic hull approximation method in Rajan et al. (1988) is ideal for higher degree polynomials with fewer roots (in our case, just one) and has cubic convergence for simple roots (better than both the bijection method and Newton’s method). Further, because of the numerical stability described in Section 2.5 below, the use of Bernstein-type polynomials in our model minimizes the errors in such root solvers based on floating–point arithmetic. Even though inverting splines are easier due to the availability of analytic expressions for roots, compared to all other other NF models, we have more efficient and more numerically stable algorithms that allow us to reduce the cost of numerical inversion in our setting.

In order to guarantee universality of triangular flows, we need to use a class of coupling functions that well-approximates increasing continuous functions. This is, in fact, the case for $B_{n}$s, and hence, we have the following theorem whose proof we postpone to Appendix 2.

The basis of all the universality proofs of NFs in the existing literature is that the learnable class of functions is dense in the class of increasing continuous functions. In contrast, the argument we present here is constructive. As a result, we can write down sequences of approximations for (known) transformations between densities explicitly; see Appendix 4.

In the case of cubic-spline NFs of Durkan et al. (2019a), it is known that for $k=1,2,3$ and $4$, when the transformation is $k$ times continuously differentiable and the bin size is $h$, the error is $O(h^{k})$ (Ahlberg et al., 1967, Chapter 2). However, we are not aware of any other instance where an error bound is available. Fortunately for us, the error of approximation of a function $f$ by its Bernstein polynomials has been extensively studied. We recall from Voronovskaya (1932) the following error bound: for $f:[0,1]\to\mathbb{R}$ twice continuously differentiable

$$\scriptsize B_{n}(f)-f=\frac{x(1-x)}{2n}f^{\prime\prime}(x)+o(n^{-1})\vspace{-5pt}$$

(5)

and this holds for an arbitrary interval $[a,b]$ with $x(1-x)$ replaced by $(x-a)(b-x)$. Since the error estimate is given in terms of the degree of the polynomials used, we can improve the optimality of our NF by avoiding unnecessarily high degree polynomials. This allows us to keep the number of trainable parameters under control in our NF model. It can be shown that the error $O(n^{-1})$ above does not necessarily improve when SOS polynomials are used instead; see Appendix 3. In our NF, at each step, the estimation is done using a univariate polynomial, and hence, the overall convergence rate is, in fact, the minimal univariate convergence rate of $O(n^{-1})$ (equivalently, the error upper bound is the maximum of univariate upper bounds), and in general, cannot be improved further regardless of how regular the density transformation is. However, our experiments (in Section 4.3) show that our model on average has a significantly smaller error than the given theoretical upper-bound.

In this section, we recall some known results in Farouki and Goodman (1996); Farouki and Rajan (1987) about the optimal stability of the Bernstein basis. The two key ideas are that smaller condition numbers lead to smaller numerical errors and that the Bernstein basis has the optimal condition numbers compared to other polynomial bases.

To illustrate this, let $p(x)$ be a polynomial on $[a,b]$ of degree $n$ expressed in terms of a basis $\{\phi_{k}\}_{k=0}^{n}$, i.e.,

$$p(x)=\sum_{k=0}^{n}c_{k}\phi_{k}(x),\,x\in[a,b].\vspace{-5pt}$$

(6)

Let $c_{k}$ be randomly perturbed, with perturbations $\delta_{k}$ where the relative error $\delta_{k}/c_{k}\in(-\varepsilon,\varepsilon)$. Then the total pointwise perturbation is

$$\delta(x)=\sum_{k=0}^{n}\delta_{k}\phi_{k}(x)\implies|\delta(x)|\leq\sum_{k=0}^{n}|\delta_{k}\phi_{k}(x)|\leq\varepsilon\sum_{k=0}^{n}|c_{k}\phi_{k}(x)|\leq\varepsilon C_{\phi}(p(x))\,,\vspace{-5pt}$$

(7)

where $C_{\phi}(p(x)):=\sum_{k=0}^{n}|c_{k}\phi_{k}(x)|$ is the condition number for the total perturbation with respect to the basis $\phi_{k}$. It is clear that $C_{\phi}(p(x))$ controls the magnitude of the total perturbation.

According to Farouki and Goodman (1996), if $\phi=\{\phi_{k}\}_{k=0}^{n}$ and $\psi=\{\psi_{k}\}_{k=0}^{n}$ are non-negative bases for polynomials of degree $n$ on $[a,b]$, and for all $j$, latter is a non-negative linear combination of the former, that is, $\psi_{j}=\sum_{k=0}^{n}M_{jk}\phi_{k}$ with $M_{jk}\geq 0$. Then, for any polynomial $p(x)$,

$$C_{\phi}(p(x))\leq C_{\psi}(p(x))\,.\vspace{-5pt}$$

(8)

For $0\leq a<b$, the Bernstein polynomials and the power monomials, $\{1,x,x^{2},\dots,x^{n}\}$, are non-negative bases on $[a,b]$. It is true that the latter is a positive linear combination of the former but not vice-versa; see Farouki and Goodman (1996). Therefore, Bernstein polynomial basis has the lowest condition number out of the two. This means that the change in the value of a polynomial caused by a perturbation of coefficients is always smaller in the Bernstein basis than in the power basis. A more involved computation gives

$$\tilde{C}_{\phi}(x_{0})=\left(\frac{m!}{|p^{(m)}(x_{0})|}\sum_{k=0}^{n}|c_{k}\phi_{k}(x_{0})|\right)^{1/m}\vspace{-5pt}$$

(9)

as the condition number that controls the computational error for a $m-$fold root $x_{0}$ of $p(x)$ in $[a,b]$; see Farouki and Rajan (1987). There, it is proved that if $\tilde{C}_{\psi}(x_{0})$ and $\tilde{C}_{\phi}(x_{0})$ denote the condition numbers for finding roots of any polynomial on $[0,1]$ in the power and the Bernstein bases on $[0,1]$, respectively, then $\tilde{C}_{\phi}(x_{0})<\tilde{C}_{\psi}(x_{0})$ for $x_{0}\in(0,1]$ and $\tilde{C}_{\phi}(0)=\tilde{C}_{\psi}(0)$. This means that the change in the value of a root of a polynomial caused by a perturbation of coefficients is always smaller in the Bernstein basis than in the power basis.

In fact, a universal statement is true: Among all non-negative basis on a given interval, the Bernstein polynomial basis is optimally stable in the sense that no other non-negative basis gives smaller condition numbers for the values of polynomials (see (Peña, 1997, Theorem 2.3)) and no other basis expressible as non-negative combinations of the Bernstein basis gives smaller condition numbers for the roots of polynomials (see (Farouki and Goodman, 1996, Section 5.3)).

In particular, $B_{n}$s are systematically more stable than the polynomials in the power form when determining roots (for example, when inverting) and evaluation (for example, when finding image points). As a result, when polynomials are used to construct NFs (for example, Q-NSF based on quadratic or cubic splines, SOS based on some of square polynomials and our NF based on Bernstein-type polynomials, ours yields the most numerically stable NF, i.e., it is theoretically impossible for them to be more robust. Our experiments in Section 4.2, while confirming this, demonstrate that our method outperforms even the NFs that are not based on polynomials.

We conducted experiments on four datasets from the UCI machine-learning repository and BSDS300 dataset. Table 1 compares the obtained test log-likelihood against recent flow-based models. As illustrated, our model achieves competitive results on all of the five datasets. We observe that our model consistently reported a lower standard deviation which may be attributed to the robustness of our model.

We also applied our method to two low-dimensional image datasets, CIFAR10 & MNIST. The results are reported in Table 2. Among the methods that do not use multi-scale convolutional architectures, we obtain the optimal results. In addition, we tested our model on several toy datasets (shown in Figure 3). Note that these 2D datasets contain multiple modes, sharp jumps and are not fully supported. So, the densities are not that obvious to learn. Despite the difficulties, our model is able to estimate the given distributions in a satisfactory manner.

In order to experimentally verify that Bernstein-type NFs are more numerically stable than other polynomial based NFs (as claimed in Section 2.5) we use a standard idea in the literature; see Taylor (1982).

We add i.i.d. noise, sampled from a Uniform$[0,10^{-2}]$, to the five datasets included in Table 1, and measure the change in the test log-likelihood as a fraction of the standard deviation (so that the change is in terms of standard deviations). In practice, values of the experimented datasets are rescaled to a magnitude around unity. In signal processing, a good SNR is considered to be above 40DB which is used in most real-world cases. Here, we have chosen a noise order of $10^{-2}$ because our intention is to demonstrate that a SNR level below or even around that range can affect the performance of NFs.

For a fair comparison, we train all the models from scratch on the noise-free train set using the codes provided by the authors, strictly following the instructions in the original works to the best of our ability. Then, we test the models on the noise-free test set. We run the above experiment $5$ times to obtain the standard deviation $\sigma$ and mean $\mu$ of the test log-likelihood. Next, we add noise to the training set, retrain the model, and obtain the test log-likelihood $y$ on the noise-free test set. Finally, we obtain the metric $\frac{|y-\mu|}{\sigma}$ which we report in Table 3.

As expected, Bernstein NF demonstrate the lowest relative change in performance, implying the robustness against random initial errors. In fact, other models are not robust: even small initial errors consistently (at least in 4 out of 5 datasets) created changes larger than $1.645\,\sigma$ (corresponding to the two $5$% tails of the distribution of errors) where $\sigma$ is the original standard deviation of error. In comparison, the change in our NF consistently (4 out of 5) was well-within the acceptable range and is at most $1.3\,\sigma$. In the remaining dataset, change in our model is $2.3\,\sigma$ while in all other models it’s more than $6\,\sigma.$

The degree $n\,(\geq 5)$ Bernstein approximation of $f\in C^{3}[0,1]$ has an error upper-bound

$$E_{n}=n^{-1}\|\rho^{2}f^{(2)}\|_{\infty}+n^{-3/2}\|\rho^{3}f^{(3)}\|_{\infty}\vspace{-5pt}$$

(10)

where $\rho(x)=\sqrt{x(1-x)}$ (Bustamante, 2017, Chapter 4). Now, we verify this using a Kumarswamy$(2,5)$ distribution as the prior and Uniform$[0,1]$ as the target. Let $f(x)=1-(1-x^{2})^{5}$ and $B_{n}$ be the learned degree $n$ Bernstein-type polynomial. The average error, $\int_{0}^{1}|f(x)-B_{n}(x)|\,dx,$ obtained using the learned $B_{n}$s and $E_{n}$ satisfying

$$1.25n^{-1}+5n^{-3/2}<E_{n}<1.25n^{-1}+5.5n^{-3/2}\vspace{-5pt}$$

(11)

are plotted in Figure 3. It shows that the observed (average) error is smaller than this theoretical upper-bound. In the NF, we have used a single layer and increased the degree of the polynomial from 10 to 100. The NF model was stable even when the degree 100 polynomial was used. So, this experiment also demonstrates that our model is, in fact, stable even when higher degree polynomials are used (as claimed in Section 2.2).

NFs learn an invertible mapping between a prior and a more complex distribution (the target) in the same dimension. Typically, the prior is chosen to be a Gaussian with identity covariance or uniform on the unit cube, and the target is the one we intend to learn. Below, we present a summary of related ideas and refer the readers to Jaini et al. (2019) and Kobyzev et al. (2020) for a comprehensive discussion.

More formally, let z and x be sampled data from the prior with density $P_{z}$ and the target distribution with density $P_{x}$, respectively. Then, NFs learn a transformation $f$ such that $f(\textbf{z})=\textbf{x}$ which is differentiable and invertible with a differentiable inverse. Such transformations are called diffeomorphisms and they allow the estimation of the probability density $P_{x}(\textbf{x})$ via the change of variables formula $P_{x}(\textbf{x})=P_{z}(f^{-1}\textbf{x})|\textbf{J}_{f}(f^{-1}\textbf{x})|^{-1}$ where $\textbf{J}_{f}$ is the Jacobian determinant of $f$.

Given an independent and identically distributed (i.i.d.) sample $\{x_{1},\dots,x_{n}\}$ with law $P_{x}$, learning the target density $P_{x}$ and the transformation $f$ (within an expressive function class $\mathfrak{F}$) is done simultaneously via minimizing the Kullback-Leibler (KL) divergence between $P_{x}$ and the pushforward of $P_{z}$ under $f$ denoted by $f_{*}P_{z}$,

$$\min_{f\in\mathfrak{F}}\text{KL}(P_{x}\|f_{*}P_{z})=-\max_{f\in\mathfrak{F}}\int\log\frac{P_{z}(f^{-1}\textbf{x})}{|\textbf{J}_{f}(f^{-1}\textbf{x})|}\,\cdot\,P_{x}(\textbf{x})\,d\textbf{x}.\vspace{-5pt}$$

(12)

Density estimation using (12) requires efficient calculation of the Jacobian as well as $f^{-1}$. Both can be achieved via constraining $f$ to be an increasing triangular map. That is, taking $P_{x}(\textbf{x})$ to be a multivariate distribution where $\textbf{x}=(x_{1},x_{2},\dots,x_{d})$, and the prior $P_{z}(\textbf{z})$ where $\textbf{z}=(z_{1},z_{2},\dots z_{d})$, the components of x are expressed as $x_{j}=f_{j}(z_{1},z_{2},\dots,z_{j})$ for suitably defined transformations $f_{j},j=1,2,\dots,d$ where $f_{j}$ is increasing with respect to $z_{j}$. From now on, we denote $(z_{1},z_{2},\dots,z_{j})$ by $z_{<j+1}$. In this case, the Jacobian determinant is the product $\prod_{j=1}^{d}\partial_{z_{j}}f_{j}$. Also, because $f_{j}$ is increasing in $z_{j}$, inversion can be done recursively starting from $f^{-1}_{1}$.