xlstm

XLSTM: Erweiterter langer Kurzzeitgedächtnis

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

Um

XLSTM ist eine neue wiederkehrende neuronale Netzwerkarchitektur, die auf Ideen des ursprünglichen LSTM basiert. Durch exponentielles Gating mit angemessener Normalisierung und Stabilisierungstechniken und einem neuen Matrix -Speicher überwindet es die Einschränkungen des ursprünglichen LSTM und zeigt eine vielversprechende Leistung bei der Sprachmodellierung im Vergleich zu Transformatoren oder Zustandsraummodellen.

xlstm groß 7b

Wir haben ein 7B -Parameter XLSTM -Sprachmodell trainiert

Wir haben die XLSTM -Architektur in Bezug auf den Trainingsdurchsatz und die Stabilität optimiert. Der Code für die aktualisierte Architektur befindet sich in xlstm/xlstm_large .

Die Modellgewichte sind auf Huggingface unter https://huggingface.co/nx-ai/xlstm-7b erhältlich.

Minimale Installation

Erstellen Sie eine Conda -Umgebung aus der environment_pt220cu121.yaml . Installieren Sie nur den Modellcode (dh das Modul xlstm ) als Paket:

Installieren Sie über PIP:

pip install xlstm

Klon aus Github:

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

Zur Verwendung des 7B XLSTM -Modells installieren Sie mlstm_kernels über:

 pip install mlstm_kernels

Anforderungen

Dieses Paket basiert auf Pytorch und wurde auf Versionen >=1.8 getestet. Für die CUDA-Version von SLSTM benötigen Sie Rechnungsfunktionen> = 8.0, siehe https://developer.nvidia.com/cuda-gpus. Installieren Sie für eine gut getestete Umgebung die environment_pt220cu121.yaml AS:

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

Für das XLSTM Large 7B -Modell benötigen wir unser Paket mlstm_kernels (Todo GitHub Link), das schnelle Kernel für das XLSTM liefert.

Modelle aus dem XLSTM -Papier

In diesem Abschnitt wird erläutert, wie die Modelle aus dem XLSTM -Papier verwendet werden.

Verwendung

Für nichtsprachige Anwendungen oder für die Integration in andere Architekturen können Sie den xLSTMBlockStack und für Sprachmodellierung oder andere tokenbasierte Anwendungen verwenden, können Sie das xLSTMLMModel verwenden.

xlstm Blockstapel

Der xLSTMBLockStack ist für die Verwendung als alternatives Rückgrat in vorhandenen Projekten gedacht. Es ähnelt einem Stapel von Transformatorblöcken, verwendet jedoch XLSTM -Blöcke:

 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

Wenn Sie mit YAML -Zeichenfolgen / -dateien zur Konfiguration arbeiten, können Sie auch DACITE zum Erstellen der KonfigurationsdataClasses verwenden. Dies ist das gleiche wie der Snippet oben:

 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 -Sprachmodell

Das xLSTMLMModel ist ein Wrapper um den xLSTMBlockStack , der den Token -Einbettung und den LM -Kopf hinzufügt.

 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 )

Experimente

Die synthetischen Experimente zeigen, dass die Vorteile von SLSTM gegenüber MLSTM und umgekehrt am besten die Paritätsaufgabe und die Multiquery Associativ-Rückrufaufgabe sind. Die Paritätsaufgabe kann nur mit staatlichen Funktionen gelöst werden, die durch die Speichermischung von SLSTM bereitgestellt werden. Die Multiquery Associativ-Rückruf-Aufgaben misst die Auswendigmachungsfunktionen, bei denen das Matrixemory und die staatliche Expansion von MLSTM sehr vorteilhaft sind. In Kombination machen sie sich bei beiden Aufgaben gut.

Um jeweils auszuführen, führen Sie die main.py im experimentenordner wie folgt aus:

 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

Beachten Sie, dass die Trainingsschleife keine frühzeitige Stopp- oder Testbewertung enthält.

Zitat

Wenn Sie diese Codebasis verwenden oder unsere Arbeit auf andere Weise wertvoll finden, zitieren Sie bitte das XLSTM -Papier:

 @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}, 
}