{"id":15601041,"url":"https://github.com/lucidrains/en-transformer","last_synced_at":"2025-04-05T08:09:40.024Z","repository":{"id":47645543,"uuid":"342772820","full_name":"lucidrains/En-transformer","owner":"lucidrains","description":"Implementation of E(n)-Transformer, which incorporates attention mechanisms into Welling's E(n)-Equivariant Graph Neural Network","archived":false,"fork":false,"pushed_at":"2024-06-02T15:19:47.000Z","size":84,"stargazers_count":219,"open_issues_count":8,"forks_count":30,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-03-29T07:09:52.503Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","deep-learning","equivariance","transformer"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-02-27T04:46:13.000Z","updated_at":"2025-03-20T19:46:18.000Z","dependencies_parsed_at":"2024-10-23T01:33:44.762Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/En-transformer","commit_stats":{"total_commits":110,"total_committers":1,"mean_commits":110.0,"dds":0.0,"last_synced_commit":"98a9f5fe0c1005e5c6349d290fb273a2c344f9b8"},"previous_names":[],"tags_count":82,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FEn-transformer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FEn-transformer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FEn-transformer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2FEn-transformer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/En-transformer/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247305935,"owners_count":20917208,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanism","deep-learning","equivariance","transformer"],"created_at":"2024-10-03T02:12:49.848Z","updated_at":"2025-04-05T08:09:39.993Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## E(n)-Equivariant Transformer\n\nImplementation of E(n)-Equivariant Transformer, which extends the ideas from Welling's \u003ca href=\"https://github.com/lucidrains/egnn-pytorch\"\u003eE(n)-Equivariant Graph Neural Network\u003c/a\u003e with attention mechanisms and ideas from transformer architecture.\n\nUpdate: Used for \u003ca href=\"https://www.biorxiv.org/content/10.1101/2023.07.15.549154v2\"\u003edesigning of CDR loops in antibodies!\u003c/a\u003e\n\n## Install\n\n```bash\n$ pip install En-transformer\n```\n\n## Usage\n\n```python\nimport torch\nfrom en_transformer import EnTransformer\n\nmodel = EnTransformer(\n    dim = 512,\n    depth = 4,                       # depth\n    dim_head = 64,                   # dimension per head\n    heads = 8,                       # number of heads\n    edge_dim = 4,                    # dimension of edge feature\n    neighbors = 64,                  # only do attention between coordinates N nearest neighbors - set to 0 to turn off\n    talking_heads = True,            # use Shazeer's talking heads https://arxiv.org/abs/2003.02436\n    checkpoint = True,               # use checkpointing so one can increase depth at little memory cost (and increase neighbors attended to)\n    use_cross_product = True,        # use cross product vectors (idea by @MattMcPartlon)\n    num_global_linear_attn_heads = 4 # if your number of neighbors above is low, you can assign a certain number of attention heads to weakly attend globally to all other nodes through linear attention (https://arxiv.org/abs/1812.01243)\n)\n\nfeats = torch.randn(1, 1024, 512)\ncoors = torch.randn(1, 1024, 3)\nedges = torch.randn(1, 1024, 1024, 4)\n\nmask = torch.ones(1, 1024).bool()\n\nfeats, coors = model(feats, coors, edges, mask = mask)  # (1, 1024, 512), (1, 1024, 3)\n```\n\nLetting the network take care of both atomic and bond type embeddings\n\n```python\nimport torch\nfrom en_transformer import EnTransformer\n\nmodel = EnTransformer(\n    num_tokens = 10,       # number of unique nodes, say atoms\n    rel_pos_emb = True,    # set this to true if your sequence is not an unordered set. it will accelerate convergence\n    num_edge_tokens = 5,   # number of unique edges, say bond types\n    dim = 128,\n    edge_dim = 16,\n    depth = 3,\n    heads = 4,\n    dim_head = 32,\n    neighbors = 8\n)\n\natoms = torch.randint(0, 10, (1, 16))    # 10 different types of atoms\nbonds = torch.randint(0, 5, (1, 16, 16)) # 5 different types of bonds (n x n)\ncoors = torch.randn(1, 16, 3)            # atomic spatial coordinates\n\nfeats_out, coors_out = model(atoms, coors, edges = bonds) # (1, 16, 512), (1, 16, 3)\n```\n\nIf you would like to only attend to sparse neighbors, as defined by an adjacency matrix (say for atoms), you have to set one more flag and then pass in the `N x N` adjacency matrix.\n\n```python\nimport torch\nfrom en_transformer import EnTransformer\n\nmodel = EnTransformer(\n    num_tokens = 10,\n    dim = 512,\n    depth = 1,\n    heads = 4,\n    dim_head = 32,\n    neighbors = 0,\n    only_sparse_neighbors = True,    # must be set to true\n    num_adj_degrees = 2,             # the number of degrees to derive from 1st degree neighbors passed in\n    adj_dim = 8                      # whether to pass the adjacency degree information as an edge embedding\n)\n\natoms = torch.randint(0, 10, (1, 16))\ncoors = torch.randn(1, 16, 3)\n\n# naively assume a single chain of atoms\ni = torch.arange(atoms.shape[1])\nadj_mat = (i[:, None] \u003c= (i[None, :] + 1)) \u0026 (i[:, None] \u003e= (i[None, :] - 1))\n\n# adjacency matrix must be passed in\nfeats_out, coors_out = model(atoms, coors, adj_mat = adj_mat) # (1, 16, 512), (1, 16, 3)\n```\n\n## Edges\n\nIf you need to pass in continuous edges\n\n```python\nimport torch\nfrom en_transformer import EnTransformer\nfrom en_transformer.utils import rot\n\nmodel = EnTransformer(\n    dim = 512,\n    depth = 1,\n    heads = 4,\n    dim_head = 32,\n    edge_dim = 4,\n    num_nearest_neighbors = 0,\n    only_sparse_neighbors = True\n)\n\nfeats = torch.randn(1, 16, 512)\ncoors = torch.randn(1, 16, 3)\nedges = torch.randn(1, 16, 16, 4)\n\ni = torch.arange(feats.shape[1])\nadj_mat = (i[:, None] \u003c= (i[None, :] + 1)) \u0026 (i[:, None] \u003e= (i[None, :] - 1))\n\nfeats1, coors1 = model(feats, coors, adj_mat = adj_mat, edges = edges)\n```\n\n## Example\n\nTo run a protein backbone coordinate denoising toy task, first install `sidechainnet`\n\n```bash\n$ pip install sidechainnet\n```\n\nThen\n\n```bash\n$ python denoise.py\n```\n\n## Todo\n\n- [ ] add https://arxiv.org/abs/2112.05682 for researchers to stretch to even bigger molecules\n\n## Citations\n\n```bibtex\n@misc{satorras2021en,\n    title \t= {E(n) Equivariant Graph Neural Networks}, \n    author \t= {Victor Garcia Satorras and Emiel Hoogeboom and Max Welling},\n    year \t= {2021},\n    eprint \t= {2102.09844},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n\n```bibtex\n@misc{shazeer2020talkingheads,\n    title   = {Talking-Heads Attention}, \n    author  = {Noam Shazeer and Zhenzhong Lan and Youlong Cheng and Nan Ding and Le Hou},\n    year    = {2020},\n    eprint  = {2003.02436},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n\n```bibtex\n@misc{liu2021swin,\n    title   = {Swin Transformer V2: Scaling Up Capacity and Resolution},\n    author  = {Ze Liu and Han Hu and Yutong Lin and Zhuliang Yao and Zhenda Xie and Yixuan Wei and Jia Ning and Yue Cao and Zheng Zhang and Li Dong and Furu Wei and Baining Guo},\n    year    = {2021},\n    eprint  = {2111.09883},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n\n```bibtex\n@inproceedings{Kim2020TheLC,\n    title   = {The Lipschitz Constant of Self-Attention},\n    author  = {Hyunjik Kim and George Papamakarios and Andriy Mnih},\n    booktitle = {International Conference on Machine Learning},\n    year    = {2020},\n    url     = {https://api.semanticscholar.org/CorpusID:219530837}\n}\n```\n\n```bibtex\n@article {Mahajan2023.07.15.549154,\n    author  = {Sai Pooja Mahajan and Jeffrey A. Ruffolo and Jeffrey J. Gray},\n    title   = {Contextual protein and antibody encodings from equivariant graph transformers},\n    elocation-id = {2023.07.15.549154},\n    year    = {2023},\n    doi     = {10.1101/2023.07.15.549154},\n    publisher = {Cold Spring Harbor Laboratory},\n    URL     = {https://www.biorxiv.org/content/early/2023/07/29/2023.07.15.549154},\n    eprint  = {https://www.biorxiv.org/content/early/2023/07/29/2023.07.15.549154.full.pdf},\n    journal = {bioRxiv}\n}\n```\n\n```bibtex\n@article{Bondarenko2023QuantizableTR,\n    title   = {Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing},\n    author  = {Yelysei Bondarenko and Markus Nagel and Tijmen Blankevoort},\n    journal = {ArXiv},\n    year    = {2023},\n    volume  = {abs/2306.12929},\n    url     = {https://api.semanticscholar.org/CorpusID:259224568}\n}\n```\n\n```bibtex\n@inproceedings{Arora2023ZoologyMA,\n    title   = {Zoology: Measuring and Improving Recall in Efficient Language Models},\n    author  = {Simran Arora and Sabri Eyuboglu and Aman Timalsina and Isys Johnson and Michael Poli and James Zou and Atri Rudra and Christopher R'e},\n    year    = {2023},\n    url     = {https://api.semanticscholar.org/CorpusID:266149332}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fen-transformer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fen-transformer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fen-transformer/lists"}