adding support for multimers via residue_index hack

[sokrypton]

No description provided.

[sokrypton]

I'm not sure if I implemented RoPE() residue index offset correctly.

Setting residue index offset for relative-positional embedding works, but if you touch RoPE, it crashes!

Here is what I tried for RoPE():

# default
position = torch.arange(total_length).cuda()

# custom residue index (with 200 offset between chains)
position = residue_index

# repeating residue index
position = torch.arange(subunit_length).repeat(subunits).cuda()

[sokrypton]

position = 500 + torch.arange(total_length)

Also works!
So it's not an absolute thing... but it seems you need to preserve the relative encoding for OmegaFold to work.

I'm beginning to suspect that the reason why homo-oligomers work is that there are "fused" examples in the language model training set.

[RuiWang1998]

Hi @sokrypton,

Thanks for this info, we'll take a look into it.

Best

[RuiWang1998]

Hi Dr. @sokrypton,

Would you care to give us an example input for this pr?

[sokrypton]

For hetero-dimer:
https://www.rcsb.org/structure/7M5F

>H1065
AKNSLTTKSLFKEMTIQGIKFTPENVVGAAKDNSGKIIFLEKGNSKSGLQHIVEEHGDQFAQIGVSEARIPDVVMKAVTDGKIVGYQGAGAGRPIYETMIDGKKYNIAVTVGSNGYVVGANLRGSVK:MKEIKLMADYHCYPLWGTTPDDFGDISPDELPISLGLKNSLEAWAKRYDAILNTDDPALSGFKSVEEEKLFIDDGYKLAELLQEELGSAYKVIYHADY

For homo-oligomer:

>tmp
PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK:PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK

>tmp
PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK:PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK:PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK:PIAQIHILEGRSDEQKETLIREVSEAISRSLDAPLTSVRVIITEMAKGHFGIGGELASK

[RuiWang1998]

Hi Dr. @sokrypton,

We have reproduced your results and agree that it is really peculiar. We are going to take a deeper look into the reason for this.

As for the training set of the language model. We directly take them from Uniref 50 sequences.

[sokrypton]

I guess... since the LM was trained on single-chains. There is no reason to expect it to generalize to proteins it hasn't seen before, especially protein multimers. I suspect it sometimes works for protein multimers when the multi-chain protein looks like a multi-domain protein in the LM training set.

[RuiWang1998]

Hi,

We have been looking into this issue and found one problem with RoPE not necessarily regarding this problem, namely that it is symmetric, which may not be what we want and we are phasing it out soon.

But still this problem persists and we do not really have an idea yet.

[JackMaguire]

@sokrypton @RuiWang1998 As a hungry user, are the problems with the current branch technical or scientific?