xlstm

XLSTM : 긴 단기 메모리 확장

종이 : https://arxiv.org/abs/2405.04517

에 대한

XLSTM은 원래 LSTM의 아이디어를 기반으로 한 새로운 반복 신경 네트워크 아키텍처입니다. 적절한 정규화 및 안정화 기술과 새로운 매트릭스 메모리를 사용하여 지수 게이팅을 통해 원래 LSTM의 한계를 극복하고 변압기 또는 상태 공간 모델과 비교할 때 언어 모델링에 대한 유망한 성능을 보여줍니다.

xlstm 큰 7b

우리는 7B 매개 변수 XLSTM 언어 모델을 교육했습니다

우리는 훈련 처리량 및 안정성 측면에서 XLSTM 아키텍처를 최적화했습니다. 업데이트 된 아키텍처 코드는 xlstm/xlstm_large 에 있습니다.

모델 가중치는 https://huggingface.co/nx-ai/xlstm-7b에서 huggingface에서 사용할 수 있습니다.

최소 설치

environment_pt220cu121.yaml 에서 콘다 환경을 만듭니다. 모델 코드 만 패키지로 만 설치하십시오 (예 : 모듈 xlstm ).

PIB를 통해 설치 :

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 Large 7B 모델의 경우 XLSTM에 빠른 커널을 제공하는 mlstm_kernels (TODO ADD GITHUB 링크) 패키지가 필요합니다.

XLSTM 용지의 모델

이 섹션에서는 XLSTM 용지의 모델을 사용하는 방법을 설명합니다.

용법

비 언어 애플리케이션 또는 기타 아키텍처에 통합하려면 xLSTMBlockStack 및 언어 모델링 또는 기타 토큰 기반 응용 프로그램에 xLSTMLMModel 사용할 수 있습니다.

XLSTM 블록 스택

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를 사용하여 구성 데이터 클래스를 만들 수도 있습니다. 위의 스 니펫과 동일합니다.

 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 )

실험

MLSTM에 대한 SLSTM의 이점을 보여주는 합성 실험은 패리티 작업과 다중 쿼리 연관성 리콜 작업입니다. 패리티 작업은 SLSTM의 메모리 혼합으로 제공되는 상태 추적 기능으로 만 해결할 수 있습니다. 다중 쿼리 연관성 리콜 작업은 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}, 
}