Add sample weight support to change loss aggregation

[han-ol]

Adresses #314.

I'll add a first commit to aid discussion.

sample_weights is a new optional adapter output and is propagated through the compute metrics methods of ContinuousApproximator, InferenceNetwork(s) including FlowMatching, FreeFormFlow and CouplingFlow.

[han-ol]

As this affects all networks and is a conversation starter at this early stage, I did not do comprehensive tests to confirm that everything works as intended.

When we revamp the integration tests for the generative networks in the near future, it should be no problem to cover usage with and without sample weights.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

Attention: Patch coverage is 78.26087% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...flow/experimental/free_form_flow/free_form_flow.py	25.00%	3 Missing ⚠️
bayesflow/networks/flow_matching/flow_matching.py	0.00%	2 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Files with missing lines	Coverage Δ
bayesflow/approximators/continuous_approximator.py	69.67% <100.00%> (+0.50%)	⬆️
bayesflow/networks/coupling_flow/coupling_flow.py	100.00% <100.00%> (ø)
bayesflow/networks/inference_network.py	90.69% <100.00%> (+1.22%)	⬆️
bayesflow/networks/point_inference_network.py	79.38% <100.00%> (-8.25%)	⬇️
bayesflow/networks/flow_matching/flow_matching.py	82.40% <0.00%> (ø)
...flow/experimental/free_form_flow/free_form_flow.py	73.87% <25.00%> (-0.68%)	⬇️

... and 11 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[LarsKue]

keras.Metric has support for sample weights. I feel like we should leverage this rather than trying to reinvent the wheel here. @han-ol can you check if there is a feasible way to implement your desired functionality using keras.Metric?

[han-ol]

True, sample weights are also supported by keras already, but I think we still have to do some changes to fully support them.

As far as I can tell, the proposed changes are not reimplementing anything that keras.Metrics already provides, because it concerns how sample weights get propagated through the approximator, including to any keras.Metrics that should be weighted.

When fitting with a PyDataset, keras accepts sample weights as a third tuple item (x,y,sample_weight).
So it is my impression that adding sample_weight next to inference_variable, inference_condition and summary condition is in line with both the keras and the bayesflow way.

We should mimic keras behavior regarding metrics insofar as to only pass weights to metrics in the self.weighted_metrics attribute. For reference: https://keras.io/api/models/model_training_apis/#compile-method

@LarsKue, do you have a direction in mind to further exploit the preexisting keras implementation of sample_weight?

[LarsKue]

@han-ol Thank you for clarifying. I think I see the idea behind the changes now.

Since you mentioned that the changes are only a conversation starter for now, could you illustrate some use cases (w.r.t. user code) and perhaps also let us know when you think these changes are production-ready?

@stefanradev93 What do you think about the current state of this PR?

[stefanradev93]

I think it's a nice to have option, especially in terms of potential future adaptations.

[han-ol]

I made some experiments with weighting already, so I can definitely show it in use if you are interested. However, I personally cannot provide an example where weighting gives a tangible benefit over other approaches yet.

Having the infrastructure already and agreeing on UI here sooner rather than later is good, IMO, because this is straightforward to include now, and can then find its way into the docs, tests, etc. at a time when we are (re)writing old and new docs anyways.

Abstractly speaking, this sort of sample weights is useful

• whenever we have a simulator that generates samples from a proposal distribution and importance weights wrt a target distribution,
• for weighted scoring rules (when you care mostly about some part of parameter space)

Regarding timeline:
If there is no objection to the general implementation approach, we can merge this in the next days. Prior to merging I would like to have

• approval of the people who implemented each of the changed generative network loss functions
• your opinion on how to aggregate: the PR currently adds to each inference network its own aggregate() method (just once due to inheritance) that accepts losses and optional weights, then reduces with keras.ops.mean(). On the one hand, this is nice because it let's users overwrite just the aggregate method if they have special sample_weights. On the other hand, it is simpler to have one aggregate function in bayesflow/utils. I would probably change it to the second (simpler) solution here.

Another todo for the functionality is

• tests

Please let me know if I should add tests prior to merging, or if it makes more sense to add tests during or after refactor of the pre-existing network and integration tests.

[paul-buerkner]

I can confirm these are highly relevant usecases for weighting and we will need them sooner rather than later. So I would be very happy to see this PR being merged soon!

[stefanradev93]

@LarsKue @paul-buerkner Should we integrate this functionality?

[paul-buerkner]

yes I think so. I think it could be useful among others for Philipps project Stefan Radev ***@***.***> schrieb am Mi., 12. März 2025, 16:22:

…

@LarsKue <https://github.com/LarsKue> @paul-buerkner <https://github.com/paul-buerkner> Should we integrate this functionality? — Reply to this email directly, view it on GitHub <#316 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ADCW2AADNMD5UMJAVRMMSWL2UA7KTAVCNFSM6AAAAABXCEOOHOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOMJYGA3DOMBWGE> . You are receiving this because you were mentioned.Message ID: ***@***.***> [image: stefanradev93]*stefanradev93* left a comment (bayesflow-org/bayesflow#316) <#316 (comment)> @LarsKue <https://github.com/LarsKue> @paul-buerkner <https://github.com/paul-buerkner> Should we integrate this functionality? — Reply to this email directly, view it on GitHub <#316 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ADCW2AADNMD5UMJAVRMMSWL2UA7KTAVCNFSM6AAAAABXCEOOHOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOMJYGA3DOMBWGE> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[stefanradev93]

Ok, let's merge dev into this branch and I will accept the merge.

[paul-buerkner]

Merging done. Quick question: Is "sample_weights" the best name for it? We could also just use "weights" right? Or would this be ambigous?

[LarsKue]

@paul-buerkner I think both would be fine. sample_weight (without the s) is what keras calls it.

[paul-buerkner]

Let's use singular then. Because sample_weight will always be only a single variable, right (data dim = 1)? Not multiple variables as is the case for inference_variables etc. @han-ol if you agree, would you mind changing that?

[han-ol]

Yes, agreed.

I changed the names "sample_weights" -> "sample_weight" everywhere and also changed the default adapter to expect just one string that is renamed to "sample_weight" (rather than a Sequence of multiple, that is concatenated into "sample_weights").

I added tests as well and support for point inference networks, so the PR is ready to undergo review at this point.

[Copilot]

Copilot reviewed 13 out of 13 changed files in this pull request and generated 1 comment.