Block Recurrent Transformer PyTorch

repository·main·Indexed 42 days ago

https://github.com/lucidrains/block-recurrent-transformer-pytorch

A PyTorch implementation of the Block Recurrent Transformer designed for long-range sequence modeling of 60k tokens and beyond. It utilizes recurrent states and compressed memories via the BlockRecurrentTransformer class to process sequences in segments while maintaining long-term dependencies.

Tokens: 1K
Snippets: 3
Records: 4
Agent score: 3%

What's inside Block Recurrent Transformer

Install block-recurrent-transformer-pytorch
main
Install the package using pip:
```
$ pip install block-recurrent-transformer-pytorch
```
https://github.com/lucidrains/block-recurrent-transformer-pytorch/blob/main/README.md
Run the Enwik8 training test
main
To run the provided training script for testing on the Enwik8 dataset:
1. Install requirements: pip install -r requirements.txt
2. Execute the training script: python train.py
```
$ pip install -r requirements.txt
$ python train.py
```
https://github.com/lucidrains/block-recurrent-transformer-pytorch/blob/main/README.md
Configure BlockRecurrentTransformer parameters
main
When initializing BlockRecurrentTransformer, you can tune the following parameters:
- num_tokens: Vocabulary size.
- dim: Model dimensions.
- depth: Number of layers.
- dim_head: Dimensions of each attention head.
- heads: Number of attention heads.
- max_seq_len: The total receptive field of the transformer.
- block_width: The block size. The block furthest forwards becomes the new cached xl memories.
- num_state_vectors: Number of state vectors.
- recurrent_layers: A tuple specifying the indices of the recurrent layers for states with fixed simple gating.
- use_compressed_mem: Boolean, whether to use compressed memories of a single block width.
- compressed_mem_factor: Compression factor for compressed memories.
- use_flash_attn: Boolean, enables Flash Attention (requires PyTorch 2.0+).
https://github.com/lucidrains/block-recurrent-transformer-pytorch/blob/main/README.md

Use the BlockRecurrentTransformer class

main

The BlockRecurrentTransformer is the main model class. It supports long-range dependencies by passing xl_memories and states between forward passes. This allows the model to process sequences in segments while maintaining a recurrent state and memory of previous blocks.

import torch
from block_recurrent_transformer_pytorch import BlockRecurrentTransformer

model = BlockRecurrentTransformer(
    num_tokens = 20000,             # vocab size
    dim = 512,                      # model dimensions
    depth = 6,                      # depth
    dim_head = 64,                  # attention head dimensions
    heads = 8,                      # number of attention heads
    max_seq_len = 1024,             # the total receptive field of the transformer
    block_width = 512,              # block size
    num_state_vectors = 512,        # number of state vectors
    recurrent_layers = (4,),        # where to place the recurrent layer(s) for states with fixed simple gating
    use_compressed_mem = False,     # whether to use compressed memories of a single block width
    compressed_mem_factor = 4,      # compression factor of compressed memories
    use_flash_attn = True           # use flash attention, if on pytorch 2.0
)

seq = torch.randint(0, 2000, (1, 1024))

# Initial forward pass
out, mems1, states1 = model(seq)

# Subsequent forward passes using memories and states from previous steps
out, mems2, states2 = model(seq, xl_memories = mems1, states = states1)
out, mems3, states3 = model(seq, xl_memories = mems2, states = states2)