Linformer PyTorch Implementation

repository·master·Indexed 43 days ago

A PyTorch implementation of the Linformer architecture providing linear complexity self-attention for Transformers. Includes the LinformerLM for language modeling, a general Linformer module, and LinformerSelfAttention for standalone layers supporting standard and contextual attention modes.

Tokens: 983
Snippets: 4
Records: 4
Agent score: 5%

What's inside Linformer

Install linformer via pip
master
Install the linformer package using pip to use the Linformer implementation in PyTorch.
```
$ pip install linformer
```
https://github.com/lucidrains/linformer/blob/master/README.md

Use LinformerLM for language modeling

master

The LinformerLM class provides a full language model implementation. Note that Linformer assumes a fixed sequence length and does not support auto-regressive tasks.

import torch
from linformer import LinformerLM

model = LinformerLM(
    num_tokens = 20000,
    dim = 512,
    seq_len = 4096,
    depth = 12,
    heads = 8,
    dim_head = 128,        # dimension of each head in multi-head attention
    k = 256,               # projection dimension for keys/values along sequence dimension
    one_kv_head = True,    # share one key/value head across all heads
    share_kv = False,      # share the same projection for keys and values
    reversible = True      # make network reversible, like Reformer
)

x = torch.randint(0, 20000, (1, 4096))
output = model(x) # returns shape (1, 4096, 20000)

Use the Linformer module

master

The Linformer class is a transformer-style module that processes input tensors through Linformer layers.

import torch
from linformer import Linformer

model = Linformer(
    dim = 512,
    seq_len = 4096,
    depth = 12,
    heads = 8,
    k = 256,
    one_kv_head = True,
    share_kv = True
)

x = torch.randn(1, 4096, 512)
output = model(x) # returns shape (1, 4096, 512)

Use LinformerSelfAttention for single attention layers

master

The LinformerSelfAttention class allows you to use Linformer's linear complexity attention as a standalone layer. It supports two modes:

Standard Self-Attention: Pass a single tensor x.
Contextual Attention: Pass a source tensor x and a context tensor. In this mode, the seq_len parameter is validated against the length of the context keys rather than the source sequence x.

import torch
from linformer import LinformerSelfAttention

# Standard Self-Attention
attn = LinformerSelfAttention(
    dim = 512,
    seq_len = 4096,
    heads = 8,
    k = 256,
    one_kv_head = True,
    share_kv = True
)

x = torch.randn(1, 4096, 512)
output = attn(x) # (1, 4096, 512)

# Contextual Attention
attn_context = LinformerSelfAttention(
    dim = 512,
    seq_len = 8192,
    heads = 8,
    k = 256,
    one_kv_head = True,
    share_kv = True
)

x_src = torch.randn(1, 2048, 512)
context = torch.randn(1, 8192, 512)
output_context = attn_context(x_src, context) # (1, 2048, 512)