Flexible Default Settings Across Inference Networks

[stefanradev93]

Adding this WIP in anticipation of the integrators PR to discuss dynamically setting default settings for subnets across different inference networks. The current straightforward proposal makes use of a default class dict for each inference net like:

MLP_DEFAULT_CONFIG = {
    "widths": (256, 256, 256),
    "activation": "mish",
    "kernel_initializer": "he_normal",
    "residual": True,
    "dropout": 0.05,
    "spectral_normalization": False,
}

The default config is only considered if the default subnet is an mlp and not a user-provided backbone. In the former case, the fields of the default config are updated with user-provided fields for subnet_kwargs (if any) and the inference network is spawned. A working draft can be inspected, for instance, in free_form_flow. Let me know what you think @paul-buerkner @LarsKue

[paul-buerkner]

I think this proposal looks sensible to me. We should definitely (a) have very good defaults that (b) can be selectively updated by the user while retaining all other default settings. I think your proposal achieves both in a very simple manner.

[stefanradev93]

This is ready to merge upon your review @LarsKue / @paul-buerkner.

[stefanradev93]

@LarsKue Interesting, the cycle consistency test seems to be failing for flow matching, even though it's fitting a couple of benchmarks...The reason seems to be a numerical one. The output values are close, but not within the tolerance threshold. I suspect the reason is that the defaults are now creating a significantly larger network. Can we increase the precision of the cycle consistency test for FM?

[paul-buerkner]

I cannot judge the details of the defaults themselves (for this I trust you, @stefanradev93, as the neural network whisperer). But the general structure of updating default configs looks good to me. I approve but will wait for @LarsKue to check the test.

[LarsKue]

The default for Flow Matching is now adaptive step size, which can mess with cycle consistency. I am already looking into improved test parametrization, which could allow us to circumvent this issue by using fixed step size or increasing the tolerance for adaptive step size only.

[stefanradev93]

An even simpler solution would be to use a tiny network in the tests.

[LarsKue]

An even simpler solution would be to use a tiny network in the tests.

I don't think this necessarily solves the issue with the cycle consistency, although it is of course desirable with respect to speeding the tests up.

The problem is that since the step size is adaptive, we necessarily end up evaluating any curved velocity field at different time points when integrating forward than backward.

[LarsKue]

I would also like to consider using a general NetworkBuilder or NetworkFactory pattern instead of directly using configs, to allow users to create networks for their use case. I think this would also tie in nicely with coupling flows, which are currently quite limited in user-friendly customization. We could still use these defaults in the builder pattern, of course.

[stefanradev93]

Yes, the factory would be nice to have, and will also solve the get_/from_config shuffling. For now, let's use fixed step size / smaller networks for the tests or relax the cycle consistency constraints for FM.

[codecov-commenter]

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

Codecov Report

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

Files with missing lines	Patch %	Lines
...w/networks/consistency_models/consistency_model.py	16.66%	5 Missing ⚠️

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

Files with missing lines	Coverage Δ
...etworks/coupling_flow/couplings/single_coupling.py	93.22% <100.00%> (+0.62%)	⬆️
bayesflow/networks/flow_matching/flow_matching.py	82.56% <100.00%> (+0.49%)	⬆️
...ayesflow/networks/free_form_flow/free_form_flow.py	74.54% <100.00%> (+2.54%)	⬆️
bayesflow/networks/inference_network.py	89.47% <100.00%> (+0.28%)	⬆️
...w/networks/consistency_models/consistency_model.py	18.11% <16.66%> (+0.07%)	⬆️

... and 2 files with indirect coverage changes