آشنایی با نحوه کار SSM و Mamba و نحوه اجرای آنها در Keras و TensorFlow.
مقاله با عنوان «Mamba: مدل‌سازی توالی‌های زمانی خطی با فضاهای حالت انتخابی» که در ۱ دسامبر ۲۰۲۳ به arXiv ارسال شد، رویکرد جالبی برای مدل‌سازی توالی ارائه می‌دهد. نویسندگان – آلبرت گو، تری دائو – “Mamba” را معرفی کردند که از مدل های فضای حالت “انتخابی” (SSM) برای دستیابی به نتایجی استفاده می کند که با عملکرد مدل ترانسفورماتور اکنون همه جا حاضر رقابت می کند.
ترانسفورماتورها اخیراً با ظهور مدل های زبان بزرگ (LLM) مانند LLaMa-2، GPT-4، Claude، Gemini و غیره محبوب شده اند، اما آنها از مشکل پنجره زمینه رنج می برند. مشکل ترانسفورماتورها در ماهیت آنها نهفته است، مکانیزم جلب توجه چند سر.
مشکل اصلی توجه چند سر از این واقعیت ناشی می شود که برای یک دنباله ورودی طول n، مقیاس پیچیدگی زمانی و فضایی با O(n²) است. این طول پنجره زمینه LLM را محدود می کند. زیرا برای افزایش آن به میزان 10 برابر، باید میزان سخت افزار مورد نیاز (بیشتر GPU VRAM) را 100 برابر کنیم.
مامبا، از سوی دیگر، ترازو O(n)!، یعنی. به صورت خطی.
این مقیاس خطی همان چیزی است که محققان را به این حدس و گمان سوق داد که مامبا می تواند آینده مدل سازی توالی باشد.
هسته مدل مامبا از مفهوم مدل های فضای حالت می آید. مدل های فضای حالت، مانند ترانسفورماتورها و RNN ها، توالی های اطلاعاتی مانند متن، سیگنال های صوتی، فریم های ویدئویی، توالی های DNA و غیره را پردازش می کنند.
مدل‌های فضای حالت از ایده توصیف یک سیستم فیزیکی به عنوان مجموعه‌ای از ورودی‌ها، خروجی‌ها و متغیرها ناشی می‌شوند. این متغیرها عبارتند از: آ ب پ ت. فرآیند SSM شامل محاسبه یک بردار حالت داخلی h
h
متغیرها کدامند؟
متغیرهای A، B، C و D پارامترهای آموخته شده هستند، و می توان آنها را اینگونه توصیف کرد:
- الف: برای محاسبه حالت پنهان جدید، وضعیت پنهان قبلی (h) چقدر باید در نظر گرفته شود
- ب: برای محاسبه حالت پنهان جدید چه مقدار ورودی (x) باید در نظر گرفته شود.
- ج: در محاسبه خروجی (y) چقدر باید حالت پنهان جدید را در نظر گرفت.
- د: در محاسبه خروجی (y) چه مقدار ورودی (x) باید در نظر گرفته شود.
D در پایان محاسبات می آید و بر نحوه محاسبه حالت پنهان تأثیری ندارد. بنابراین معمولاً خارج از ssm در نظر گرفته می شود و می توان آن را یک اتصال پرش به حساب آورد.
انتقال از فضاهای پیوسته به فضاهای گسسته
فرمول فوق برای سیستمی اعمال می شود که ورودی و خروجی آن به یک فضای پیوسته تعلق دارد. اما در مواردی مانند مدل‌سازی زبان، که ورودی و خروجی به فضاهای گسسته تعلق دارند (مقادیر رمز در فرهنگ لغت). همچنین پیدا کردن h
با نگه داشتن مرتبه صفر، هر بار که یک ورودی دریافت می شود، مدل مقدار خود را تا دریافت ورودی بعدی حفظ می کند. این منجر به یک فضای ورودی پیوسته می شود.
این طول نگه داشتن توسط یک پارامتر جدید به نام، اندازه گام ∆. می توان آن را به عنوان وضوح ورودی در نظر گرفت. در حالت ایده آل، ∆ باید بی نهایت کوچک باشد.
از نظر ریاضی، حفظ مرتبه صفر را می توان به صورت زیر توصیف کرد:
در نهایت، می‌توانیم یک SSM جداگانه ایجاد کنیم، مانند:
از آنجایی که D با اتصال پرش غیر SSM استفاده می شود، خروجی را می توان به موارد زیر کاهش داد:
در SSM، حالت پنهان تا دریافت ورودی بعدی به جلو منتقل می شود. این شبیه به نحوه عملکرد شبکه های عصبی مکرر است.
این قالب تکراری SSM را می توان درست مانند RNN باز کرد. اما بر خلاف RNN ها که تکراری و کند هستند، SSM می تواند توالی ورودی را به صورت موازی پردازش کند (درست مانند ترانسفورماتورها) و این باعث می شود فرآیندهای یادگیری سریعتر شود.
توجه داشته باشید که ‘D’ در یک پیوند پرش استفاده می شود که خارج از SSM است.
بینش کلیدی در مورد چگونگی سرعت یادگیری SSM استفاده از متغیرها است الف، ب، ج در یک هسته کانولوشنال از پیش محاسبه شده Maarten Grootendorst توضیح بسیار خوبی در مورد چگونگی ساخت این هسته متعارف “convolutional” نوشت. اما در اینجا یک توضیح ساده ریاضی وجود دارد.
به راه خروج فکر کن آقای. برای طول دنباله ای از کخروجی برای y(k) ارائه خواهد شد (با فرض h0 = صفر):
به طور مشابه، y3 را می توان به صورت زیر نشان داد:
با برون یابی مدل، yk را می توان به صورت زیر نشان داد:
این فرمول را می توان به موارد زیر کاهش داد:
نماد ضرب با ظاهر خنده دار یک عملیات کانولوشن را نشان می دهد که در آن هسته کانولوشن K است. توجه داشته باشید که K به آن بستگی ندارد. ایکس، بنابراین، K را می توان در یک هسته کانولوشن از پیش محاسبه کرد و روند را سریعتر کرد.
به همان اندازه که قدرت محاسباتی SSM به نظر می رسد، بسیار زیبا به نظر می رسد ماه در معیارهایی مانند دقت در مقایسه با ترانسفورماتورها.
مشکل اصلی در متغیرهای ∆، A، B و C نهفته است. معلوم می‌شود که از آنجایی که ماتریس‌های یکسانی را برای هر ورودی اعمال می‌کنیم، آنها واقعاً نمی‌توانند زمینه توالی را مدیریت کنند.
SSMها در نحوه مدیریت داده ها انعطاف ناپذیر هستند[4]
پس چه چیزی در مورد Mamba خاص است؟ در mamba از فرآیندی به نام SSM “انتخابی” استفاده می کنیم که در آن متغیرهای ∆، B و C بر اساس ورودی محاسبه می شوند. 🤔 این کار را با عبور دادن جریان ورودی از لایه های خطی و گرفتن خروجی برای Δ، B و C انجام می دهیم.
اما پس از آن، ورودی ∆، B و C را وابسته می‌کند، به این معنی که نمی‌توان آنها را از قبل محاسبه کرد، پیچیدگی سریع در اینجا کار نخواهد کرد. اما نویسندگان روشی را مورد بحث قرار می دهند که مبتنی بر آن است اسکن انجمنی موازی
اسکن انجمنی موازی
اسکن انجمنی موازی یک تکنیک قدرتمند است که در محاسبات موازی برای انجام عملیات جمع پیشوند، که یک عملیات تجمعی روی دنباله ای از اعداد است، استفاده می شود. این عملیات “تداعی” است، به این معنی که نحوه گروه بندی اعداد در عملیات نتیجه را تغییر نمی دهد.
در زمینه مدل Mamba، تعریف یک عملگر انجمنی عناصر و عملگرهای انجمنی را برای یک عملیات اسکن انجمنی موازی به دست می‌دهد. این اجازه می دهد تا حل مسئله موازی برای کل بازه زمانی، و در نتیجه پیچیدگی زمانی لگاریتمی در تعداد زیر بازه ها.
الگوریتم سخت افزار
همراه با اسکن انجمنی، نویسندگان همچنین یک الگوریتم آگاه از سخت افزار را پیشنهاد می کنند که در آن از ویژگی های عجیب پردازنده های گرافیکی Nvidia مربوط به سرعت HBM و SRAM استفاده می کنند. آنها ادعا می کنند که محاسبه حالت های SSM می تواند توسط موارد زیر تسریع شود:
- نگه داشتن حالت پنهان و A در ظرفیت سریعتر اما کمتر SRAM،
- در حین محاسبه Δ، B و C، در ظرفیت کندتر اما بزرگتر HBM.
- سپس Δ، B و C را به آنها منتقل می کنند SRAMحالت پنهان جدید را محاسبه کنید SRAM.
- و سپس Δ, B & C را به عقب بنویسید HBM.
در بخش پیاده سازی، من در مورد نحوه کار با الگوریتم سخت افزار صحبت نمی کنم، بلکه فقط از اسکن انجمنی موازی استفاده می کنم.
با در نظر گرفتن همه اینها، بیایید معماری Mamba را با استفاده از Keras و TensorFlow بررسی و پیاده سازی کنیم.
معماری مامبا پس از مطالعه مقاله و تجزیه و تحلیل کد، می تواند به چندین جزء کلیدی تقسیم شود که به شرح زیر است:
معماری مامبا از چندین لایه انباشته از “بلوک های مامبا” تشکیل شده است. که، با قضاوت در تصویر بالا، از چند جزء تشکیل شده است. نکته مهم دیگری که باید به آن توجه شود این است که نویسندگان خروجی Selective SSM را به ورودی اصلی اضافه کرده و سپس اعمال می کنند. عادی سازی لایه به آن. این نرمال سازی می تواند نرمال سازی لایه یا عادی سازی 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)
_________________________________________________________________
برای پردازش آسان‌تر، بیایید داده‌های خود را در یک از قبل نشانه گذاری کنیم آرایه های بی رنگسپس آنها را به اشیاء 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
این مدل را می توان به مدل زبان و الگوریتم های مشابه تبدیل کرد جستجوی پرتو، نمونه برداری top-k، نمونه برداری حریصانه و غیره. می تواند برای تولید زبان استفاده شود.
این کد را می توان در Github من یافت.
بسیاری از کدها از پیاده سازی رسمی mamba الهام گرفته شده است[2] و اجرای پایتورچ دیگری به نام «مامبا-تینی»[3]
ممنون که خواندید.
