Neural Diffeomorphic Non-uniform B-spline Flows
This work addresses the need for efficient and differentiable transformations in physics simulations, offering an incremental improvement over existing spline and smooth normalizing flow methods.
The authors tackled the problem of modeling complex probability distributions with normalizing flows that require second derivatives for physics applications, proposing diffeomorphic non-uniform B-spline flows that are at least twice continuously differentiable and bi-Lipschitz continuous, resulting in better solutions for force matching in Boltzmann generators compared to previous spline flows and faster performance than smooth normalizing flows.
Normalizing flows have been successfully modeling a complex probability distribution as an invertible transformation of a simple base distribution. However, there are often applications that require more than invertibility. For instance, the computation of energies and forces in physics requires the second derivatives of the transformation to be well-defined and continuous. Smooth normalizing flows employ infinitely differentiable transformation, but with the price of slow non-analytic inverse transforms. In this work, we propose diffeomorphic non-uniform B-spline flows that are at least twice continuously differentiable while bi-Lipschitz continuous, enabling efficient parametrization while retaining analytic inverse transforms based on a sufficient condition for diffeomorphism. Firstly, we investigate the sufficient condition for Ck-2-diffeomorphic non-uniform kth-order B-spline transformations. Then, we derive an analytic inverse transformation of the non-uniform cubic B-spline transformation for neural diffeomorphic non-uniform B-spline flows. Lastly, we performed experiments on solving the force matching problem in Boltzmann generators, demonstrating that our C2-diffeomorphic non-uniform B-spline flows yielded solutions better than previous spline flows and faster than smooth normalizing flows. Our source code is publicly available at https://github.com/smhongok/Non-uniform-B-spline-Flow.