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

It is the assignment for COMP 576 in Rice University

PyoMyo - Python Opensource Myo library

training script for space time memory network

SmallInitEmb - LayerNorm(SmallInit(Embedding)) in a Transformer to improve convergence

【Arxiv】Exploring Separable Attention for Multi-Contrast MR Image Super-Resolution

Jupyter Dock is a set of Jupyter Notebooks for performing molecular docking protocols interactively, as well as visualizing, converting file formats and analyzing the results.

Pytorch implementation for ACMMM2021 paper "I2V-GAN: Unpaired Infrared-to-Visible Video Translation".

Self-Supervised Learning with Kernel Dependence Maximization

MOOSE (Multi-organ objective segmentation) a data-centric AI solution that generates multilabel organ segmentations to facilitate systemic TB whole-person research

Fairness Metrics: All you need to know

[ICCV'21] Pri3D: Can 3D Priors Help 2D Representation Learning?

Resources for the "Evaluating the Factual Consistency of Abstractive Text Summarization" paper

Code for the ICCV 2021 Workshop paper: A Unified Efficient Pyramid Transformer for Semantic Segmentation.

DNA sequence classification by Deep Neural Network

Official re-implementation of the Calibrated Adversarial Refinement model described in the paper Calibrated Adversarial Refinement for Stochastic Semantic Segmentation

This repository is for DSA and CP scripts for reference.

Avalanche RL: an End-to-End Library for Continual Reinforcement Learning

Axel - 3D printed robotic hands and they controll with Raspberry Pi and Arduino combo

Monocular 3D Object Detection: An Extrinsic Parameter Free Approach (CVPR2021)

Viewmaker Networks: Learning Views for Unsupervised Representation Learning