To familiarize yourself with Jax/Flax, I would recommend this notebookby University of Amsterdam's Deep learning course (which also has otherresources included). The Jax documentation is also a very good read. Ifyou are already familiar with numpy and another deep learning library,learning the basics of Jax should take less than a few hours, or you cantry learning it as you go. I personally struggle to focus on multiplethings at once so I needed to set aside a few days for each topicseparately. To get an understanding of normalizing flows (withoutreading a 50+ page paper) you can take a look at one or more of thefollowing links:

LilianWeng's blog: Good rundown of the theory and many of the commonfunctions.

Eric Jang'sblog: Theory and a tensor-flow (🤮) implementation. I found the codehard to follow since a lot of magic happens in tensor-flow. Eric hassince made a pure Jax tutorial aswell, which I highly recommend once you're comfortable with the theory(and if you don't want to bother with flax/optax).

UVADLCnotebook: Helpful, but much of the content focuses on dealing withthe specifics of image data processing and advanced architectures, whichcan make learning the fundamentals very difficult. I've simplified theircode for parts of this tutorial.

NormalizingFLows using Pyro library: Pyro tutorials are great in that theyfocus on conveying a fundamental understanding of the topic rather thanthe specifics of the library, but I think implementing your own solutionin a simplified setting is necessary before using any libraries.

I will briefly go over the basic theory behind normalizing flows as arefresher, and also to establish the variable names that will be used inthe code.

Change of Variables

To reiterate the example given by Eric Jang, lets say you have acontinuous uniform random variable and its probability density function(PDF) where is a real valued output of . Lets say we applied a function to the outputs of . The result can be treated as a newrandom variable with its own PDFfunction . Where is the output of , and . Change of variables theorem helpsdefine .

Let's say we give the range of, we know what the PDF for would look like, since it'suniform.

What would this distribution look like for ? Well, it should be clear that giventhe range of , would have the range . In this case, it's easy tocalculate without change ofvariables theorem. Intuitively, is uniform and covers the values between 1 and 3. Since the area underthe curve should sum up to 1, then the distribution for should look something like the figurebelow.

But can we prove this mathematically? Yes!

This is an example in 1D, but using change of variables theorem todetermine changes in probability distributions is the basis fornormalizing flows. The term is really thedeterminant of the Jacobin of a multi-variable function , which in a 1D environment is equal tothe derivative of . If you want toknow what a Jacobian is, Khan academy has a fewshort videos plus a quiz which I recommend (it should take about anhour at most if you already know how to take simple derivatives).

Code for 1D Example:

Defining the problem setting and solution in code might seem trivialif we know the answers already. But I found it a helpful exercise beforeexpanding this idea to a 2 or more dimensional setting.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# let's say we have a valid probability distribution where x is between [0,1] and p(x) = 1 

from scipy.stats import rv_continuous # scipy has a base class for continuous distributions. You do not have to inherit from it.

# function f, and its reverse
def fn(x, reverse = False):
    if reverse:
        return x/2 - 1
    else:
        return 2*x + 1

class x_dist(rv_continuous):
    def _pdf(self, x):
        return np.where((x>0) & (x<1),1,0)
    
px = x_dist()
x = np.linspace(-0.5,1.5,100)
h = plt.fill_between(x,px.pdf(x),color="blue")
plt.legend(handles=[h], labels=["pdf of X"])

1
2
3
4
5
6
7
8
9

# now we know that if x is being projected by f(x), then px will be projected by 1/2
# f(x) = y
# pydy = pxdx
# py = px (dx/(dy) = px (dx/(d(f(x))) = px/(f') = px/2 = 0.5px
# so Y PDF will look like this: 
y = fn(x)
py = 0.5 * px.pdf(x)
h = plt.fill_between(y,py,color = "g",)
plt.legend(handles=[h], labels=["pdf of Y"])

Normalizing Flows in 2D withJax

Let's define the same 2D dataset as Eric Jang's problem. We generate1024 samples, and each sample has 2 points. There is a relationshipbetween the two points; to create a generative model, we want to find aseries of invertible functions which can untangle the "complicated"relationships, and project each point to a unimodal normal distribution.If successful, generating new samples is trivial: we take 2 samples froma gaussian distribution and stack them as , and send backwards through the flow. In other words, transform by reversing the order of functionsand sending it through the inverse of each function. Animplication of this is that dimensionality remains the same innormalizing flows.

1
2
3
4
5
6
7
8

# define a target distribution
def get_points():
    x2 = 0 + 4 * np.random.normal()
    x1 = 0.25 * x2**2 + np.random.normal()
    return x1,x2
x_samples = np.array([get_points() for i in range(1024)])
t_dist = jnp.array(x_samples)
sns.scatterplot(x = x_samples[:,0],y = x_samples[:,1])

Regardless of your target distribution, after going through the flow,you want to final distribution to be a multi-variable normaldistribution with no covariance/correlation. We assume the marginaldistribution of each point after the transformations to be normal, andcalculate the loss based on this assumption. This means that duringtraining, the parameters of the functions are adjusted such that thepoints are projected into a multi-variate normal distribution, sosomething like this:

But how do we calculate this loss values? This is a critical stepthat did not click with me right away. If my explanation isn't clearhere please read one of the recommended posts mentioned early on. Let'sassume the flow only has 1 function, .

We know that , which can be rewritten byreplacing with thedeterminant of the Jacobian of and taking the log of both sides: If we have k (for this example, ) samples of (i.e, is the input) , and and we want togenerate new samples that fit ,what we do is project each to avalue using aparameterized function , andcalculate the PDF of the two points in under a normal distribution. From therewe can calculate (whichyou can define as the sum of probabilities for each of the two points)using the formula above. We want to maximize , the probability of the samples,therefore is minimizedduring training.

Defining Functions

We implement leakyRelu and realNVP, two common functions/layers usedin normalizing flows. When defining layers, we have to define how ittransforms an input, and also how to calculate the log determinant ofthe jacobian (LDJ). LDJs are passed through and added up in every layer,and are needed to calculated the projected probabilities. Theprobabilities are then used to calculate the loss of the network.

Relu

In a LeaklyRelu layer, if an input , it is multiplied by a value , else, it is unchanged. The LDJfor each individual point in a sample is:

Here, is either 1 or 2 sinceeach sample is in 2D. Since the final goal is to calculate the lossfunction, you can sum up the LDJ values before passing it on to thesubsequent layer.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

# using Flax's nn module
class LeakyRelu(nn.Module):        
    layer_type = "leaky_relu" # a name that will be used when plotting intermediate layers
    def setup(self):
        self.alpha = self.param('alpha', nn.initializers.uniform(0.9), (1,)) # alpha is randomly sampled from a uniform distribution between 0 and 1
    def __call__(self, z, ldj,rng, reverse=False):
        rng, new_rng = jax.random.split(rng, 2)
        if not reverse:
            dj = jnp.where(z>0, 1, self.alpha)
            z = jnp.where(z>0, z, self.alpha * z)
            ldj += jnp.log(jnp.absolute(dj)).sum(axis=[1])
            return z, ldj, new_rng
        else:
            z = jnp.where(z>0, z, z/self.alpha)
            dj = jnp.where(z>0, 1, 1/self.alpha)
            ldj -= jnp.log(jnp.absolute(dj)).sum(axis=[1]) 
            return z, ldj, new_rng
    

You can visualize the effect of the layer.

1
2
3
4
5
6
7
8
9
10
11

# initialize an lru layer 
lru = LeakyRelu()
rng, inp_rng, init_rng = jax.random.split(rng, 3)
params = lru.init(init_rng,x_samples,0,inp_rng)

# go forward and check LDJ
z = t_dist
ldj = jnp.zeros(len(z))
sns.scatterplot(x= z[:,0], y = z[:,1],label="original")
z,ldj,rng = lru.apply(params,z,0,rng,reverse=False)
sns.scatterplot(x= z[:,0], y = z[:,1],label="forward")

Inspecting the the params and LDJ valuesshould look something like this:

1
2
3
4
5
6
7

FrozenDict({
    params: {
        alpha: DeviceArray([0.73105997], dtype=float32),
    },
})
ldj list: [-0.31325978 -0.31325978  0.         ... -0.31325978 -0.31325978
  0.        ]

where the outputs can take 3 values: 0, , or (if both points were lessthan 0). Note that .

Coupling Layers

A coupling layer is a layer where a subset of points in a sample (soeither or ) are unaffected by thetransformation. However, this unchanged subset determines the change inthe remaining points. Here, isthe indices that will remain unchanged. These indices can be given toany function (typically a neural network) that returns a scalingparameter and a transformparameter . and are used to scale and shift theunmasked points.

This setup allows for coupling, or transformation of some pointsaccording to the value of other points. But we actually do not need toknow the inverse of the neural network, or whatever method we use tocalculate and . The Jacobian is simply a triangularmatrix with the diagonal equal to 1 at the masked indices, and equal to at the unmaskedindices. So the LDJ is simply the product the diagonal, or where j is the numberof unmasked indices (which is always 1 in a 2D example, since we leave 1out and change the remaining).

I've modified an example provided by University of Amsterdam deeplearning course for the 2D example. A simple neural network (also calleda hyper-network) generates the scale and transform parameters for one ofthe two points based on the value of the other. The coupling layer thentransforms the sample and calculates the LDJ.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

# modified from https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial2/Introduction_to_JAX.html
# and https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial11/NF_image_modeling.html
class SimpleNet(nn.Module):
    num_hidden : int   # Number of hidden neurons
    num_outputs : int  # Number of output neurons

    def setup(self):
        self.linear1 = nn.Dense(features=self.num_hidden)
        self.linear2 = nn.Dense(features=self.num_outputs)

    def __call__(self, x):
        # Perform the calculation of the model to determine the prediction
        x = self.linear1(x)
        x = nn.tanh(x)
        x = self.linear2(x)
        return x
    
class CouplingLayer(nn.Module):
    network : nn.Module  # NN to use in the flow for predicting mu and sigma
    mask : np.ndarray  # Binary mask where 0 denotes that the element should be transformed, and 1 not.
    c_in : int  # Number of input channels
    layer_type = "Scale and Shift"
    def setup(self):
        self.scaling_factor = self.param('scaling_factor',
                                         nn.initializers.zeros,
                                         (self.c_in,))

    def __call__(self, z, ldj, rng, reverse=False):
        # Apply network to masked input
        z_in = z * self.mask
        nn_out = self.network(z_in)
        s, t = nn_out.split(2, axis=-1)
        # Stabilize scaling output
        s_fac = jnp.exp(self.scaling_factor).reshape(1, -1)
        s = nn.tanh(s / s_fac) * s_fac
        # Mask outputs (only transform the second part)
        s = s * (1 - self.mask)
        t = t * (1 - self.mask)
        # Affine transformation
        if not reverse:
            # Whether we first shift and then scale, or the other way round,
            # is a design choice, and usually does not have a big impact
            z = (z + t) * jnp.exp(s)
            ldj += s.sum(axis=[1])
        else:
            z = (z * jnp.exp(-s)) - t
            ldj -= s.sum(axis=[1])
        return z, ldj, rng

MultiLayer Network

Finally, we want to create a model with multiple layers, train it,and sample from it when we are done training. This is very simple to dowith the networks we defined in flax, and the optimizers provided byoptax . The sampling function could be even simpler if we didn't want tosee the intermediate transformations. Notice that during sampling, wesimply generate two points from a normal distribution and pass it backthrough the flows.

Architecture

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# multi layer network
class PointFlow(nn.Module):
    flows : Sequence[nn.Module]  
    def __call__(self, z, rng, intermediate=False,training=True):
        ldj = 0 
        rng, rng_new = jax.random.split(rng,2)
        for flow in self.flows:
            z, ldj, rng = flow(z, ldj, rng, reverse=False)
        return z, ldj, rng_new
    
    def sample(self,rng,num = 1, intermediate = False):
        ldj = 0
        rng, new_rng = jax.random.split(rng,2)
        z = jax.random.normal(rng,shape=[num,2])
        intermediate_layers = [z]
        for flow in reversed(self.flows):
            z, ldj, rng = flow(z,ldj,rng,reverse=True)
            if intermediate: # if we want to see the intermediate results
                intermediate_layers.append(z)
        return z, rng, intermediate_layers

The PointsFlow model just needs a list of flow layers,we've already defined Relu and Coupling layers.

1
2
3
4
5
6
7
8

flow_layers = []
for i in range(9):
    flow_layers += [LeakyRelu()]
    flow_layers += [CouplingLayer(network = SimpleNet(num_hidden=4,num_outputs=2),mask = jnp.array([0,1]) if i%2 else jnp.array([1,0]),c_in=1)]

    
model = PointFlow(flow_layers)
params = model.init(init_rng,inp,inp_rng,intermediate=True)

Training and Sampling

Then we need to define the loss calculation and optimizer. I foundthat adamw worked much better than SGD and adam. The loss function iswhere we calculate the negative log loss, which requires andunderstanding of the change of variable rules which I discussed in notenough detail earlier. Try to derive it yourself on paper to make sureyou understand how it works (Lilian Wang's blog was very helpful for mein this regard)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

optimizer = optax.adamw(learning_rate=0.0001)
model_state = train_state.TrainState.create(apply_fn=model.apply,
                                            params=params,
                                            tx=optimizer)
# loss calculation
def loss_calc(state, params, rng, batch):
    rng, init_rng = jax.random.split(rng,2)
    z, ldj, rng  = state.apply_fn(params,batch,rng)
    # loss is ldj + log(probability of points in a normal distribution)
    log_pz = jax.scipy.stats.norm.logpdf(z).sum(axis=[1])
    log_px = ldj + log_pz
    nll = -log_px
    return nll.mean(), rng
    
nll, rng = loss_calc(model_state,params,rng,inp) # run it once and see the output 

Finally, we define a train_step function. If you have a GPU, thejitted function runs incredibly fast. Just in time compilation is one ofthe best features of Jax.

1
2
3
4
5
6
7
8
9
10
11
12
13

@jax.jit
def train_step(state,rng,batch):
    rng, init_rng = jax.random.split(rng,2)
    grad_fn = jax.value_and_grad(loss_calc,  # Function to calculate the loss
                             argnums=1,  # Parameters are second argument of the function
                             has_aux=True  # Function has additional outputs, here rng. Which you don't even need now that I think about it. 
                            )
    # Determine gradients for current model, parameters and batch
    (loss,rng), grads = grad_fn(state, state.params, rng, batch)
    # Perform parameter update with gradients and optimizer
    state = state.apply_gradients(grads=grads)
    # Return state and any other value we might want
    return state, rng, loss

Training is simple :)

1
2
3
4

state = model_state
for i in range(50000):
    state,rng,loss = train_step(state,rng,X)
    print("iter %d patience %d loss %f"%(i,patience, loss) , end="\r")

And so is sampling:

1
2
3
4
5

layers = ["random sample"] + [model.flows[i].layer_type for i in range(len(model.flows))][::-1] # names of generation layers (backward order of training layers)

for i,out in enumerate(mid_layers):
    sns.scatterplot(x= out[:,0], y = out[:,1],label = "out of" + layers[i])
    plt.show()