Implementation of Lie Transformer, Equivariant Self-Attention, in Pytorch

Last update: Oct 26, 2022

Overview

Lie Transformer - Pytorch (wip)

Implementation of Lie Transformer, Equivariant Self-Attention, in Pytorch. Only the SE3 version will be present in this repository, as it may be needed for Alphafold2 replication.

Install

$ pip install lie-transformer-pytorch

Usage

import torch
from lie_transformer_pytorch import LieTransformer

model = LieTransformer(
    dim = 512,
    depth = 2,
    heads = 8,
    dim_head = 64,
    liftsamples = 4
)

coors = torch.randn(1, 64, 3)
features = torch.randn(1, 64, 512)
mask = torch.ones(1, 64).bool()

out = model(features, coors, mask = mask) # (1, 256, 512) <- 256 = (seq len * liftsamples)

Todo

Figure out location based attention as described in section 3.2 I think it may be as simple as https://github.com/lucidrains/lie-transformer-pytorch/blob/main/lie_transformer_pytorch/lie_transformer_pytorch.py#L262-L265 . Welcoming any feedback

Credit

This repository is largely adapted from LieConv, cited below!

Citations

@misc{hutchinson2020lietransformer,
    title       = {LieTransformer: Equivariant self-attention for Lie Groups}, 
    author      = {Michael Hutchinson and Charline Le Lan and Sheheryar Zaidi and Emilien Dupont and Yee Whye Teh and Hyunjik Kim},
    year        = {2020},
    eprint      = {2012.10885},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG}
}

@misc{finzi2020generalizing,
    title   = {Generalizing Convolutional Neural Networks for Equivariance to Lie Groups on Arbitrary Continuous Data}, 
    author  = {Marc Finzi and Samuel Stanton and Pavel Izmailov and Andrew Gordon Wilson},
    year    = {2020},
    eprint  = {2002.12880},
    archivePrefix = {arXiv},
    primaryClass = {stat.ML}
}

Comments

Help needed
It seems like we need equivariance (combined with attention) for alphafold2, so I am currently working on getting this and https://github.com/lucidrains/se3-transformer-pytorch ready. This is in preparation for a full replication for protein folding in silico as more details emerge.

I could use some help, for anyone who is more of an expert in Lie Group theory out there, to get this library to a state where it is usable.

Figure out location based attention as described in section 3.2 I think it may be as simple as https://github.com/lucidrains/lie-transformer-pytorch/blob/main/lie_transformer_pytorch/lie_transformer_pytorch.py#L265-L269 .

Any feedback would be appreciated!
help wanted
opened by lucidrains 1

Implementation of the 😇 Attention layer from the paper, Scaling Local Self-Attention For Parameter Efficient Visual Backbones

HaloNet - Pytorch Implementation of the Attention layer from the paper, Scaling Local Self-Attention For Parameter Efficient Visual Backbones. This re

189 Nov 22, 2022

Implementation of a memory efficient multi-head attention as proposed in the paper, "Self-attention Does Not Need O(n²) Memory"

Memory Efficient Attention Pytorch Implementation of a memory efficient multi-head attention as proposed in the paper, Self-attention Does Not Need O(

180 Jan 5, 2023

EGNN - Implementation of E(n)-Equivariant Graph Neural Networks, in Pytorch

EGNN - Pytorch Implementation of E(n)-Equivariant Graph Neural Networks, in Pytorch. May be eventually used for Alphafold2 replication. This

259 Jan 4, 2023

The official TensorFlow implementation of the paper Action Transformer: A Self-Attention Model for Short-Time Pose-Based Human Action Recognition

Action Transformer A Self-Attention Model for Short-Time Human Action Recognition This repository contains the official TensorFlow implementation of t

20 Jan 3, 2023

The official implementation of ELSA: Enhanced Local Self-Attention for Vision Transformer

Implementation of the method proposed in the paper "Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation"

Neural Descriptor Fields (NDF) PyTorch implementation for training continuous 3D neural fields to represent dense correspondence across objects, and u

167 Jan 6, 2023

Implementation of Lie Transformer, Equivariant Self-Attention, in Pytorch

Related tags

Overview

Lie Transformer - Pytorch (wip)

Install

Usage

Todo

Credit

Citations

You might also like...

Implementation of the 😇 Attention layer from the paper, Scaling Local Self-Attention For Parameter Efficient Visual Backbones

Implementation of a memory efficient multi-head attention as proposed in the paper, "Self-attention Does Not Need O(n²) Memory"

EGNN - Implementation of E(n)-Equivariant Graph Neural Networks, in Pytorch

The official TensorFlow implementation of the paper Action Transformer: A Self-Attention Model for Short-Time Pose-Based Human Action Recognition

The official implementation of ELSA: Enhanced Local Self-Attention for Vision Transformer

Equivariant CNNs for the sphere and SO(3) implemented in PyTorch

git《Self-Attention Attribution: Interpreting Information Interactions Inside Transformer》(AAAI 2021) GitHub:

Official implementation of the paper "Topographic VAEs learn Equivariant Capsules"

Implementation of the method proposed in the paper "Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation"

Comments

Help needed

Releases(0.0.17)

0.0.17(Feb 19, 2021)

0.0.16(Feb 19, 2021)

0.0.15(Feb 9, 2021)

0.0.14(Feb 9, 2021)

0.0.12(Feb 9, 2021)

0.0.11a(Feb 9, 2021)

0.0.10a(Feb 3, 2021)

0.0.10(Feb 3, 2021)

0.0.9(Feb 2, 2021)

0.0.8(Feb 2, 2021)

0.0.6(Jan 23, 2021)

0.0.5(Jan 23, 2021)

0.0.4(Jan 23, 2021)

0.0.3(Jan 22, 2021)

0.0.2(Jan 22, 2021)

0.0.1(Jan 22, 2021)

Owner

Phil Wang

Implementation for paper: Self-Regulation for Semantic Segmentation

In the case of your data having only 1 channel while want to use timm models

Code for A Volumetric Transformer for Accurate 3D Tumor Segmentation

Implementation of our paper "DMT: Dynamic Mutual Training for Semi-Supervised Learning"

Research on Tabular Deep Learning (Python package & papers)

Suite of 500 procedurally-generated NLP tasks to study language model adaptability

Multi-View Radar Semantic Segmentation

Training DiffWave using variational method from Variational Diffusion Models.

A plug-and-play library for neural networks written in Python

Deep Learning for humans

Official pytorch implementation of "Scaling-up Disentanglement for Image Translation", ICCV 2021.

Your interactive network visualizing dashboard

Official implementation of "Accelerating Reinforcement Learning with Learned Skill Priors", Pertsch et al., CoRL 2020

Matching python environment code for Lux AI 2021 Kaggle competition, and a gym interface for RL models.

Video Background Music Generation with Controllable Music Transformer (ACM MM 2021 Oral)

🤖 Project template for your next awesome AI project. 🦾

A PyTorch implementation of Radio Transformer Networks from the paper "An Introduction to Deep Learning for the Physical Layer".

Ansible Automation Example: JSNAPY PRE/POST Upgrade Validation

A face dataset generator with out-of-focus blur detection and dynamic interval adjustment.

AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty