Stibor

Package to easily define normalizing flows and neural flows for Pytorch.

• Normalizing flows define complicated high-dimensional densities as transformations of random variables.
• Neural flows define continuous time dynamics with invertible neural networks.

Install package and dependencies

pip install git+https://github.com/mbilos/stribor.git

Normalizing flows

Base densities

• Normal st.Normal and st.UnitNormal and st.MultivariateNormal
• Uniform st.UnitUniform
• Other distributions from torch.distributions

Invertible transformations

• Activation functions
  • ELU st.ELU
  • Leaky ReLU st.LeakyReLU
• Affine
  • Element-wise transformation st.Affine
  • Fixed (non-learnable) element-wise transformation st.AffineFixed
  • Linear layer with PLU factorization st.AffinePLU
  • Matrix exponential st.MatrixExponential
• Coupling layer that can be combined with any element-wise transformation st.Coupling
• Continuous normalizing flows st.ContinuousNormalizingFlow
  • Differential equations with stochastic trace estimation:
    • st.net.DiffeqMLP
    • st.net.DiffeqDeepset
    • st.net.DiffeqSelfAttention
  • Differential equations with fixed zero trace:
    • st.net.DiffeqZeroTraceMLP
    • st.net.DiffeqZeroTraceDeepSet
    • st.net.DiffeqZeroTraceAttention
  • Differential equations with exact trace computation:
    • st.net.DiffeqExactTraceMLP
    • st.net.DiffeqExactTraceDeepSet
    • st.net.DiffeqExactTraceAttention
• Cummulative sum st.Cumsum and difference st.Diff
  • Across single column st.CumsumColumn and st.DiffColumn
• Permutations
  • Flipping the indices st.Flip
  • Random permutation of indices st.Permute
• Sigmoid st.Sigmoid and logit st.Logit function
• Spline (quadratic or cubic) element-wise transformation st.Spline

Example

To define a normalizing flow, define a base distribution and a series of transformations, e.g.:

import stribor as st
import torch

dim = 2
base_dist = st.UnitNormal(dim)

transforms = [
    st.Coupling(
        flow=st.Affine(dim, latent_net=st.net.MLP(dim, [64], dim)),
        mask='ordered_right_half'
    ),
    st.ContinuousNormalizingFlow(
        dim,
        net=st.net.DiffeqMLP(dim + 1, [64], dim)
    )
]

flow = st.Flow(base_dist, transforms)

x = torch.rand(1, dim)
y, ljd = flow(x)
y_inv, ljd_inv = flow.inverse(y)

Run tests

pytest