Reproducing proteingym

[sohviluukkonen]

Hi!

I'm trying to run your proteingym pipeline, and I'm having a couple of issues. I'm using the proteingym_branch and your ProtMamba_long_fondation model.

Firstly, in proteingym.py you are supposed to import prepare_target from tests/proteingym/utils.py but that function doesn't exist in there, so I'm currently importing it from ProtMamba_ssm/utils.py. Is that ok, or are this functions supposed to be different?

Secondly, in ProtMamba_ssm/modules.py it seems that you have updated the mamba version as you have changed
from mamba_ssm.modules.mamba_simple import Block from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn
to
from mamba_ssm.modules.block import Block from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn.
I'm still using the old Mamba version. Do you think this is problematic?

With these two (minor) changes I'm getting very different results from the ones reported in the paper with an average Spearman correlation of ~0.10. Here are Spearman values per protein that I find: ProtMamba_long_foundation_proteingym_msalength_200_spearman.json
Any idea where this huge difference could come from?

Thanks for your help!

[CyrilMa]

Hi @sohviluukkonen,

I have updated the code to make it more clear (https://github.com/Bitbol-Lab/ProtMamba-ssm/blob/proteingym_branch/tests/proteingym/proteingym.py). This should solve a lot of issues. In particular, one problem that could happen is if you don't have checkpoint_mixer = False when you load the model.

Otherwise, it's not clear for me if using the old version or not is problematic, maybe we could investigate that if your problem is not solved later.

Hope this will help!