atomicarchitects / equiformer_v2

[ICLR 2024] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
https://arxiv.org/abs/2306.12059
MIT License
218 stars 28 forks source link

Using Equiformer for Encoding Protein 3D Structure #13

Closed mahdip72 closed 4 months ago

mahdip72 commented 4 months ago

Thanks for putting together this awesome repo!

I'm working on a VQVAE model for the backbone atoms of protein 3D structures and thinking about using Equiformer V2 as my first part the encoder. I used the original SE(3) model before, but it was super computationally heavy.

My plan is to quantize protein 3D structures residue-wise based on the number of carbon alpha atoms and only pass the coordinates to the encoder, leaving out the amino acid type info.

Do you have any suggestions or thoughts on how well this setup might work?

yilunliao commented 4 months ago

Hi @mahdip72

Thanks for your interest in EquiformerV2.

I don't have too much thoughts though.

For modeling protein, I would first go with what is the common practice (i.e., how they model the residues). But I am not familiar with that line of work.

As for your application (i.e., VQVAE), I assume it is something related to reconstruction or maybe denoising diffusion models. Here is an evidence showing higher degrees are indeed better: https://github.com/lucidrains/equiformer-pytorch/issues/6 Also in Table 12 in the paper, we show that higher degrees are helpful for denoising. But you might start with just lower degrees to make sure your overall methods are correct and then use higher degrees to get the gain.

Let me know if you have more concrete problems.