Significant memory increase with migration events compared to momi2

[chichizhao]

Hi there!

Recently, I've been migrating my demographic inference workflow from momi2 to momi3 (JAX-based) to take advantage of its improved computational capabilities. However, I've noticed a significant increase in memory requirements — several times higher than momi2 — especially when modeling migration events between demes.

My demographic model involves 5 demes (NW, BA, EB, EA, MA) with migration events between them, analyzed across ~216M genetic sites. Under momi2, this model runs successfully with manageable memory usage. Under momi3, the memory footprint grows substantially, making it difficult to run on standard hardware.

I'm raising this issue to ask:

Is the increased memory usage with migration events (particularly continuous migration) expected behavior in momi3 compared to momi2?
Are there recommended strategies or configurations for handling multi-deme models with migration events while keeping memory manageable?
Is there ongoing work on sparse tensor support in JAX that might help optimize memory for continuous migration scenarios?
Thank you!

[terhorst]

Hi @chichizhao , thanks for the report. Could you clarify for me how you're running momi2 and momi3? momi2 does not support continuous migration so for an even comparison you'd want to use a demography consisting of pulse migrations only. In that case I would expect memory usage to be similar, not exactly the same as the backend for momi3 is completely different (rewritten in Jax).

For continuous migration, it requires solving a system of ODEs, and for large nbrs of demes unfortunately memory usage can be quite substantial, particularly if you are computing gradients. We are working on an approximate method for analyzing large demographies with a lot of migration; stay tuned! :-)

Also, development of momi3 has been folded into jthlab/demestats see https://demestats.readthedocs.io/en/latest/momi3/index.html, so for the most up-to-date info please check there!