search

Bitbol-Lab / ProtMamba-ssm

ProtMamba: a homology-aware but alignment-free protein state space model
https://www.biorxiv.org/content/10.1101/2024.05.24.595730v1
Apache License 2.0
49 stars 7 forks source link

Reproducing proteingym #13

Open sohviluukkonen opened 1 month ago

sohviluukkonen commented 1 month ago

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 commented 1 month ago

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!