Better NN architectures

[manuelgloeckler]

🚀 Feature Request

Unlike NPE, score and flow matching approaches are highly sensitive to the neural network architecture used to estimate the corresponding vector field. It is well-known that diffusion models, for example, tend to work best with specialized architectures like U-Nets for images (and other similarly specialized structures).

Describe the solution you'd like

To address this, the following steps should be completed:

• [ ] Identify suitable neural network architectures that could improve performance for score and flow matching.
• [ ] Implement these architectures and their corresponding network builders (see here).
• [ ] Test whether these new architectures outperform the default MLPs/ResNets. To do so, use the newly introduced mini-sbibm feature (#1325) with pytest --bm or pytest --bm -n for multiprocessing (if you have a capable CPU).

Additional Context

The aim is to enhance the performance of score and flow matching by incorporating more suitable neural network architectures, providing better results than the default MLPs/ResNets currently used.

[jaivardhankapoor]

I'd like to tackle this issue. How can I assign myself?

[manuelgloeckler]

Assigned you :)

[janfb]

I'd like to tackle this issue. How can I assign myself?

@jaivardhankapoor you should have received an invitation to become external collaborator with Triage rights. This should give you the rights to assign yourself. Can you please check whether this works, e.g., unassigning and re-assigning yourself?

[jaivardhankapoor]

I'd like to tackle this issue. How can I assign myself?

@jaivardhankapoor you should have received an invitation to become external collaborator with Triage rights. This should give you the rights to assign yourself. Can you please check whether this works, e.g., unassigning and re-assigning yourself?

Yes, I got access now. Thanks :)