xlstm

XLSTM: memoria extendida a largo plazo

Documento: https://arxiv.org/abs/2405.04517

Acerca de

XLSTM es una nueva arquitectura de red neuronal recurrente basada en ideas del LSTM original. A través de la activación exponencial con técnicas apropiadas de normalización y estabilización y una nueva memoria de matriz, supera las limitaciones del LSTM original y muestra un rendimiento prometedor en el modelado de idiomas en comparación con transformadores o modelos de espacio de estado.

xlstm grande 7b

Entrenamos un modelo de lenguaje de parámetro 7b XLSTM

Hemos optimizado la arquitectura XLSTM en términos de rendimiento de capacitación y estabilidad. El código para la arquitectura actualizada se encuentra en xlstm/xlstm_large .

Los pesos del modelo están disponibles en Huggingface en https://huggingface.co/nx-ai/xlstm-7b.

Instalación mínima

Cree un entorno de conda desde el archivo environment_pt220cu121.yaml . Instale solo el código de modelo (es decir, el módulo xlstm ) como paquete:

Instalar a través de PIP:

pip install xlstm

Clon de Github:

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

Para usar el modelo 7B XLSTM, instale mlstm_kernels a través de:

 pip install mlstm_kernels

Requisitos

Este paquete se basa en Pytorch y se probó para las versiones >=1.8 . Para la versión CUDA de SLSTM, necesita capacidad de cómputo> = 8.0, consulte https://developer.nvidia.com/cuda-gpus. Para un entorno bien probado, instale el environment_pt220cu121.yaml como:

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

Para el modelo XLSTM grande 7b requerimos nuestro paquete mlstm_kernels (TODO ADD GitHub), que proporciona núcleos rápidos para el XLSTM.

Modelos del papel XLSTM

Esta sección explica cómo usar los modelos del papel XLSTM.

Uso

Para aplicaciones que no son idiomas o para integrarse en otras arquitecturas, puede usar el xLSTMBlockStack y para el modelado de idiomas u otras aplicaciones basadas en token, puede usar el xLSTMLMModel .

pila de bloques XLSTM

El xLSTMBLockStack está destinado a su uso como columna vertebral alternativa en proyectos existentes. Es similar a una pila de bloques de transformador, pero usa bloques 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

Si está trabajando con cadenas / archivos YAML para la configuración, también puede usar DACITE para crear las DataClasses de configuración. Esto es lo mismo que el fragmento de arriba:

 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

modelo de idioma xlstm

El xLSTMLMModel es un envoltorio alrededor del xLSTMBlockStack que agrega la incrustación de token y la cabeza 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 )

Experimentos

Los experimentos sintéticos que muestran los beneficios de SLSTM sobre MLSTM y Vice Versa Best son la tarea de paridad y la tarea de recuperación asociativa de múltiples query. La tarea de paridad solo se puede resolver con las capacidades de seguimiento de estado proporcionadas por la mezcla de memoria de SLSTM. Las capacidades de memorización de la tarea de recuperación asociativa múltiple miden, donde la memoria matricial y la expansión estatal de MLSTM es muy beneficiosa. En combinación les va bien en ambas tareas.

Para ejecutar cada uno, ejecute main.py en la carpeta de experimentos como:

 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

Tenga en cuenta que el bucle de entrenamiento no contiene detención temprana o evaluación de pruebas.

Citación

Si usa esta base de código, o de otro modo encuentra nuestro trabajo valioso, cite el documento 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}, 
}