xlstm

XLSTM: расширенная длинная кратковременная память

Бумага: https://arxiv.org/abs/2405.04517

О

XLSTM - это новая рецидивирующая архитектура нейронной сети, основанная на идеях оригинального LSTM. Благодаря экспоненциальному стробированию с соответствующими методами нормализации и стабилизации и новой матричной памятью, оно преодолевает ограничения исходного LSTM и показывает многообещающую производительность на языковом моделировании по сравнению с трансформаторами или моделями пространства состояний.

xlstm большой 7b

Мы обучили параметр 7b xlstm языковой модель

Мы оптимизировали архитектуру XLSTM с точки зрения обучения пропускной способности и стабильности. Код для обновленной архитектуры находится в xlstm/xlstm_large .

Веса модели доступны на huggingface по адресу https://huggingface.co/nx-ai/xlstm-7b.

Минимальная установка

Создайте среду Conda из File environment_pt220cu121.yaml . Установите только код модели (то есть модуль xlstm ) в качестве пакета:

Установите через PIP:

pip install xlstm

Клон от GitHub:

git clone https://github.com/NX-AI/xlstm.git
cd xlstm
pip install -e .

Для использования модели 7B XLSTM установить mlstm_kernels через:

 pip install mlstm_kernels

Требования

Этот пакет основан на Pytorch и был протестирован для версий >=1.8 . Для версии CUDA SLSTM вам нужна возможность вычисления> = 8.0, см. Https://developer.nvidia.com/cuda-gpus. Для хорошо проверенной среды установите environment_pt220cu121.yaml как:

conda env create -n xlstm -f environment_pt220cu121.yaml
conda activate xlstm

Для большой модели XLSTM 7B нам требуется наш пакет mlstm_kernels (TODO ADD GITHUB LINK), который обеспечивает быстрые ядра для XLSTM.

Модели из бумаги XLSTM

В этом разделе объясняется, как использовать модели из бумаги XLSTM.

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

Для приложений без языка или для интеграции в другие архитектуры вы можете использовать xLSTMBlockStack и для языкового моделирования или других приложений, основанных на токенах, вы можете использовать xLSTMLMModel .

xlstm block Stack

xLSTMBLockStack предназначен для использования в качестве альтернативной основы в существующих проектах. Это похоже на стопку блоков трансформатора, но использует блоки XLSTM:

 import torch

from xlstm import (
    xLSTMBlockStack ,
    xLSTMBlockStackConfig ,
    mLSTMBlockConfig ,
    mLSTMLayerConfig ,
    sLSTMBlockConfig ,
    sLSTMLayerConfig ,
    FeedForwardConfig ,
)

cfg = xLSTMBlockStackConfig (
    mlstm_block = mLSTMBlockConfig (
        mlstm = mLSTMLayerConfig (
            conv1d_kernel_size = 4 , qkv_proj_blocksize = 4 , num_heads = 4
        )
    ),
    slstm_block = sLSTMBlockConfig (
        slstm = sLSTMLayerConfig (
            backend = "cuda" ,
            num_heads = 4 ,
            conv1d_kernel_size = 4 ,
            bias_init = "powerlaw_blockdependent" ,
        ),
        feedforward = FeedForwardConfig ( proj_factor = 1.3 , act_fn = "gelu" ),
    ),
    context_length = 256 ,
    num_blocks = 7 ,
    embedding_dim = 128 ,
    slstm_at = [ 1 ],

)

xlstm_stack = xLSTMBlockStack ( cfg )

x = torch . randn ( 4 , 256 , 128 ). to ( "cuda" )
xlstm_stack = xlstm_stack . to ( "cuda" )
y = xlstm_stack ( x )
y . shape == x . shape

Если вы работаете со строками / файлами YAML для конфигурации, вы также можете использовать Dacite для создания DATACLASES COUNT. Это то же самое, что и фрагмент выше:

 from omegaconf import OmegaConf
from dacite import from_dict
from dacite import Config as DaciteConfig
from xlstm import xLSTMBlockStack , xLSTMBlockStackConfig

xlstm_cfg = """ 
mlstm_block:
  mlstm:
    conv1d_kernel_size: 4
    qkv_proj_blocksize: 4
    num_heads: 4
slstm_block:
  slstm:
    backend: cuda
    num_heads: 4
    conv1d_kernel_size: 4
    bias_init: powerlaw_blockdependent
  feedforward:
    proj_factor: 1.3
    act_fn: gelu
context_length: 256
num_blocks: 7
embedding_dim: 128
slstm_at: [1]
"""
cfg = OmegaConf . create ( xlstm_cfg )
cfg = from_dict ( data_class = xLSTMBlockStackConfig , data = OmegaConf . to_container ( cfg ), config = DaciteConfig ( strict = True ))
xlstm_stack = xLSTMBlockStack ( cfg )

x = torch . randn ( 4 , 256 , 128 ). to ( "cuda" )
xlstm_stack = xlstm_stack . to ( "cuda" )
y = xlstm_stack ( x )
y . shape == x . shape

xlstm языковая модель

xLSTMLMModel - это обертка вокруг xLSTMBlockStack , которая добавляет встраивание токена и головку LM.

 from omegaconf import OmegaConf
from dacite import from_dict
from dacite import Config as DaciteConfig
from xlstm import xLSTMLMModel , xLSTMLMModelConfig

xlstm_cfg = """ 
vocab_size: 50304
mlstm_block:
  mlstm:
    conv1d_kernel_size: 4
    qkv_proj_blocksize: 4
    num_heads: 4
slstm_block:
  slstm:
    backend: cuda
    num_heads: 4
    conv1d_kernel_size: 4
    bias_init: powerlaw_blockdependent
  feedforward:
    proj_factor: 1.3
    act_fn: gelu
context_length: 256
num_blocks: 7
embedding_dim: 128
slstm_at: [1]
"""
cfg = OmegaConf . create ( xlstm_cfg )
cfg = from_dict ( data_class = xLSTMLMModelConfig , data = OmegaConf . to_container ( cfg ), config = DaciteConfig ( strict = True ))
xlstm_stack = xLSTMLMModel ( cfg )

x = torch . randint ( 0 , 50304 , size = ( 4 , 256 )). to ( "cuda" )
xlstm_stack = xlstm_stack . to ( "cuda" )
y = xlstm_stack ( x )
y . shape [ 1 :] == ( 256 , 50304 )

Эксперименты

Синтетические эксперименты демонстрируют преимущества SLSTM по сравнению с MLSTM, и наоборот, лучшая задача по паритету и задача Multi-Query Assocative Assive. Задача паритета может быть решена только с помощью возможностей отслеживания состояний, предоставляемых смешением памяти SLSTM. Многопрофильная задача Assocative Assepetion Измеряет возможности запоминания, где матричная память и расширение MLSTM в состоянии матрицы очень полезны. В сочетании они преуспевают в обеих задачах.

Чтобы запустить каждый, запустите main.py в папке экспериментов, как:

 python experiments/main.py --config experiments/parity_xLSTM01.yaml   # xLSTM[0:1], sLSTM only
python experiments/main.py --config experiments/parity_xLSTM10.yaml   # xLSTM[1:0], mLSTM only
python experiments/main.py --config experiments/parity_xLSTM11.yaml   # xLSTM[1:1], mLSTM and sLSTM

Обратите внимание, что обучающая петля не содержит ранней остановки или оценки тестов.

Цитирование

Если вы используете эту кодовую базу или иным образом найдите нашу работу ценной, пожалуйста, укажите бумагу XLSTM:

 @inproceedings{beck:24xlstm,
      title={xLSTM: Extended Long Short-Term Memory}, 
      author={Maximilian Beck and Korbinian Pöppel and Markus Spanring and Andreas Auer and Oleksandra Prudnikova and Michael Kopp and Günter Klambauer and Johannes Brandstetter and Sepp Hochreiter},
      booktitle = {Thirty-eighth Conference on Neural Information Processing Systems},
      year={2024},
      url={https://arxiv.org/abs/2405.04517}, 
}