Bidirectional Cross Attention

The BidirectionalCrossAttention module allows for updating both a source sequence and a target sequence in a single step using shared query/key attention. This is useful for multi-modal tasks (e.g., video and audio) where both sequences need to attend to each other.

Key features:

• Updates both sequences simultaneously.
• Supports masking for both the primary sequence and the context sequence.
• The attended output for each sequence maintains the same shape as its respective input.

import torch
from bidirectional_cross_attention import BidirectionalCrossAttention

# Example sequences (e.g., video and audio)
video = torch.randn(1, 4096, 512)
audio = torch.randn(1, 8192, 386)

# Masks for the sequences
video_mask = torch.ones((1, 4096)).bool()
audio_mask = torch.ones((1, 8192)).bool()

# Initialize the module
joint_cross_attn = BidirectionalCrossAttention(
    dim = 512,
    heads = 8,
    dim_head = 64,
    context_dim = 386
)

# Perform bidirectional cross attention
video_out, audio_out = joint_cross_attn(
    video,
    audio,
    mask = video_mask,
    context_mask = audio_mask
)

# Verify output shapes match input shapes
assert video_out.shape == video.shape
assert audio_out.shape == audio.shape