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

LightHuBERT: Lightweight and Configurable Speech Representation Learning with Once-for-All Hidden-Unit BERT

In this project we investigate the performance of the SetCon model on realistic video footage. Therefore, we implemented the model in PyTorch and tested the model on two example videos.

Pure python PEMDAS expression solver without using built-in eval function

Task-related Saliency Network For Few-shot learning

Recognize numbers from an (28 x 28) image using neural networks

Code for the paper "A Study of Face Obfuscation in ImageNet"

This repository contains all code and data for the Inside Out Visual Place Recognition task

NALSM: Neuron-Astrocyte Liquid State Machine

A `Neural = Symbolic` framework for sound and complete weighted real-value logic

Implementation of paper "Decision-based Black-box Attack Against Vision Transformers via Patch-wise Adversarial Removal"

TensorFlow CNN for fast style transfer

Twins: Revisiting the Design of Spatial Attention in Vision Transformers

An Efficient Implementation of Analytic Mesh Algorithm for 3D Iso-surface Extraction from Neural Networks

Official Pytorch implementation for "End2End Occluded Face Recognition by Masking Corrupted Features, TPAMI 2021"

Meta Language-Specific Layers in Multilingual Language Models

Pretrained models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet.

QQ Browser 2021 AI Algorithm Competition Track 1 1st Place Program

The implementation of ICASSP 2020 paper "Pixel-level self-paced learning for super-resolution"

fastgradio is a python library to quickly build and share gradio interfaces of your trained fastai models.

Expand human face editing via Global Direction of StyleCLIP, especially to maintain similarity during editing.