xlstm

XLSTM: ذاكرة طويلة الأجل طويلة المدى

ورقة: https://arxiv.org/abs/2405.04517

عن

XLSTM هي بنية جديدة للشبكة العصبية المتكررة تعتمد على أفكار LSTM الأصلية. من خلال البوابات الأسية مع تقنيات التطبيع والتثبيت المناسبة وذاكرة مصفوفة جديدة ، فإنه يتغلب على قيود LSTM الأصلية ويظهر أداءًا واعد على نمذجة اللغة مقارنةً بالمحولات أو نماذج مساحة الحالة.

XLSTM كبير 7 ب

قمنا بتدريب نموذج لغة معلمة XLSTM 7B

لقد قمنا بتحسين بنية XLSTM من حيث إنتاجية التدريب والاستقرار. يقع الرمز للهندسة المعمارية المحدثة في xlstm/xlstm_large .

تتوفر أوزان النموذج على Luggingface على https://huggingface.co/nx-ai/xlstm-7b.

الحد الأدنى من التثبيت

قم بإنشاء بيئة كوندا من ملفات الملفات environment_pt220cu121.yaml . قم بتثبيت رمز النموذج فقط (أي الوحدة xlstm ) كحزمة:

التثبيت عبر PIP:

pip install xlstm

استنساخ من جيثب:

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 كتلة مكدس

يهدف 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 )

التجارب

إن التجارب الاصطناعية تُظهر فوائد SLSTM على MLSTM والعكس بالعكس هي مهمة التكافؤ ومهمة الاستدعاء الترابطية المتعددة. لا يمكن حل مهمة التكافؤ إلا من خلال إمكانات تتبع الحالة التي توفرها عملية خلط الذاكرة لـ 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}, 
}