فهم كيفية عمل SSMs وMamba، بالإضافة إلى كيفية البدء في تنفيذها في Keras وTensorFlow.
قدمت الورقة في 1 ديسمبر 2023 على موقع arXiv، بعنوان “مامبا: نمذجة تسلسل الزمن الخطي مع مساحات الحالة الانتقائية” نهجًا مثيرًا للاهتمام لنمذجة التسلسل. قدم المؤلفان – ألبرت جو، وتري داو – “مامبا” الذي يستخدم نماذج فضاء الحالة “الانتقائية” (SSM) لتحقيق نتائج تتنافس مع أداء نموذج المحولات المنتشر الآن في كل مكان.
شهدت المحولات شعبية مؤخرًا مع ظهور نماذج اللغات الكبيرة (LLMs) مثل LLaMa-2، وGPT-4، وClaude، وGemini، وما إلى ذلك، ولكنها تعاني من مشكلة نافذة السياق. تكمن مشكلة المحولات في جوهرها، وهو آلية الانتباه المتعددة للرأس.
المشكلة الرئيسية في الاهتمام متعدد الرؤوس تنبع من حقيقة أنه بالنسبة لطول تسلسل الإدخال n، فإن التعقيد الزمني وتعقيد الفضاء يقاسان بـ O(n²). وهذا يحد من طول نافذة السياق الخاصة بـ LLM. لأنه لزيادة هذه القيمة بمقدار 10x، نحتاج إلى زيادة متطلبات الأجهزة (أبرزها GPU VRAM) بمقدار 100x.
مامبا، من ناحية أخرى، يتقدم يا(ن)!، أي خطيًا.
هذا القياس الخطي هو ما دفع الباحثين إلى التكهن بأن مامبا قد يكون مستقبل نمذجة التسلسل.
جوهر نموذج مامبا يأتي من مفهوم نماذج الدولة الفضائية. تقوم نماذج مساحة الحالة، مثل Transformers وRNN، بمعالجة تسلسلات المعلومات، مثل النص والإشارات الصوتية وإطارات الفيديو وتسلسلات الحمض النووي وما إلى ذلك.
تأتي نماذج مساحة الحالة من فكرة وصف النظام المادي كمجموعة من المدخلات والمخرجات والمتغيرات. هذه المتغيرات هي: ا ب ت ث. تتضمن عملية SSM حساب ناقل الحالة الداخلية h
ح (ر) غالبًا ما يطلق عليها الحالة “المخفية” أو “الكامنة”، وسألتزم بتسميتها بالحالة “المخفية” لمزيد من الوضوح. من المهم ملاحظة أن A وB وC وD هي معلمات مستفادة في SSM.
ما هي المتغيرات؟
المتغيرات A وB وC وD هي معلمات متعلمة، ويمكن وصفها بأنها:
- ج: ما هو المقدار الذي يجب أن تؤخذ فيه الحالة المخفية السابقة (ح) في الاعتبار لحساب الحالة المخفية الجديدة
- ب: ما المقدار الذي يجب مراعاته في الإدخال (x) لحساب الحالة المخفية الجديدة.
- ج: ما هو المقدار الذي يجب مراعاته في الحالة المخفية الجديدة عند حساب الإخراج (y).
- د: ما هو المقدار الذي يجب أخذه في الاعتبار عند حساب المدخلات (x) عند حساب المخرجات (y).
يأتي D في نهاية الحسابات ولا يؤثر على كيفية حساب الحالة المخفية. وبالتالي، عادةً ما يتم اعتباره خارج ssm، ويمكن اعتباره بمثابة اتصال تخطي.
الانتقال من المساحات المستمرة إلى المساحات المنفصلة
تنطبق الصيغة المذكورة أعلاه على نظام ينتمي فيه الإدخال والإخراج إلى مساحة مستمرة. ولكن في حالات، مثل نمذجة اللغة، حيث ينتمي الإدخال والإخراج إلى مساحات منفصلة (قيم رمزية في المفردات). وأيضا العثور على ح (ر) يمثل تحديا تحليليا. ويمكن تحقيق ذلك عن طريق إجراء تعليق الطلب صفر.
في حالة تعليق الطلب الصفري، في كل مرة يتم فيها تلقي إدخال، يحتفظ النموذج بقيمته حتى يتم تلقي الإدخال التالي. وهذا يؤدي إلى مساحة إدخال مستمرة.
يتم تحديد طول “التعليق” هذا بواسطة معلمة جديدة تسمى، حجم الخطوة ∆. يمكن اعتباره دقة الإدخال. من الناحية المثالية، ∆ يجب أن تكون متناهية الصغر.
رياضياً، يمكن وصف تعليق الطلب الصفري على النحو التالي:
وأخيرا، يمكننا إنشاء SSM منفصلة، على النحو التالي:
نظرًا لاستخدام D مع اتصال تخطي خارج SSM، يمكن تقليل الإخراج إلى:
في أجهزة SSM، يتم نقل الحالة المخفية إلى وقت تلقي الإدخال التالي. وهذا مشابه لكيفية عمل الشبكات العصبية المتكررة.
يمكن فك هذا التنسيق المتكرر لـ SSM، تمامًا مثل شبكات RNN. ولكن على عكس شبكات RNN، التي تتسم بالتكرار والبطء، يمكن لـ SSM معالجة تسلسل الإدخال بالتوازي (تمامًا مثل المحولات) وهذا يجعل عمليات التدريب أسرع.
لاحظ أنه يتم استخدام “D” في اتصال التخطي، وهو خارج SSM.
الفكرة الأساسية حول كيفية جعل SSM التدريب سريعًا هي استخدام المتغيرات أ، ب، ج في نواة تلافيفية محسوبة مسبقًا. كتب مارتن جروتندورست شرحًا جيدًا حقًا لكيفية إنشاء هذه النواة “الالتفافية” الأساسية. ولكن هنا تفسير رياضي بسيط.
النظر في الإخراج ذ. لطول تسلسل ك، الإخراج ل ص (ك) سيتم تمثيلها (بافتراض h0 = صفر):
وبالمثل، يمكن تمثيل y3 على النحو التالي:
باستقراء النمط، يمكن تمثيل yk على النحو التالي:
يمكن تقليل هذه الصيغة أيضًا إلى:
يمثل رمز الضرب ذو المظهر المضحك عملية تلافيفية، حيث تكون نواة التلافيف هي K. لاحظ أن K لا تعتمد على س, وبالتالي يمكن حساب K مسبقًا في نواة تلافيفية، مما يجعل العملية أسرع.
وبقدر جودة القدرة الحسابية لأصوات SSM، فقد تبين أنها جميلة مه في مقاييس مثل الدقة مقارنة بالمحولات.
تكمن المشكلة الأساسية في المتغيرات ∆ وA وB وC. وتبين أنه بما أننا نطبق نفس المصفوفات على كل مدخلات، فإنها لا تستطيع حقًا معالجة سياق التسلسل.
أجهزة SSM غير مرنة في طريقة معالجة البيانات[4]
إذًا ما الذي يميز مامبا؟ في mamba، نستخدم عملية تسمى SSM “الانتقائية”، حيث يتم حساب المتغيرات ∆ وB وC بناءً على المدخلات. 🤔. نقوم بذلك عن طريق تمرير المدخلات الحالية عبر الطبقات الخطية، ونأخذ المخرجات على أنها ∆ وB وC.
ولكن هذا يجعل مدخلات ∆ وB وC معتمدة، مما يعني أنه لا يمكن حسابها مسبقًا 😢، ولن يعمل الالتفاف السريع هنا. لكن المؤلفين يناقشون الطريقة التي تعتمد على المسح النقابي الموازي.
المسح النقابي الموازي
يعد المسح النقابي الموازي تقنية قوية تستخدم في الحوسبة المتوازية لإجراء عملية مجموع البادئة، وهي عملية تراكمية على سلسلة من الأرقام. هذه العملية “ترابطية”، مما يعني أن الطريقة التي يتم بها تجميع الأرقام في العملية لا تغير النتيجة.
في سياق نموذج مامبا، من خلال تحديد العامل النقابي، يتم الحصول على العناصر والعوامل الارتباطية لعملية المسح النقابي المتوازية. وهذا يسمح بحل المشاكل على كامل الفاصل الزمني بالتوازي، مما يؤدي إلى تعقيد الوقت اللوغاريتمي في عدد الفترات الفرعية.
خوارزمية تدرك الأجهزة
إلى جانب الفحص الترابطي، يقترح المؤلفون أيضًا خوارزمية تدرك الأجهزة، حيث يستخدمون المراوغات داخل وحدات معالجة الرسومات Nvidia المتعلقة بسرعة HBM وSRAM. ويجادلون بأن حساب حالات SSM يمكن تسريعه من خلال:
- الحفاظ على الحالة المخفية و A في السعة الأسرع ولكن الأقل سرام,
- أثناء الحوسبة ∆ وB وC بسعة أبطأ ولكن أكبر إتش بي إم.
- ثم يقومون بنقل ∆ وB وC إلى سرام، حساب الحالة المخفية الجديدة في الداخل سرام.
- ثم اكتب ∆، B & C مرة أخرى إلى إتش بي إم.
في قسم التنفيذ، لن أناقش كيفية العمل مع الخوارزمية المدركة للأجهزة، بل سأستخدم فقط الفحص النقابي المتوازي.
مع أخذ كل هذا في الاعتبار، دعونا نستكشف وننفذ بنية Mamba باستخدام Keras وTensorFlow.
يمكن تقسيم بنية Mamba، بعد قراءة البحث وتحليل الكود، إلى بعض المكونات الرئيسية المرتبطة على النحو التالي:
تتكون بنية مامبا من طبقات متعددة مكدسة من “كتل مامبا”. والتي، انطلاقا من الرسم التوضيحي أعلاه، تتكون من عدد لا بأس به من المكونات. شيء آخر مهم يجب ملاحظته هو أن المؤلفين يضيفون الإخراج من SSM الانتقائي إلى الإدخال الأصلي ثم يطبقون a تطبيع طبقة إليها. يمكن أن يكون هذا التسوية إما تسوية الطبقة أو تسوية RMS.
لنبدأ بجزء من برمجة Mamba. سوف نستخدم التبعيات التالية:
tensorflow[and-cuda]==2.15.0.post1 # if you want to use GPU or
tensorflow==2.15.0.post1 # if you want to only use CPU
transformers==4.36.2 # for using the bert tokenizer
einops==0.7.0 # useful to make matrix manipulation faster
datasets==2.16.1 # to load datasets
# all other modules (like numpy) will be auto installed
الواردات:
import tensorflow_datasets as tfds
import tensorflow as tffrom tensorflow import keras
from tensorflow.keras import layers, Model
from dataclasses import dataclass
from einops import rearrange, repeat
from typing import Union
from transformers import AutoTokenizer
import datasets
import math
import numpy as np
لتسهيل معالجة وسيطة النمذجة، دعونا ننشئ عملية بسيطة ModelArgs فئة البيانات كفئة التكوين. يتيح لنا ذلك تمرير متغير فئة البيانات في الوسائط عندما نقوم بتهيئة النموذج.
@dataclass
class ModelArgs:
model_input_dims: int = 64
model_states: int = 64
projection_expand_factor: int = 2
conv_kernel_size: int = 4
delta_t_min: float = 0.001
delta_t_max: float = 0.1
delta_t_scale: float = 0.1
delta_t_init_floor: float = 1e-4
conv_use_bias: bool = True
dense_use_bias: bool = False
layer_id: int = -1
seq_length: int = 128
num_layers: int = 5
dropout_rate: float = 0.2
use_lm_head: float = False
num_classes: int = None
vocab_size: int = None
final_activation = None
loss:Union[str, keras.losses.Loss] = None
optimizer: Union[str, keras.optimizers.Optimizer] = keras.optimizers.AdamW()
metrics = ['accuracy']def __post_init__(self):
self.model_internal_dim: int = int(self.projection_expand_factor * self.model_input_dims)
self.delta_t_rank = math.ceil(self.model_input_dims/16)
if self.layer_id == -1:
self.layer_id = np.round(np.random.randint(0, 1000), 4)
if self.vocab_size == None:
raise ValueError("vocab size cannot be none")
if self.use_lm_head:
self.num_classes=self.vocab_size
else:
if self.num_classes == None:
raise ValueError(f'num classes cannot be {self.num_classes}')
if self.num_classes == 1:
self.final_activation = 'sigmoid'
else:
self.final_activation = 'softmax'
if self.loss == None:
raise ValueError(f"loss cannot be {self.loss}")
حمل ال قاعدة بيرت غير مغطاة رمز مميز:
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
vocab_size = tokenizer.vocab_size
قبل أن ننفذ فئتي Mamba وSSM، نحتاج إلى تنفيذ الفحص النقابي الموازي، ويبدو الكود كما يلي:
def selective_scan(u, delta, A, B, C, D):
# first step of A_bar = exp(ΔA), i.e., ΔA
dA = tf.einsum('bld,dn->bldn', delta, A)
dB_u = tf.einsum('bld,bld,bln->bldn', delta, u, B)dA_cumsum = tf.pad(
dA[:, 1:], [[0, 0], [1, 1], [0, 0], [0, 0]])[:, 1:, :, :]
dA_cumsum = tf.reverse(dA_cumsum, axis=[1]) # Flip along axis 1
# Cumulative sum along all the input tokens, parallel prefix sum,
# calculates dA for all the input tokens parallely
dA_cumsum = tf.math.cumsum(dA_cumsum, axis=1)
# second step of A_bar = exp(ΔA), i.e., exp(ΔA)
dA_cumsum = tf.exp(dA_cumsum)
dA_cumsum = tf.reverse(dA_cumsum, axis=[1]) # Flip back along axis 1
x = dB_u * dA_cumsum
# 1e-12 to avoid division by 0
x = tf.math.cumsum(x, axis=1)/(dA_cumsum + 1e-12)
y = tf.einsum('bldn,bln->bld', x, C)
return y + u * D
بهذا يمكننا تنفيذ MambaBlock:
class MambaBlock(layers.Layer):
def __init__(self, modelargs: ModelArgs, *args, **kwargs):
super().__init__(*args, **kwargs)
self.args = modelargs
args = modelargs
self.layer_id = modelargs.layer_idself.in_projection = layers.Dense(
args.model_internal_dim * 2,
input_shape=(args.model_input_dims,), use_bias=False)
self.conv1d = layers.Conv1D(
filters=args.model_internal_dim,
use_bias=args.conv_use_bias,
kernel_size=args.conv_kernel_size,
groups=args.model_internal_dim,
data_format="channels_first",
padding='causal'
)
# this layer takes in current token 'x'
# and outputs the input-specific Δ, B, C (according to S6)
self.x_projection = layers.Dense(args.delta_t_rank + args.model_states * 2, use_bias=False)
# this layer projects Δ from delta_t_rank to the mamba internal
# dimension
self.delta_t_projection = layers.Dense(args.model_internal_dim,
input_shape=(args.delta_t_rank,), use_bias=True)
self.A = repeat(
tf.range(1, args.model_states+1, dtype=tf.float32),
'n -> d n', d=args.model_internal_dim)
self.A_log = tf.Variable(
tf.math.log(self.A),
trainable=True, dtype=tf.float32,
name=f"SSM_A_log_{args.layer_id}")
self.D = tf.Variable(
np.ones(args.model_internal_dim),
trainable=True, dtype=tf.float32,
name=f"SSM_D_{args.layer_id}")
self.out_projection = layers.Dense(
args.model_input_dims,
input_shape=(args.model_internal_dim,),
use_bias=args.dense_use_bias)
def call(self, x):
"""Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba pape.
Official Implementation:
class Mamba, https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L119
mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
"""
(batch_size, seq_len, dimension) = x.shape
x_and_res = self.in_projection(x) # shape = (batch, seq_len, 2 * model_internal_dimension)
(x, res) = tf.split(x_and_res,
[self.args.model_internal_dim,
self.args.model_internal_dim], axis=-1)
x = rearrange(x, 'b l d_in -> b d_in l')
x = self.conv1d(x)[:, :, :seq_len]
x = rearrange(x, 'b d_in l -> b l d_in')
x = tf.nn.swish(x)
y = self.ssm(x)
y = y * tf.nn.swish(res)
return self.out_projection(y)
def ssm(self, x):
"""Runs the SSM. See:
- Algorithm 2 in Section 3.2 in the Mamba paper
- run_SSM(A, B, C, u) in The Annotated S4
Official Implementation:
mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
"""
(d_in, n) = self.A_log.shape
# Compute ∆ A B C D, the state space parameters.
# A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
# ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
# and is why Mamba is called **selective** state spaces)
A = -tf.exp(tf.cast(self.A_log, tf.float32)) # shape -> (d_in, n)
D = tf.cast(self.D, tf.float32)
x_dbl = self.x_projection(x) # shape -> (batch, seq_len, delta_t_rank + 2*n)
(delta, B, C) = tf.split(
x_dbl,
num_or_size_splits=[self.args.delta_t_rank, n, n],
axis=-1) # delta.shape -> (batch, seq_len) & B, C shape -> (batch, seq_len, n)
delta = tf.nn.softplus(self.delta_t_projection(delta)) # shape -> (batch, seq_len, model_input_dim)
return selective_scan(x, delta, A, B, C, D)
وأخيرًا، هناك كتلة متبقية لتنفيذ اتصال التخطي الخارجي.
class ResidualBlock(layers.Layer):
def __init__(self, modelargs: ModelArgs, *args, **kwargs):
super().__init__(*args, **kwargs)
self.args = modelargs
self.mixer = MambaBlock(modelargs)
self.norm = layers.LayerNormalization(epsilon=1e-5)def call(self, x):
"""
Official Implementation:
Block.forward(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L297
Note: the official repo chains residual blocks that look like
[Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> ...
where the first Add is a no-op. This is purely for performance reasons as this
allows them to fuse the Add->Norm.
We instead implement our blocks as the more familiar, simpler, and numerically equivalent
[Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> ....
"""
return self.mixer(self.norm(x)) + x
وبهذا يمكننا تهيئة نموذجنا. في هذا المثال، سأوضح كيفية استخدام كتلة Mamba لإنشاء نموذج تصنيف بسيط، ولكن يمكن تعديله بسهولة ليصبح نموذجًا للغة. دعونا تحميل يستعرض موقع IMDB مجموعة البيانات لمصنف مشاعر بسيط.
from datasets import load_dataset
from tqdm import tqdmdataset = load_dataset("ajaykarthick/imdb-movie-reviews")
نقوم أولاً بإنشاء دالة تأخذ وسيطات النموذج وتعيد نموذجًا.
def init_model(args: ModelArgs):
input_layer = layers.Input(shape=(args.seq_length,), name="input_ids")
x = layers.Embedding(
args.vocab_size,
args.model_input_dims,
input_length=args.seq_length)(input_layer)for i in range(args.num_layers):
x = ResidualBlock(args, name=f"Residual_{i}")(x)
x = layers.Dropout(args.dropout_rate)(x) # for regularization
x = layers.LayerNormalization(epsilon=1e-5)(x) # normalization layer
# use flatten only if we are not using the model as an LM
if not args.use_lm_head:
x = layers.Flatten()(x)
x = layers.Dense(1024, activation=tf.nn.gelu)(x)
output_layer = layers.Dense(
args.num_classes,
activation=args.final_activation)(x)
model = Model(
inputs=input_layer,
outputs=output_layer, name="Mamba_ka_Mamba")
model.compile(
loss=args.loss,
optimizer=args.optimizer,
metrics=args.metrics
)
return model
يمكننا الآن تهيئة نموذجنا وتلخيصه:
args = ModelArgs(
model_input_dims=128,
model_states=32,
num_layers=12,
dropout_rate=0.2,
vocab_size=vocab_size,
num_classes=1,
loss="binary_crossentropy",
)
model = init_model(args)
model.summary()
Model: "Mamba_ka_Mamba"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_ids (InputLayer) [(None, 128)] 0 embedding_2 (Embedding) (None, 128, 128) 3906816
Residual_0 (ResidualBlock) (None, 128, 128) 129024
dropout_24 (Dropout) (None, 128, 128) 0
Residual_1 (ResidualBlock) (None, 128, 128) 129024
dropout_25 (Dropout) (None, 128, 128) 0
... (I have shrinked this to make it more readable)
dropout_35 (Dropout) (None, 128, 128) 0
layer_normalization_38 (La (None, 128, 128) 256
yerNormalization)
flatten_2 (Flatten) (None, 16384) 0
dense_148 (Dense) (None, 1024) 16778240
dense_149 (Dense) (None, 1) 1025
=================================================================
Total params: 22234625 (84.82 MB)
Trainable params: 22234625 (84.82 MB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________
لتسهيل المعالجة، دعونا نقوم بترميز بياناتنا مسبقًا في ملف صفائف numpy، ثم قم بتحويلها إلى كائنات tf.data.Dataset:
train_labels, test_labels = [], []
train_ids = np.zeros((len(dataset['train']), args.seq_length))
test_ids = np.zeros((len(dataset['test']), args.seq_length))for i, item in enumerate(tqdm(dataset['train'])):
text = item['review']
train_ids[i, :] = tokenizer.encode_plus(
text,
max_length=args.seq_length,
padding='max_length',
return_tensors="np")['input_ids'][0][:args.seq_length]
train_labels.append(item['label'])
for i, item in enumerate(tqdm(dataset['test'])):
text = item['review']
test_ids[i, :] = tokenizer.encode_plus(
text,
max_length=args.seq_length,
padding='max_length',
return_tensors="np")['input_ids'][0][:args.seq_length]
test_labels.append(item['label'])
del dataset # delete the original dataset to save some memory
BATCH_SIZE = 32
train_dataset = tf.data.Dataset.from_tensor_slices((train_ids, train_labels)).batch(BATCH_SIZE).shuffle(1000)
test_dataset = tf.data.Dataset.from_tensor_slices((test_ids, test_labels)).batch(BATCH_SIZE).shuffle(1000)
الآن يمكن تدريب النموذج:
history = model.fit(train_dataset, validation_data=test_dataset, epochs=10)
يمكنك اللعب باستخدام خوارزمية الاستدلال:
def infer(text: str, model: Model, tokenizer):
tokens = tokenizer.encode(
"Hello what is up",
max_length=args.seq_length,
padding='max_length', return_tensors="np")
output = model(tokens)[0, 0]
return output
يمكن تحويل هذا النموذج إلى نموذج لغة وخوارزميات مثل البحث عن الشعاع، وأخذ العينات من أعلى مستوى، وأخذ العينات الجشعة، وما إلى ذلك. يمكن استخدامها لتوليد اللغة.
يمكن العثور على هذا الرمز على جيثب الخاص بي.
الكثير من التعليمات البرمجية مستوحاة من التنفيذ الرسمي للمامبا[2] وتطبيق آخر لـ pytorch يسمى “mamba-tiny”[3]
شكرا لقرائتك.
