FLASH pytorch

ВСПЫШКА — Пайторч

Реализация варианта трансформатора, предложенного в статье «Качество трансформатора в линейном времени».

Установить

$ pip install FLASH-pytorch

Использование

Основной новой схемой в этой статье является «Единица закрытого внимания», которая, как они утверждают, может заменить многоголовое внимание, сократив его до одной головы.

Он использует квадратную активацию relu вместо softmax, активация которой впервые была замечена в статье Primer, а также использование ReLU в ReLA Transformer. Стиль шлюзования, похоже, в основном вдохновлен gMLP.

 import torch
from flash_pytorch import GAU

gau = GAU (
    dim = 512 ,
    query_key_dim = 128 ,     # query / key dimension
    causal = True ,           # autoregressive or not
    expansion_factor = 2 ,    # hidden dimension = dim * expansion_factor
    laplace_attn_fn = True   # new Mega paper claims this is more stable than relu squared as attention function
)

x = torch . randn ( 1 , 1024 , 512 )
out = gau ( x ) # (1, 1024, 512)

Затем авторы объединяют GAU с линейным вниманием Катаропулоса, используя группировку последовательностей, чтобы преодолеть известную проблему с авторегрессионным линейным вниманием.

Эту комбинацию квадратичной закрытой единицы внимания с групповым линейным вниманием они назвали FLASH.

Вы также можете использовать это довольно легко

 import torch
from flash_pytorch import FLASH

flash = FLASH (
    dim = 512 ,
    group_size = 256 ,             # group size
    causal = True ,                # autoregressive or not
    query_key_dim = 128 ,          # query / key dimension
    expansion_factor = 2. ,        # hidden dimension = dim * expansion_factor
    laplace_attn_fn = True   # new Mega paper claims this is more stable than relu squared as attention function
)

x = torch . randn ( 1 , 1111 , 512 )     # sequence will be auto-padded to nearest group size
out = flash ( x ) # (1, 1111, 512)

Наконец, вы можете использовать полноценный преобразователь FLASH, как указано в статье. Он содержит все позиционные вложения, упомянутые в статье. Абсолютное позиционное внедрение использует масштабированную синусоиду. Квадратичное внимание GAU получит одноголовое относительное позиционное смещение T5. Помимо всего этого, как внимание GAU, так и линейное внимание будут вращающимися (RoPE).

 import torch
from flash_pytorch import FLASHTransformer

model = FLASHTransformer (
    num_tokens = 20000 ,          # number of tokens
    dim = 512 ,                   # model dimension
    depth = 12 ,                  # depth
    causal = True ,               # autoregressive or not
    group_size = 256 ,            # size of the groups
    query_key_dim = 128 ,         # dimension of queries / keys
    expansion_factor = 2. ,       # hidden dimension = dim * expansion_factor
    norm_type = 'scalenorm' ,     # in the paper, they claimed scalenorm led to faster training at no performance hit. the other option is 'layernorm' (also default)
    shift_tokens = True          # discovered by an independent researcher in Shenzhen @BlinkDL, this simply shifts half of the feature space forward one step along the sequence dimension - greatly improved convergence even more in my local experiments
)

x = torch . randint ( 0 , 20000 , ( 1 , 1024 ))
logits = model ( x ) # (1, 1024, 20000) 

Тест на авторегрессии Enwik8

$ python train.py

Цитаты

 @article { Hua2022TransformerQI ,
    title   = { Transformer Quality in Linear Time } ,
    author  = { Weizhe Hua and Zihang Dai and Hanxiao Liu and Quoc V. Le } ,
    journal = { ArXiv } ,
    year    = { 2022 } ,
    volume  = { abs/2202.10447 }
}

 @software { peng_bo_2021_5196578 ,
    author    = { PENG Bo } ,
    title     = { BlinkDL/RWKV-LM: 0.01 } ,
    month     = { aug } ,
    year      = { 2021 } ,
    publisher = { Zenodo } ,
    version   = { 0.01 } ,
    doi       = { 10.5281/zenodo.5196578 } ,
    url       = { https://doi.org/10.5281/zenodo.5196578 }
}

 @inproceedings { Ma2022MegaMA ,
    title   = { Mega: Moving Average Equipped Gated Attention } ,
    author  = { Xuezhe Ma and Chunting Zhou and Xiang Kong and Junxian He and Liangke Gui and Graham Neubig and Jonathan May and Luke Zettlemoyer } ,
    year    = { 2022 }
}