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

A simple Rock-Paper-Scissors game using CV in python

SeMask: Semantically Masked Transformers for Semantic Segmentation.

Codes and Data Processing Files for our paper.

Code for "Single-view robot pose and joint angle estimation via render & compare", CVPR 2021 (Oral).

The Surprising Effectiveness of Visual Odometry Techniques for Embodied PointGoal Navigation

Neural network for digit classification powered by cuda

PyTorch Implementation of ByteDance's Cross-speaker Emotion Transfer Based on Speaker Condition Layer Normalization and Semi-Supervised Training in Text-To-Speech

Principled Detection of Out-of-Distribution Examples in Neural Networks

[v1 (ISBI'21) + v2] MedMNIST: A Large-Scale Lightweight Benchmark for 2D and 3D Biomedical Image Classification

Alpha-Zero - Telegram Group Manager Bot Written In Python Using Pyrogram

FFCV: Fast Forward Computer Vision (and other ML workloads!)

Official code for the CVPR 2021 paper "How Well Do Self-Supervised Models Transfer?"

Pytorch implementation of "Neural Wireframe Renderer: Learning Wireframe to Image Translations"

Guiding evolutionary strategies by (inaccurate) differentiable robot simulators @ NeurIPS, 4th Robot Learning Workshop

pytorch implementation of GPV-Pose

Overview of architecture and implementation of TEDS-Net, as described in MICCAI 2021: "TEDS-Net: Enforcing Diffeomorphisms in Spatial Transformers to Guarantee TopologyPreservation in Segmentations"

Official implementation of the ICCV 2021 paper "Conditional DETR for Fast Training Convergence".

Revitalizing CNN Attention via Transformers in Self-Supervised Visual Representation Learning

code for paper "Not All Unlabeled Data are Equal: Learning to Weight Data in Semi-supervised Learning" by Zhongzheng Ren*, Raymond A. Yeh*, Alexander G. Schwing.

ACL'2021: LM-BFF: Better Few-shot Fine-tuning of Language Models

code for paper "Not All Unlabeled Data are Equal: Learning to Weight Data in Semi-supervised Learning" by Zhongzheng Ren, Raymond A. Yeh, Alexander G. Schwing.