lucidrains / bidirectional-cross-attention

A simple cross attention that updates both the source and target in one step

Bidirectional Cross Attention

A simple cross attention that updates both the source and target in one step. The key insight is that one can do shared query / key attention and use the attention matrix twice to update both ways. Used for a contracting project for predicting DNA / protein binding here.

Install

$ pip install bidirectional-cross-attention

Usage

import torch
from bidirectional_cross_attention import BidirectionalCrossAttention

video = torch.randn(1, 4096, 512)
audio = torch.randn(1, 8192, 386)

video_mask = torch.ones((1, 4096)).bool()
audio_mask = torch.ones((1, 8192)).bool()

joint_cross_attn = BidirectionalCrossAttention(
    dim = 512,
    heads = 8,
    dim_head = 64,
    context_dim = 386
)

video_out, audio_out = joint_cross_attn(
    video,
    audio,
    mask = video_mask,
    context_mask = audio_mask
)

# attended output should have the same shape as input

assert video_out.shape == video.shape
assert audio_out.shape == audio.shape

Todo

• allow for cosine sim attention

Citations

As far as I know, I came up with it, but if you discover this in the literature, do let me know and I will cite it appropriately.

bidirectional-cross-attention's Issues

How to combine it into an Encoder stack

First, thanks for sharing this and other awesome repos!!

I was wondering how to create a transformer Encoder stack out of it, similarly to this: https://github.com/larngroup/DTITR (there is a nice pretty clear diagram at the top there)

Apart from "manually" creating such architecture, I was wondering, is there some way to tell another lib of yours - x_transformers, to use this custom cross-modality-cross-attention layer (or any type of related cross-modality-cross-attention which you support) to build an Encoder stack with it?

Thanks again for the awesome repos :)