Package: jax-flows
Normalizing Flows in JAXImplementations of normalizing flows (RealNVP, Glow, MAF) in theJAXdeep learning framework.What are normalizing flows?Normalizing flow modelsaregenerative models, i.e. they infer the underlying probability distribution of an observed dataset. With that distribution we can do a number of interesting things, namely sample new realistic points and query probability densities.Why JAX?A few reasons!JAX encourages a functional style. When writing a layer, I didn't want people to worry about PyTorch or TensorFlow boilerplate and how their code has to fit into "the system" (e.g. do I have to keep track ofself.traininghere?)Allyou have to worry about is writing a vanilla python function which, given an ndarray, returns the correct set of outputs. You could develop your own layers with effectively no knowledge of the encompassing framework.JAX'srandom number generation systemplaces reproducibility first. To get a sense for this, when you start to parallelize a system, centralized state-based models for PRNG a latorch.manual_seed()ortf.random.set_seed()start to yield inconsistent results. Given that randomness is such a central component to work in this area, I thought that uncompromising reproducibility would be a nice feature.JAX has a really flexible automatic differentiation system. So flexible, in fact, that you can (basically) write arbitrary python functions (including for loops, if statements, etc.) and automatically compute their jacobian with a call tojax.jacfwd. So, in theory, you could write a normalizing flow layer and automatically compute its jacobian's log determinant without having to do so manually (although we're not quite there yet).How do things work?Here's an introduction! But for a more comprehensive description, check out thedocumentation.BijectionsAbijectionis a parameterized invertible function.init_fun=flows.InvertibleLinear()params,direct_fun,inverse_fun=init_fun(rng,input_dim=5)# Transform inputstransformed_inputs,log_det_jacobian_direct=direct_fun(params,inputs)# Reconstruct original inputsreconstructed_inputs,log_det_jacobian_inverse=inverse_fun(params,transformed_inputs)assertnp.array_equal(inputs,reconstructed_inputs)We can construct a sequence of bijections usingflows.Serial. The result is just another bijection, and adheres to the exact same interface.init_fun=flows.Serial(flows.AffineCoupling()flows.InvertibleLinear(),flows.ActNorm(),)params,direct_fun,inverse_fun=init_fun(rng,input_dim=5)DistributionsAdistributionis characterized by a probability density querying function, a sampling function, and its parameters.init_fun=flows.Normal()params,log_pdf,sample=init_fun(rng,input_dim=5)# Query probability density of pointslog_pdfs=log_pdf(params,inputs)# Draw new pointssamples=sample(rng,params,num_samples)Normalizing Flow ModelsUnder this definition, a normalizing flow model is just adistribution. But to retrieve one, we have to give it abijectionand anotherdistributionto act as a prior.bijection=flows.Serial(flows.AffineCoupling(),flows.InvertibleLinear(),flows.ActNorm()flows.AffineCoupling(),flows.InvertibleLinear(),flows.ActNorm())prior=flows.Normal()init_fun=flows.Flow(bijection,prior)params,log_pdf,sample=init_fun(rng,input_dim=5)How do I train a model?The same as you always would in JAX! First, define an appropriate loss function and parameter update step.defloss(params,inputs):return-log_pdf(params,inputs).mean()@jitdefstep(i,opt_state,inputs):params=get_params(opt_state)gradient=grad(loss)(params,inputs)returnopt_update(i,gradient,opt_state)Then execute a standard JAX training loop.batch_size=32itercount=itertools.count()forepochinrange(num_epochs):npr.shuffle(X)forbatch_indexinrange(0,len(X),batch_size):opt_state=step(next(itercount),opt_state,X[batch_index:batch_index+batch_size])optimized_params=get_params(opt_state)Now that we have our trained model parameters, we can query and sample as regular.log_pdfs=log_pdf(optimized_params,inputs)samples=sample(rng,optimized_params,num_samples)Magic!Interested in contributing?Yay! Check out ourcontributing guidelines.InspirationThis repository is largely modeled after thepytorch-flowsrepository byIlya Kostrikov, thenf-jaxrepository byEric Jang, and thenormalizing-flowsrepository byTony Duan.The implementations are modeled after the work of the following papers:NICE: Non-linear Independent Components EstimationLaurent Dinh, David Krueger, Yoshua BengioarXiv:1410.8516Density estimation using Real NVPLaurent Dinh, Jascha Sohl-Dickstein, Samy BengioarXiv:1605.08803Improving Variational Inference with Inverse Autoregressive FlowDiederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max WellingarXiv:1606.04934Glow: Generative Flow with Invertible 1x1 ConvolutionsDiederik P. Kingma, Prafulla DhariwalarXiv:1807.03039Flow++: Improving Flow-Based Generative Models
with Variational Dequantization and Architecture DesignJonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, Pieter AbbeelOpenReview:Hyg74h05tXMasked Autoregressive Flow for Density EstimationGeorge Papamakarios, Theo Pavlakou, Iain MurrayarXiv:1705.07057Neural Spline FlowsConor Durkan, Artur Bekasov, Iain Murray, George PapamakariosarXiv:1906.04032And by association the following surveys:Normalizing Flows: An Introduction and Review of Current MethodsIvan Kobyzev, Simon Prince, Marcus A. BrubakerarXiv:1908.09257Normalizing Flows for Probabilistic Modeling and InferenceGeorge Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, Balaji LakshminarayananarXiv:1912.02762