CoLT5-attention

Install the package via pip:

$ pip install colt5-attention

The ConditionalRoutedAttention implements conditionally routed efficient attention. It uses a light branch with a local attention receptive field (light_window_size) and a heavy branch that receives a specific number of routed tokens for both queries and key-values.

import torch
from colt5_attention import ConditionalRoutedAttention

tokens = torch.randn(2, 32768, 512)
mask = torch.ones(2, 32768).bool()

attn = ConditionalRoutedAttention(
    dim = 512,
    light_dim_head = 64,
    light_heads = 8,
    light_window_size = 128,
    heavy_dim_head = 64,
    heavy_heads = 8,
    num_heavy_tokens_q = 1024,
    num_heavy_tokens_kv = 1024
)

attn_out = attn(tokens, mask = mask) # (2, 32768, 512)

The ConditionalRoutedTransformerBlock combines both routed attention and routed feedforward with residual connections. A stack of these blocks constitutes a CoLT5 encoder.

import torch
from colt5_attention import ConditionalRoutedTransformerBlock

tokens = torch.randn(2, 32768, 512)
mask = torch.ones(2, 32768).bool()

block = ConditionalRoutedTransformerBlock(
    dim = 512,
    light_dim_head = 64,
    light_heads = 8,
    light_window_size = 128,
    heavy_dim_head = 64,
    heavy_heads = 8,
    light_ff_mult = 0.5,
    heavy_ff_mult = 4,
    num_heavy_ff_tokens = 1024,
    num_heavy_attn_tokens_q = 1024,
    num_heavy_attn_tokens_kv = 1024
)

block_out = block(tokens, mask = mask) # (2, 32768, 512)

The ConditionalRoutedImageAttention is designed for high-resolution image feature maps. It uses a local window patch for the light branch and routed attention for the heavy branch.

Key parameters:

• channel_first: Set to True if the input has channels before spatial dimensions.

import torch
from colt5_attention import ConditionalRoutedImageAttention

attn = ConditionalRoutedImageAttention(
    dim = 32,
    light_dim_head = 64,
    light_heads = 8,
    light_window_size = 32,
    channel_first = True,
    heavy_dim_head = 64,
    heavy_heads = 8,
    num_heavy_tokens_q = 1024,
    num_heavy_tokens_kv = 1024
).cuda()

fmap = torch.randn(1, 32, 256, 256).cuda()

out = attn(fmap)

The ConditionalRoutedViT is a Vision Transformer implementation using coordinate descent routed attention and feedforward.

import torch
from colt5_attention.vit import ConditionalRoutedViT

vit = ConditionalRoutedViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    attn_num_heavy_tokens_q = 16,
    attn_num_heavy_tokens_kv = 16,
    attn_heavy_dim_head = 64,
    attn_heavy_heads = 8,
    attn_light_window_size = 4,
    attn_light_dim_head = 32,
    attn_light_heads = 4,
    ff_num_heavy_tokens = 16,
    ff_heavy_mult = 4,
    ff_light_mult = 2
)

images = torch.randn(1, 3, 256, 256)
logits = vit(images)

The topk function provides a differentiable version of topk using a small wrapper around coordinate descent.

Returns:

• values: The values, which have already been multiplied by the gates.
• indices: The topk indices.
• coor_descent_values: The coordinate descent values.
• gates: The routing gates.

import torch
from colt5_attention import topk

x = torch.randn(1024, 512)

values, indices, coor_descent_values, gates = topk(x, k = 10, fused = True)

The ConditionalRoutedFeedForward implements a routed feedforward block where a subset of tokens is routed to a 'heavy' branch with a larger hidden dimension, while the rest use a 'light' branch. The outputs are summed.

import torch
from colt5_attention import ConditionalRoutedFeedForward

tokens = torch.randn(2, 32768, 512)
mask = torch.ones(2, 32768).bool()

ff = ConditionalRoutedFeedForward(
    dim = 512,
    light_ff_mult = 0.5,      # hidden dimension ratio of light branch
    heavy_ff_mult = 4,        # hidden dimension ratio of heavy branch
    num_heavy_tokens = 1024   # heavy branch receives only 1024 routed tokens of 32768
)

ff_out = ff(tokens, mask = mask)  # (2, 32768, 512)