Amortized point estimators

[marvinschmitt]

The paper "Likelihood-free parameter estimation with neural Bayes estimators" (Sainsbury-Dale, Zammit-Mangion, & Huser, 2023) enables neural amortized point estimation, which is generally faster than fully Bayesian neural inference with conditional generative models (current standard in BayesFlow).

Implementing amortized point estimators in BayesFlow would help with fast prototyping and sanity checks (and surely other tasks, too!).

Some BayesFlow pointers/proposals:

• Implement the point estimation features in a class AmortizedPointEstimator in bayesflow.amortizers
• Implement appropriate loss functions in bayesflow.losses (depending on the target quantity)
• Implement a suitable inference network in bayesflow.inference_networks which takes the summary network output (e.g., DeepSet outputs) and regresses to the point estimate
• Summary networks are already implemented in bayesflow.summary_networks and probably work out-of-the-box

Some links:

• Paper: https://doi.org/10.1080/00031305.2023.2249522
• NeuralEstimators package (R): https://github.com/msainsburydale/NeuralEstimators
• NeuralEstimators.jl package (Julia): https://github.com/msainsburydale/NeuralEstimators.jl

[Jice-Zeng]

Great proposals! I read the paper. But I am not sure if the method can generate posterior samples, which may be helpful to detect multi-modal distributions.

[stefanradev93]

Yes, the method will not deal with multimodality out of the box. I think its main use case will be as a fast and frugal heuristic to obtain estimates very quickly.

[paul-buerkner]

Idea: can we train the summary nets based on point inference nets and then reuse the trained summary nets in combination with full posterior inference nets? Stefan Radev ***@***.***> schrieb am Di., 13. Feb. 2024, 12:43:

…

Yes, the method will not deal with multimodality out of the box. I think its main use case will be as a fast and frugal heuristic to obtain estimates very quickly. — Reply to this email directly, view it on GitHub <#121 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ADCW2AH5AXV6TW2ZYGIOC4DYTLHMNAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBQGMZDIOJQG4> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[andrewzm]

Multi-modality on the marginals should be detectable through the posterior quantiles. The main "compromise" is that there is no easy way to get information on posterior joint dependence generally. But as Stefan said, this is more of an "exploratory" step or an alternative for when the full NF framework is proving difficult to fit. I like Paul's idea. Possible issue: Summaries that are good for estimating individual parameters might not be relevant/suitable for extracting information on posterior correlations (additional summaries might be needed when doing full posterior inference).

…

On Tue, Feb 13, 2024 at 3:00 PM Paul-Christian Bürkner < ***@***.***> wrote: Idea: can we train the summary nets based on point inference nets and then reuse the trained summary nets in combination with full posterior inference nets? Stefan Radev ***@***.***> schrieb am Di., 13. Feb. 2024, 12:43: > Yes, the method will not deal with multimodality out of the box. I think > its main use case will be as a fast and frugal heuristic to obtain > estimates very quickly. > > — > Reply to this email directly, view it on GitHub > < #121 (comment)>, > or unsubscribe > < https://github.com/notifications/unsubscribe-auth/ADCW2AH5AXV6TW2ZYGIOC4DYTLHMNAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBQGMZDIOJQG4> > . > You are receiving this because you are subscribed to this thread.Message > ID: ***@***.***> > — Reply to this email directly, view it on GitHub <#121 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABEU3LL2TL35DSD6ZYOV63YTLQMHAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBQGM4DONRXGA> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[stefanradev93]

I like the idea too! We could use the point summary network or the point estimates themselves as summaries, with an auxiliary, smaller summary stream for capturing additional information relevant for the joint.

[andrewzm]

I think using the point estimates as summary statistics will be similar to when one uses expert summary statistics, which we know can work well (e.g., the Jiang et al. Statistica Sinica paper) but probably what you have at the moment is even better. To start off with, from what I know of BayesFlow, I think we should keep the statistics network as is, and add an option for the user to choose "point estimation" instead of "full inference" -- technically all is required is to swap out the invertible NN with a much simpler dense neural network with input dimension equal to the summary statistic dimension and output dimension equal to the parameter dimension. There could even be a sensible default for the architecture of this NN. The user also needs to be able to specify a loss function when choosing "point estimation", but we could just start off with squared error loss at first just to get it running.

…

On Wed, Feb 14, 2024 at 12:27 AM Stefan Radev ***@***.***> wrote: I like the idea! We could use the point summary network or the point estimates themselves as summaries, with an auxiliary, smaller summary stream for capturing additional information relevant for the joint. — Reply to this email directly, view it on GitHub <#121 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABEU3OVKIAMWU7RDJTACD3YTNS3JAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBRGUYTGNZUGM> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[paul-buerkner]

yes I fully agree. my idea was going in a different direction though but it would be more like research which we can do once the implementation is done Andrew Zammit Mangion ***@***.***> schrieb am Mi., 14. Feb. 2024, 16:10:

…

I think using the point estimates as summary statistics will be similar to when one uses expert summary statistics, which we know can work well (e.g., the Jiang et al. Statistica Sinica paper) but probably what you have at the moment is even better. To start off with, from what I know of BayesFlow, I think we should keep the statistics network as is, and add an option for the user to choose "point estimation" instead of "full inference" -- technically all is required is to swap out the invertible NN with a much simpler dense neural network with input dimension equal to the summary statistic dimension and output dimension equal to the parameter dimension. There could even be a sensible default for the architecture of this NN. The user also needs to be able to specify a loss function when choosing "point estimation", but we could just start off with squared error loss at first just to get it running. On Wed, Feb 14, 2024 at 12:27 AM Stefan Radev ***@***.***> wrote: > I like the idea! We could use the point summary network or the point > estimates themselves as summaries, with an auxiliary, smaller summary > stream for capturing additional information relevant for the joint. > > — > Reply to this email directly, view it on GitHub > < #121 (comment)>, > or unsubscribe > < https://github.com/notifications/unsubscribe-auth/AABEU3OVKIAMWU7RDJTACD3YTNS3JAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBRGUYTGNZUGM> > . > You are receiving this because you are subscribed to this thread.Message > ID: ***@***.***> > — Reply to this email directly, view it on GitHub <#121 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ADCW2AAU4KQIMG6AP4GYYPTYTRIOFAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBTGE2DEMRWGU> . You are receiving this because you commented.Message ID: ***@***.***>

[paul-buerkner]

and I did not mean using the point estimates as summary statistics. I agree that would not be very useful in general. Paul Buerkner ***@***.***> schrieb am Mi., 14. Feb. 2024, 16:51:

…

yes I fully agree. my idea was going in a different direction though but it would be more like research which we can do once the implementation is done Andrew Zammit Mangion ***@***.***> schrieb am Mi., 14. Feb. 2024, 16:10: > I think using the point estimates as summary statistics will be similar > to > when one uses expert summary statistics, which we know can work well > (e.g., > the Jiang et al. Statistica Sinica paper) but probably what you have at > the > moment is even better. To start off with, from what I know of BayesFlow, > I > think we should keep the statistics network as is, and add an option for > the user to choose "point estimation" instead of "full inference" -- > technically all is required is to swap out the invertible NN with a much > simpler dense neural network with input dimension equal to the summary > statistic dimension and output dimension equal to the parameter > dimension. > There could even be a sensible default for the architecture of this NN. > The > user also needs to be able to specify a loss function when choosing > "point > estimation", but we could just start off with squared error loss at first > just to get it running. > > On Wed, Feb 14, 2024 at 12:27 AM Stefan Radev ***@***.***> > wrote: > > > I like the idea! We could use the point summary network or the point > > estimates themselves as summaries, with an auxiliary, smaller summary > > stream for capturing additional information relevant for the joint. > > > > — > > Reply to this email directly, view it on GitHub > > < > #121 (comment)>, > > > or unsubscribe > > < > https://github.com/notifications/unsubscribe-auth/AABEU3OVKIAMWU7RDJTACD3YTNS3JAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBRGUYTGNZUGM> > > > . > > You are receiving this because you are subscribed to this > thread.Message > > ID: ***@***.***> > > > > — > Reply to this email directly, view it on GitHub > <#121 (comment)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/ADCW2AAU4KQIMG6AP4GYYPTYTRIOFAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBTGE2DEMRWGU> > . > You are receiving this because you commented.Message ID: > ***@***.***> >

[ukoethe]

I also think we need to do it in two steps:

• train summary statistics jointly with a normalizing flow as usual
• train point estimates (with an additional small network) from the frozen summary network output

Thus, only inference will potentially be faster, not training. I suspect that attemtps to train summary statistics directly via point estimates might result in poor summaries, because it will be harder for the summary network to capture the uncertainty.

[ukoethe]

A luxury method to train point estimates would enhance the point estimator network with additional outputs for the (co)variance of the point estimate and the fraction of the full posterior explained by the MAP mode.

These outputs could be trained by destillation from the SBI posterior: (1) Find the MAP mode in the NF posterior output => target value of the MAP point estimate. (2) Fit a Gaussian to the MAP mode (=Laplace approximation) to get the (co)variance and fraction of probability mass explained by the MAP mode.

The destillation approach has the additional benefit that the point estimator would return the true MAP estimates (subject to the accuracy of the normalizing flow) instead of the ground truth parameterizations of the training runs.

[stefanradev93]

Summarizing the ideas so far:

1. Train point estimators instead of generative networks - simply attaches a configurable MLP to the summary network and uses some norm (L1, L2, L-Infinity, https://www.tensorflow.org/api_docs/python/tf/norm) as a loss.

• Advantage: fast and pleasant to iterate over for larger problems.
• Disadvantage: no full Bayesian inference (by design).

2. Train point estimators and use pre-trained summary network in tandem with a full generative network - uses a frozen summary network, possibly in combination with another, smaller and trainable summary network, to train a full workflow.

• Advantage: the first summary network is trained fast and already captures some important information. May be equivalent to good hand-crafted summary statistics.
• Disadvantage: summary statistics may not be good, speed gains diminished due to two-step approach.

3. Train the default approach and distil it using a point estimator - trains the full workflow (i.e., summary and generative net) and attaches a configurable MLP to the pretrained summary network, possibly learning a summary of the approximate posterior (e.g., MAP).

• Advantage: the summary statistics will already be good for learning the full posterior (assuming everything has converged just fine), so they will also be good for the point estimator. Point estimator can then be trained hyperefficiently.
• Disadvantage: speed gains diminished due to two-step approach, the utility of the point estimates is unclear if full workflow is already good.

Crucially, all three are enabled, once the functionality is there. Does anyone want to implement the blueprints Marvin suggested in the first post? Otherwise, I would self-assign in a week or two.

[andrewzm]

Looks good, I think 1. is where we should start as that will help establish the core benefits, and then move from there. Matt and I will be happy to provide comments and check changes/provide suggestions as they come through; can also contribute to code but probably someone who knows the internals should drive it.

…

On Thu, Feb 15, 2024 at 12:24 AM Stefan Radev ***@***.***> wrote: Summarizing the ideas so far: 1. *Train point estimators instead of generative networks* - simply attaches a configurable MLP to the summary network and uses some norm (L1, L2, L-Infinity, https://www.tensorflow.org/api_docs/python/tf/norm) as a loss. - *Advantage*: fast and pleasant to iterate over for larger problems. - *Disadvantage*: no full Bayesian inference (by design). 1. *Train point estimators and use pre-trained summary network in tandem with a full generative network* - uses a frozen summary network, possibly in combination with another, smaller and trainable summary network, to train a full workflow. - *Advantage*: the first summary network is trained fast and already captures some important information. May be equivalent to good hand-crafted summary statistics. - *Disadvantage*: summary statistics may not be good, speed gains diminished due to two-step approach. 1. *Train the default approach and distil it using a point estimator* - trains the full workflow (i.e., summary and generative net) and attaches a configurable MLP to the pretrained summary network. - *Advantage*: the summary statistics will already be good for learning the full posterior (assuming everything has converged just fine), so they will also be good for the point estimator. Point estimator can then be trained hyperefficiently. - *Disadvantage*: speed gains diminished due to two-step approach, the utility of the point estimates is unclear if full workflow is already good. Crucially, all three are enabled, once the functionality is there. Does anyone want to implement the blueprints Marvin suggested in the first post? Otherwise, I would self-assign in a week or two. — Reply to this email directly, view it on GitHub <#121 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABEU3MBQXPGPA7ZOIRJAVTYTS3HJAVCNFSM6AAAAABC5CBZ3CVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNBTG43DANRQHA> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[stefanradev93]

Hi all, I have implemented and tested a first working version of an amortized point estimator, which can be found in the Development branch.

An example use case is:

import bayesflow as bf
from bayesflow.benchmarks import Benchmark

# Trivial Gaussian benchmark or a real model
benchmark = Benchmark('gaussian_linear', mode='posterior')

# An easy-to-use MLP with residual connections
inference_net = bf.helper_networks.ConfigurableMLP(input_dim=10, dropout_rate=0.05)

# Can be any norm in [1, 2, np.inf] and can also use a summary network
amortizer = bf.amortizers.AmortizedPointEstimator(inference_net, norm_ord=2)

# Training can happen as usual
trainer = bf.trainers.Trainer(amortizer=amortizer, configurator=benchmark.configurator)
data = benchmark.generative_model(5000)
h = trainer.train_offline(data, epochs=50, batch_size=32)

# Quick point estimates can be obtained by simply calling the .estimate() method
test_data = benchmark.configurator(benchmark.generative_model(100))
estimates = amortizer.estimate(test_data)

Let me know what you think.