Incorrect implementation of TargetConditionalFlowMatcher

[Rattko]

Hey,

I've been experimenting with Conditional Flow Matching for a paper I'm currently working on and I've noticed that you have an incorrect implementation of TargetConditionalFlowMatcher, in particular, the sample_xt and compute_conditional_flow methods. The current implementation (rewritten 1:1 in Jax) completely fails to learn even the simplest distributions.

I believe that sample_xt should use x0 instead of epsilon and compute_conditional_flow should return x1 - (1 - self.sigma) * x0 instead of what it currently returns.

Hope this helps ;)

[atong01]

Hi,

I think the current implementation here is correct. The TargetConditionalFlowMatcher explicitly does not depend on x0, as is specified in the Lipman et al. paper.

Note that this only works for initial distributions which are standard normal. Perhaps this is why it is not working in your case?

Thanks,
Alex

[Rattko]

Hmm, my prior distribution is the Standard normal and yet I couldn't make the training work with your implementation. I'll try to create a minimal working example and get back to you

[atong01]

Cool thanks! Let me know