LongNet

LongNet is a Transformer variant designed to scale sequence lengths up to 1 billion tokens.

Core Mechanism: It uses dilated attention, which expands the attentive field exponentially as the distance between tokens grows. This allows for linear computational complexity and a logarithmic dependency between tokens, solving the bottleneck of standard Transformers where sequence length scaling is computationally prohibitive.

Key Advantages:

• Linear Complexity: Scales efficiently to extremely long sequences.
• Drop-in Replacement: Dilated attention can replace standard attention in existing Transformer architectures.
• Distributed Training: Can be used as a distributed trainer for extremely long sequences.

To perform a simple training run on the enwiki8 dataset:

1. Clone the repository.
2. Install dependencies using pip install -r requirements.txt.
3. Execute the training script with python3 train.py.

Install the longnet package using pip to access the Dilated Attention and LongNetTransformer implementations.

pip install longnet

The LongNetTransformer is a complete transformer model designed for training. It includes dilated transformer blocks with Feedforwards, LayerNorm, SWIGLU, and parallel transformer blocks.

Parameters:

• num_tokens: Vocabulary size.
• dim: Model dimension.
• depth: Number of layers.
• dim_head: Dimension of each head.
• heads: Number of attention heads.
• ff_mult: Feedforward multiplier.

import torch
from long_net.model import LongNetTransformer

longnet = LongNetTransformer(
    num_tokens=20000,
    dim=512,
    depth=6,
    dim_head=64,
    heads=8,
    ff_mult=4,
)

tokens = torch.randint(0, 20000, (1, 512))
logits = longnet(tokens)
print(logits)

The DilatedAttention class provides the core dilated attention mechanism. It can be used as a drop-in replacement for standard attention to handle longer sequences with linear computational complexity.

Parameters:

• dim: The dimension of the model.
• heads: Number of attention heads.
• dilation_rate: The rate of dilation.
• segment_size: The size of the segments used for attention.
• qk_norm (optional): Boolean to enable Query-Key normalization.

import torch
from long_net import DilatedAttention

# model config
dim = 512
heads = 8
dilation_rate = 2
segment_size = 64

# input data
batch_size = 32
seq_len = 8192

# create model and data
model = DilatedAttention(dim, heads, dilation_rate, segment_size, qk_norm=True)
x = torch.randn((batch_size, seq_len, dim))

output = model(x)
print(output)