xlstm

XLSTM: Mémoire à long terme prolongée

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

À propos

XLSTM est une nouvelle architecture de réseau neuronal récurrente basée sur les idées du LSTM original. Grâce à une déclenchement exponentielle avec des techniques de normalisation et de stabilisation appropriées et une nouvelle mémoire matricielle, il surmonte les limites du LSTM d'origine et montre des performances prometteuses sur la modélisation du langage par rapport aux transformateurs ou aux modèles d'espace d'état.

xlstm grand 7b

Nous avons formé un modèle de langue XLSTM paramètre 7B

Nous avons optimisé l'architecture XLSTM en termes de débit de formation et de stabilité. Le code de l'architecture mis à jour est situé dans xlstm/xlstm_large .

Les poids du modèle sont disponibles sur HuggingFace sur https://huggingface.co/nx-ai/xlstm-7b.

Installation minimale

Créez un environnement conda à partir de l' environment_pt220cu121.yaml de fichier_pt220cu121.yaml. Installez le code de modèle uniquement (c'est-à-dire le module xlstm ) en tant que package:

Installer via PIP:

pip install xlstm

Clone de GitHub:

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

Pour utiliser le modèle 7B XLSTM, installez mlstm_kernels via:

 pip install mlstm_kernels

Exigences

Ce package est basé sur Pytorch et a été testé pour les versions >=1.8 . Pour la version CUDA de SLSTM, vous avez besoin de capacité de calcul> = 8.0, voir https://developer.nvidia.com/cuda-gpus. Pour un environnement bien testé, installez l' environment_pt220cu121.yaml comme:

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

Pour le modèle 7B XLSTM, nous avons besoin de notre package mlstm_kernels (TODO ADD GITHUB), qui fournit des noyaux rapides pour le XLSTM.

Modèles du papier xlstm

Cette section explique comment utiliser les modèles à partir du papier XLSTM.

Usage

Pour les applications non linguistiques ou pour l'intégration dans d'autres architectures, vous pouvez utiliser le xLSTMBlockStack et pour la modélisation de la langue ou d'autres applications à jeton, vous pouvez utiliser le xLSTMLMModel .

pile de blocs xlstm

Le xLSTMBLockStack est destiné à une utilisation comme squelette alternative dans les projets existants. Il est similaire à une pile de blocs de transformateurs, mais utilise des blocs 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 vous travaillez avec des chaînes / fichiers YAML pour la configuration, vous pouvez également utiliser la dacite pour créer les classes de configuration. C'est la même chose que l'extrait ci-dessus:

 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

modèle de langue xlstm

Le xLSTMLMModel est un emballage autour du xLSTMBlockStack qui ajoute l'incorporation de jeton et la tête 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 )

Expériences

Les expériences synthétiques montrent que les avantages du SLSTM par rapport à MLSTM et vice versa Best sont la tâche de parité et la tâche de rappel associative multi-Quey. La tâche de parité ne peut être résolue qu'avec des capacités de suivi de l'État fournies par le mélange de mémoire de SLSTM. La tâche de rappel associative multi-requier mesure les capacités de mémorisation, où l'expansion de la mémoire matricielle et de l'État de MLSTM est très bénéfique. En combinaison, ils réussissent bien sur les deux tâches.

Pour exécuter chacun, exécutez le main.py dans le dossier des expériences comme:

 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

Notez que la boucle de formation ne contient pas d'évaluation d'arrêt ou de test précoce.

Citation

Si vous utilisez cette base de code ou si vous trouvez notre travail précieux, veuillez citer le papier 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}, 
}