Package: bayes-vi
Get StartedInstall the packagehttps://pypi.org/project/bayes-vi/using pip:pipinstallbayes-viCheck out the documentation athttps://maxgrtz.github.io/bayesian-inference/!Importsimportcollectionsimportnumpyasnpimportpandasaspdimportmatplotlib.pyplotaspltimportarvizasaz# use ggplot styles for graphsplt.style.use('ggplot')importtensorflowastfimporttensorflow_probabilityastfptfd=tfp.distributionstfb=tfp.bijectorstfk=tf.keras# set tf logger to log level ERROR to avoid warningstf.get_logger().setLevel('ERROR')# import probabilistic modelsfrombayes_vi.modelimportModel# import utilsfrombayes_vi.utilsimportto_ordered_dictfrombayes_vi.utils.datasetsimportmake_dataset_from_dffrombayes_vi.utils.symplectic_integratorsimportLeapfrogIntegrator# mcmc importsfrombayes_vi.inference.mcmcimportMCMCfrombayes_vi.inference.mcmc.transition_kernelsimportHamiltonianMonteCarlo,NoUTurnSampler,RandomWalkMetropolisfrombayes_vi.inference.mcmc.stepsize_adaptation_kernelsimportSimpleStepSizeAdaptation,DualAveragingStepSizeAdaptation# vi importsfrombayes_vi.inference.viimportVIfrombayes_vi.inference.vi.surrogate_posteriorsimportADVI,NormalizingFlowfrombayes_vi.inference.vi.flow_bijectorsimportHamiltonianFlow,AffineFlow,make_energy_fn,make_scale_fn,make_shift_fn1. Example WorkflowA Bayesian modeling worklow using this package could look like this.1.1. Create a DatasetCreate atf.data.Datasetfrom your data. For this example consider some univariate Gaussian fake data.num_datapoints=20loc=7scale=3true_dist=tfd.Normal(loc=loc,scale=scale)y=true_dist.sample(num_datapoints)The package provides a utility functionmake_dataset_from_dffor easy construction of a dataset from apd.DataFrame.data=pd.DataFrame(dict(y=y))dataset=make_dataset_from_df(data,target_names=['y'],format_features_as='dict')ax=data.plot(kind='kde',title='Univariate Gaussian Test Data')1.2. Define a Bayesian ModelTo define a BayesianModel, simply provide ancollections.OrderedDictof prior distributions astfp.distribution.Distirbutions and a likelihood function, which takes the parameters as inputs by name.priors=collections.OrderedDict(loc=tfd.Normal(loc=0.,scale=10.),scale=tfd.HalfNormal(scale=10.))deflikelihood(loc,scale):returntfd.Normal(loc=loc,scale=scale)model=Model(priors=priors,likelihood=likelihood)It is possible to provide a list of constraining bijectors on construction of the model, which allows for computations on an unconstrained space.If no such list is defined, default bijectors for each parameter are chosen.1.3. Run Inference Algorithm of ChoiceWith this simple setup, you can now use either MCMC or variational inference to do Bayesian inference.1.3.1. Run MCMC (e.g. with No-U-Turn Sampler with dual averaging stepsize adaptation)Note that the transition kernel may have various additional parameters. Here we use just the defaults and only provide the stepsize adaptation kernel.NUM_CHAINS=2NUM_SAMPLES=4000NUM_BURNIN_STEPS=1000stepsize_adaptation_kernel=DualAveragingStepSizeAdaptation(num_adaptation_steps=int(NUM_BURNIN_STEPS*0.8))kernel=NoUTurnSampler(stepsize_adaptation_kernel=stepsize_adaptation_kernel)mcmc=MCMC(model=model,dataset=dataset,transition_kernel=kernel)mcmc_result=mcmc.fit(num_chains=NUM_CHAINS,num_samples=NUM_SAMPLES,num_burnin_steps=NUM_BURNIN_STEPS,progress_bar=True,)99.08% [5000/5000 00:15<00:00]Themcmc_resultis a wrapper containing the sample resultsmcmc_result.samplesand additional traced metrics specific to the kernelmcmc_result.trace.To visualize the results, one can use existing libraries like arviz (https://arviz-devs.github.io/arviz/index.html).Later versions of this package should allow for a seamless integration with arviz, but for now we have to reformat the sample results to conform to arviz conventions.# adapt sample results to arviz required format --- later versions should allow for seamless integration with arvizposterior_samples=to_ordered_dict(model.param_names,tf.nest.map_structure(lambdav:v.numpy().T,mcmc_result.samples))az.summary(posterior_samples)meansdhdi_3%hdi_97%mcse_meanmcse_sdess_meaness_sdess_bulkess_tailr_hatloc6.8370.5465.8187.8530.0080.0055210.05210.05240.04550.01.0scale2.4440.4261.7083.2160.0060.0054639.04273.05315.04424.01.0ax=az.plot_trace(posterior_samples)1.3.2. Run Variational Inference (e.g. meanfield ADVI)Instead of MCMC methods, we might want to use variational inference.Choose a surrogate posterior, e.g. meanfieldADVI, a discrepancy function fromtfp.vi, e.g.tfp.vi.kl_reverse, and select some optimization parameters (number of optimzation steps, sample size to approximate variational loss, optimizer, learning rate).NUM_STEPS=1000SAMPLE_SIZE=10LEARNING_RATE=1e-2surrogate_posterior=ADVI(model=model,mean_field=True)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,meanfield_advi_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 00:02<00:00 avg loss: 50.336]ax=plt.plot(meanfield_advi_losses)df=pd.DataFrame(approx_posterior.sample(5000))df.describe(percentiles=[0.03,0.5,0.97]).Tcountmeanstdmin3%50%97%maxloc5000.06.8196530.5112644.7102795.8445156.8253407.7931838.429766scale5000.02.4204030.3674141.1745331.7604952.4210833.1177243.836521ax=df.plot(kind='kde')2. Normalizing Flows for Variational InferenceContinuing the simple example above, this is how you can flexibly define surrogate posteriors based on Normalizing flows for variational Bayesian inference.By default, the base distribution is a standard diagonal Gaussian.This example uses the predefined trainableAffineFlowbijector, but any parameterizedtfp.bijectors.Bijectorcan be used. Just define your own!
You can construct more expressive flows by making use of thetfp.bijectors.Chainbijector to define compositions of flows.ndims=model.flat_unconstrained_param_event_ndims# this is the number of dimensions of the unconstrained parameter space2.1. Affine Flowsflow_bijector=AffineFlow(ndims)surrogate_posterior=NormalizingFlow(model=model,flow_bijector=flow_bijector)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,affine_flow_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 00:02<00:00 avg loss: 50.336]2.2. Continuous FlowsTo define continuous normalizing flows, you may use is thetfp.bijectors.FFJORDbijector.defget_continuous_flow_bijector(unconstrained_event_dims):state_fn=tfk.Sequential()state_fn.add(tfk.layers.Dense(64,activation=tfk.activations.tanh))state_fn.add(tfk.layers.Dense(unconstrained_event_dims))state_fn.build((None,unconstrained_event_dims+1))state_time_derivative_fn=lambdat,state:state_fn(tf.concat([tf.fill((state.shape[0],1),t),state],axis=-1))returntfb.FFJORD(state_time_derivative_fn,ode_solve_fn=tfp.math.ode.DormandPrince(first_step_size=0.1).solve,trace_augmentation_fn=tfb.ffjord.trace_jacobian_hutchinson)flow_bijector=get_continuous_flow_bijector(ndims)surrogate_posterior=NormalizingFlow(model=model,flow_bijector=flow_bijector)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,cnf_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 01:02<00:00 avg loss: 50.439]2.3. ComparisonHave a look at the optimzation behaviors in comparison.ax=pd.DataFrame({'meanfield_advi':meanfield_advi_losses,'affine_flow':affine_flow_losses,'cnf':cnf_losses}).plot(ylim=(-10,500))3. Augmented Normalizing Flows for Variational InferenceThe package allows to easily define augmented normalizing flows. This obviously has no advantages for such a simple example model.Simply choose a number of extra dimensionsextra_ndimsand define the flow bijector on the augmented space, i.e. on dimensionsndims + extra_ndims.The optimization problem is solved on the augmented space by lifting the target distribution with the conditionalposterior_lift_distribution, which by default is simplylambda _: tfd.MultivariateNormalDiag(loc=tf.zeros(extra_ndims), scale_diag=tf.ones(extra_ndims)).ndims=model.flat_unconstrained_param_event_ndims# this is the number of dimensions of the unconstrained parameter spaceextra_ndims=2# dimensions of the augmentation spaceflow_bijector=AffineFlow(ndims+extra_ndims)surrogate_posterior=NormalizingFlow(model=model,flow_bijector=flow_bijector,extra_ndims=extra_ndims)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,affine_flow_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 00:04<00:00 avg loss: 50.310]defget_continuous_flow_bijector(unconstrained_event_dims):state_fn=tfk.Sequential()state_fn.add(tfk.layers.Dense(64,activation=tfk.activations.tanh))state_fn.add(tfk.layers.Dense(unconstrained_event_dims))state_fn.build((None,unconstrained_event_dims+1))state_time_derivative_fn=lambdat,state:state_fn(tf.concat([tf.fill((state.shape[0],1),t),state],axis=-1))returntfb.FFJORD(state_time_derivative_fn,ode_solve_fn=tfp.math.ode.DormandPrince(first_step_size=0.1).solve,trace_augmentation_fn=tfb.ffjord.trace_jacobian_hutchinson)flow_bijector=get_continuous_flow_bijector(ndims+extra_ndims)surrogate_posterior=NormalizingFlow(model=model,flow_bijector=flow_bijector,extra_ndims=extra_ndims)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,cnf_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 01:08<00:00 avg loss: 50.633]ax=pd.DataFrame({'meanfield_advi':meanfield_advi_losses,'affine_flow':affine_flow_losses,'cnf':cnf_losses}).plot(ylim=(-10,500))4. Hamiltonian Normalizing FlowsAs another example consider Hamiltonian Normalizing Flows (HNF) as motivated by Toth, P., Rezende, D. J., Jaegle, A., Racanière, S., Botev, A., & Higgins, I. (2019). Hamiltonian Generative Networks. ArXiv.http://arxiv.org/abs/1909.13789.HNFs are based on defining augmented flows as the Hamiltonian dynamics of a family of Hamiltonian systems.To compute such dynamics, we use the symplectic leapfrog integrator.The package provides an implementation of a general Hamiltonian Flow bijector.4.1. Verify Hamiltonian Flow BijectorTo test the implementation the Hamiltonian Flow Bijector, consider the Hamiltonian of a 1d harmonic oscillator with mass m=1 and constant k=1:H(q,p) = T(p) + U(q) = 0.5 p^2 + 0.5 q^2.We know that this should yield an oscillating dynamic, which manifests as a circle on phase space.T=lambdap:0.5*p**2U=lambdaq:0.5*q**2# event dims can be ignored in this example, because they are only used to construct trainable default energy function if none are provided explicitlyhamiltonian_flow=HamiltonianFlow(event_dims=1,potential_energy_fn=U,kinetic_energy_fn=T,symplectic_integrator=LeapfrogIntegrator(),step_sizes=0.5,num_integration_steps=1)initial_state=tf.constant([1.0,0.0])path=[initial_state.numpy()]for_inrange(20):path.append(hamiltonian_flow.forward(path[-1]).numpy())path=np.stack(path)fig,[ax1,ax2]=plt.subplots(1,2,figsize=(6,3))ax1.plot(path[:,0])ax1.set_title('configuration plot')ax2.plot(path[:,0],path[:,1])ax2.set_title('phase space plot')plt.tight_layout()4.2. Hamiltonian Normalizing FlowsIn the case of HNFs the augmentation space has the same dimension as the parameter space.ndims=model.flat_unconstrained_param_event_ndims# this is the number of dimensions of the unconstrained parameter spaceextra_ndims=ndims# dimensions of the augmentation spaceFor simplicity, we use the default MLP based potential and kinetic energy functions.flow_bijector=HamiltonianFlow(ndims,step_sizes=tf.Variable(0.1),num_integration_steps=2)surrogate_posterior=NormalizingFlow(model=model,flow_bijector=flow_bijector,extra_ndims=extra_ndims)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,hnf_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 00:07<00:00 avg loss: 54.512]ax=pd.DataFrame({'meanfield_advi':meanfield_advi_losses,'affine_flow':affine_flow_losses,'cnf':cnf_losses,'hnf':hnf_losses}).plot(ylim=(-10,500))An interesting observation is that HNFs become more stable and perform better, if the base distribution is trainable.This compensates for the problem, that volume preserving flows like HNFs can only permute the densities of the base distribution.flow_bijector=HamiltonianFlow(ndims,step_sizes=tf.Variable(0.1),num_integration_steps=2)base_distribution=tfd.MultivariateNormalDiag(loc=tf.Variable([0.0]*2*ndims),scale_diag=tf.Variable([0.1]*2*ndims))surrogate_posterior=NormalizingFlow(model=model,flow_bijector=flow_bijector,base_distribution=base_distribution,extra_ndims=extra_ndims)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,hnf_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 00:07<00:00 avg loss: 50.746]ax=pd.DataFrame({'meanfield_advi':meanfield_advi_losses,'affine_flow':affine_flow_losses,'cnf':cnf_losses,'hnf':hnf_losses}).plot(ylim=(-10,500))5. Regression Model ExampleAs a final example, this is how to define regression models.5.1. Create a DatasetAgain, create atf.data.Datasetfrom your data or, in this case, generate some fake dataset.num_datapoints=50beta0=7beta1=3scale=10x=tf.random.normal((num_datapoints,),mean=0.0,stddev=10.0)true_dist=tfd.Normal(loc=beta0+x*beta1,scale=scale)y=true_dist.sample()data=pd.DataFrame(dict(x=x,y=y))dataset=make_dataset_from_df(data,target_names=['y'],feature_names=['x'],format_features_as='dict')ax=data.plot(x='x',y='y',kind='scatter',title='Linear Regression Test Data')1.2. Define a Bayesian ModelTo define theModel, again provide ancollections.OrderedDictof prior distributions astfp.distribution.Distirbutions and a likelihood function, which takes the parameters and the features as inputs.# define priorspriors=collections.OrderedDict(beta_0=tfd.Normal(loc=0.,scale=10.),beta_1=tfd.Normal(loc=0.,scale=10.),scale=tfd.HalfNormal(scale=10.),)# define likelihooddeflikelihood(beta_0,beta_1,scale,features):linear_response=beta_0+beta_1*features['x']returntfd.Normal(loc=linear_response,scale=scale)model=Model(priors=priors,likelihood=likelihood)1.3. Run Inference Algorithm of ChoiceYou can now again use either MCMC or variational inference to do Bayesian inference.1.3.1. Run MCMC (e.g. with Random Walk Metropolis kernel)NUM_CHAINS=2NUM_SAMPLES=4000NUM_BURNIN_STEPS=1000kernel=RandomWalkMetropolis()mcmc=MCMC(model=model,dataset=dataset,transition_kernel=kernel)mcmc_result=mcmc.fit(num_chains=NUM_CHAINS,num_samples=NUM_SAMPLES,num_burnin_steps=NUM_BURNIN_STEPS,progress_bar=True,)92.70% [4635/5000 00:01<00:00]# adapt sample results to arviz required format --- later versions should allow for seamless integration with arvizposterior_samples=to_ordered_dict(model.param_names,tf.nest.map_structure(lambdav:v.numpy().T,mcmc_result.samples))az.summary(posterior_samples)meansdhdi_3%hdi_97%mcse_meanmcse_sdess_meaness_sdess_bulkess_tailr_hatbeta_06.7441.2964.5619.2640.1240.088109.0109.0110.0146.01.01beta_12.8590.1442.5703.1100.0060.004596.0591.0600.0689.01.00scale9.2170.9247.72410.9780.0630.045215.0214.0220.0395.01.01ax=az.plot_trace(posterior_samples)1.3.2. Run Variational Inference (e.g. full rank ADVI)NUM_STEPS=1000SAMPLE_SIZE=10LEARNING_RATE=1e-2surrogate_posterior=ADVI(model=model)vi=VI(model=model,dataset=dataset,surrogate_posterior=surrogate_posterior,discrepancy_fn=tfp.vi.kl_reverse)approx_posterior,meanfield_advi_losses=vi.fit(optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE),num_steps=NUM_STEPS,sample_size=SAMPLE_SIZE,progress_bar=True)100.00% [1000/1000 00:04<00:00 avg loss: 188.817]ax=plt.plot(meanfield_advi_losses)df=pd.DataFrame(approx_posterior.sample(5000))df.describe(percentiles=[0.03,0.5,0.97]).Tcountmeanstdmin3%50%97%maxbeta_05000.06.8775421.2591722.3234224.4992396.8543239.26364111.693712beta_15000.02.8636480.1367082.3824102.6097252.8617143.1209233.358475scale5000.09.0054770.8705405.7141167.3802058.99294810.63355212.097922ax=df.plot(kind='kde')