赞
踩
在深度学习领域,序列建模仍然是一项具有挑战性的任务,通常由 LSTM 和 Transformers 等模型来解决。然而,这些模型的计算量很大,因此在实际应用场景中,这些模型方法仍存在巨大的缺陷。而Mamba 是一个线性时间序列建模框架,其旨在提高序列建模的效率和有效性。本博文将深入探讨使用 PyTorch 实现 Mamba 的过程,讨论这一创新方法背后的技术问题和代码
1.1 Transformer:
Transformer因其注意机制而闻名。借助于Transformer操作特性,特征序列中的任何部分都可以与其他部分进行动态交互,尤其是因果注意力特征,能够很好的捕获因果特征的信息。因此,Transformer能够处理好序列中的每一个元素,相应的,Transformer的计算代价和内存成本也都很高,与序列长度(L²)的平方成正比。
1.2 递归神经网络(RNN):
RNN 是按照序列顺序更新隐藏状态,它只考虑当前输入特征和上一个隐藏状态信息。这种方法允许它们以恒定的内存成本处理无限长的序列。然而,RNN 的简单性也变相的成为一个缺点,即限制了其记忆长期依赖关系的能力。此外,尽管有 LSTM 这样的创新,RNN 中的时间反向传播(BPTT)机制可能会占用大量内存,并可能出现梯度消失或爆炸的问题。
1.3 状态空间模型(S4):
状态空间模型具有良好的特性。它们提供了一种计算代价和内存成本的平衡,比 RNNs 更有效地捕捉长程依赖性,同时比 Transformers 更节省内存。
图1|序列建模网络架构发展©️【深蓝AI】
1.4 Mamba架构的方法思路:
●选择性状态空间: Mamba 以状态空间模型的概念为基础,引入了一种新的模型架构设计思路。它利用选择性状态空间,能更高效、更有效地捕捉长序列中的相关信息。
●线性时间复杂性: 与 Transformers 不同,Mamba 的运行时间与序列长度成线性关系。这一特性使其特别适用于超长序列的任务,而传统的模型在这方面会很吃力。
图2|Mamba引入选择性状态空间©️【深蓝AI】
Mamba 通过其 “选择性状态空间”(Selective State Spaces)的概念,为传统的状态空间模型引入了一个新颖的架构。这种方法稍微放宽了标准状态空间模型的僵化状态转换,使其更具适应性和灵活性,有点类似于 LSTM。不过,Mamba 保留了状态空间模型的高效计算特性,使其能够一次性完成整个序列的前向传递。
2.1 导入必须的库文件
在简单介绍完Mamba架构之后,为大家带来Mamba的代码实现过程,首先导入必须的库。
# PyTorch相关的库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn import functional as F
from einops import rearrange
from tqdm import tqdm
# 系统相关的库
import math
import os
import urllib.request
from zipfile import ZipFile
from transformers import AutoTokenizer
torch.autograd.set_detect_anomaly(True)
2.2 设置标识和训练设备
这里主要针对是否使用GPU,以及Mamba的选择设定对应的表示、以及所使用的设备。
# 配置标识和超参数
USE_MAMBA = 1
DIFFERENT_H_STATES_RECURRENT_UPDATE_MECHANISM = 0
# 设定所用设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
2.3 设置初始化超参数
这一小节定义了模型维度(d_model)、状态大小、序列长度和批次大小等超参数。
# 人为定义的超参数
d_model = 8
state_size = 128 # 状态大小
seq_len = 100 # 序列长度
batch_size = 256 # 批次大小
last_batch_size = 81 # 最后一个批次大小
current_batch_size = batch_size
different_batch_size = False
h_new = None
temp_buffer = None
2.4 定义S6模块
S6 模块是 Mamba 架构中的一个复杂组件,它主要由一系列线性变换和离散化过程组成,用于处理输入的特征序列。它在捕捉序列的时间动态特征方面起着至关重要的作用,而时间动态特征是语言建模等序列建模任务的一个关键方面。
# 定义S6模块 class S6(nn.Module): def __init__(self, seq_len, d_model, state_size, device): super(S6, self).__init__() # 一系列线性变换 self.fc1 = nn.Linear(d_model, d_model, device=device) self.fc2 = nn.Linear(d_model, state_size, device=device) self.fc3 = nn.Linear(d_model, state_size, device=device) # 设定一些超参数 self.seq_len = seq_len self.d_model = d_model self.state_size = state_size self.A = nn.Parameter(F.normalize(torch.ones(d_model, state_size, device=device), p=2, dim=-1)) # 参数初始化 nn.init.xavier_uniform_(self.A) self.B = torch.zeros(batch_size, self.seq_len, self.state_size, device=device) self.C = torch.zeros(batch_size, self.seq_len, self.state_size, device=device) self.delta = torch.zeros(batch_size, self.seq_len, self.d_model, device=device) self.dA = torch.zeros(batch_size, self.seq_len, self.d_model, self.state_size, device=device) self.dB = torch.zeros(batch_size, self.seq_len, self.d_model, self.state_size, device=device) # 定义内部参数h和y self.h = torch.zeros(batch_size, self.seq_len, self.d_model, self.state_size, device=device) self.y = torch.zeros(batch_size, self.seq_len, self.d_model, device=device) # 离散化函数 def discretization(self): # 离散化函数定义介绍在Mamba论文中的28页 self.dB = torch.einsum("bld,bln->bldn", self.delta, self.B) #dA = torch.matrix_exp(A * delta) # matrix_exp() only supports square matrix self.dA = torch.exp(torch.einsum("bld,dn->bldn", self.delta, self.A)) #print(f"self.dA.shape = {self.dA.shape}") #print(f"self.dA.requires_grad = {self.dA.requires_grad}") return self.dA, self.dB # 前行传播 def forward(self, x): # 参考Mamba论文中算法2 self.B = self.fc2(x) self.C = self.fc3(x) self.delta = F.softplus(self.fc1(x)) # 离散化 self.discretization() if DIFFERENT_H_STATES_RECURRENT_UPDATE_MECHANISM: # 如果不使用'h_new',将触发本地允许错误 global current_batch_size current_batch_size = x.shape[0] if self.h.shape[0] != current_batch_size different_batch_size = True # 缩放h的维度匹配当前的批次 h_new = torch.einsum('bldn,bldn->bldn', self.dA, self.h[:current_batch_size, ...]) + rearrange(x, "b l d -> b l d 1") * self.dB else: different_batch_size = False h_new = torch.einsum('bldn,bldn->bldn', self.dA, self.h) + rearrange(x, "b l d -> b l d 1") * self.dB # 改变y的维度 self.y = torch.einsum('bln,bldn->bld', self.C, h_new) # 基于h_new更新h的信息 global temp_buffer temp_buffer = h_new.detach().clone() if not self.h.requires_grad else h_new.clone() return self.y else: # 将会触发错误 # 设置h的维度 h = torch.zeros(x.size(0), self.seq_len, self.d_model, self.state_size, device=x.device) y = torch.zeros_like(x) h = torch.einsum('bldn,bldn->bldn', self.dA, h) + rearrange(x, "b l d -> b l d 1") * self.dB # 设置y的维度 y = torch.einsum('bln,bldn->bld', self.C, h) return y
S6 模块继承于 nn.Module,是 Mamba 算法模型的关键部分,负责处理离散化过程和前向传播。
2.5 定义MambaBlock模块
MambaBlock 模块是一个定制的神经网络模块,是 Mamba 模型的关键部件,它封装了处理输入数据的多个网络层和操作函数。MambaBlock 模块代表一个复杂的神经网络模块,包括线性投影、卷积、激活函数、自定义 S6 模块和残差连接。该模块是 Mamba 模型的基本组成部分,通过一系列转换处理输入序列,以捕捉数据中的相关模式和特征。这些不同网络层和操作函数的组合使 MambaBlock 能够有效处理复杂的序列建模任务。
# 定义MambaBlock模块 class MambaBlock(nn.Module): def __init__(self, seq_len, d_model, state_size, device): super(MambaBlock, self).__init__() self.inp_proj = nn.Linear(d_model, 2*d_model, device=device) self.out_proj = nn.Linear(2*d_model, d_model, device=device) # 残差连接 self.D = nn.Linear(d_model, 2*d_model, device=device) # 设置偏差属性 self.out_proj.bias._no_weight_decay = True # 初始化偏差 nn.init.constant_(self.out_proj.bias, 1.0) # 初始化S6模块 self.S6 = S6(seq_len, 2*d_model, state_size, device) # 添加1D卷积 self.conv = nn.Conv1d(seq_len, seq_len, kernel_size=3, padding=1, device=device) # 添加线性层 self.conv_linear = nn.Linear(2*d_model, 2*d_model, device=device) # 正则化 self.norm = RMSNorm(d_model, device=device) # 前向传播 def forward(self, x): # 参考Mamba论文中的图3 x = self.norm(x) x_proj = self.inp_proj(x) # 1D卷积操作 x_conv = self.conv(x_proj) x_conv_act = F.silu(x_conv) # Swish激活 # 线性操作 x_conv_out = self.conv_linear(x_conv_act) # S6模块操作 x_ssm = self.S6(x_conv_out) x_act = F.silu(x_ssm) # Swish激活 # 残差连接 x_residual = F.silu(self.D(x)) x_combined = x_act * x_residual x_out = self.out_proj(x_combined) return x_out
MambaBlock 模块是另一个封装了 Mamba 核心功能的模块,包括输入投影、一维卷积和 S6 模块。
2.6 定义Mamba模型
Mamba 类代表 Mamba 模型的整体架构,由一系列 MambaBlock 模块组成。每个模块负责处理输入的序列数据,一个模块的输出作为下一个模块的输入。这种顺序处理使模型能够捕捉输入数据中的复杂模式和关系,从而有效地完成顺序建模的任务。多个模块的堆叠是深度学习架构中常见的设计,因为它能让模型学习数据的分层表示特征。
# 定义Mamba模型
class Mamba(nn.Module):
def __init__(self, seq_len, d_model, state_size, device):
super(Mamba, self).__init__()
self.mamba_block1 = MambaBlock(seq_len, d_model, state_size, device)
self.mamba_block2 = MambaBlock(seq_len, d_model, state_size, device)
self.mamba_block3 = MambaBlock(seq_len, d_model, state_size, device)
def forward(self, x):
x = self.mamba_block1(x)
x = self.mamba_block2(x)
x = self.mamba_block3(x)
return x
该类定义了整个 Mamba 模型,将多个 MambaBlock 模块链接在一起,构成整体算法模型的架构。
2.7 定义RMSNorm模块
RMSNorm 模块是一个自定义的归一化层,继承了 PyTorch 的 nn.Module。该层用于对神经网络的激活值进行归一化操作,这有助于加快训练速度。
class RMSNorm(nn.Module):
def __init__(self,
d_model: int,
eps: float = 1e-5,
device: str ='cuda'):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(d_model, device=device))
def forward(self, x):
output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
return output
RMSNorm 模块是用于归一化的均方根网络层,是神经网络架构中的一种常用技术。
本节介绍如何在简单的数据样本上实例化和使用 Mamba 算法模型。
# 创建模拟数据
x = torch.rand(batch_size, seq_len, d_model, device=device)
# 创建Mambda算法模型
mamba = Mamba(seq_len, d_model, state_size, device)
# 定义rmsnorm模块
norm = RMSNorm(d_model)x = norm(x)
# 前向传播
test_output = mamba(x)
print(f"test_output.shape = {test_output.shape}")
3.1 数据准备和训练函数
Enwiki8Dataset 类是一个自定义数据集处理程序,它继承自 PyTorch 的 Dataset 类,专门用于为序列建模任务(如语言建模)而构建的数据集。
# 定义填充函数
def pad_sequences_3d(sequences, max_len=None, pad_value=0):
# 获得张量的维度大小
batch_size, seq_len, feature_size = sequences.shape
if max_len is None:
max_len = seq_len + 1
# 初始化 padded_sequences
padded_sequences = torch.full((batch_size, max_len, feature_size), fill_value=pad_value, dtype=sequences.dtype, device=sequences.device)
# 填充每个序列
padded_sequences[:, :seq_len, :] = sequences
return padded_sequences
train 函数用于训练 Mamba 算法模型。下面看看它的组成部分:
def train(model, tokenizer, data_loader, optimizer, criterion, device, max_grad_norm=1.0, DEBUGGING_IS_ON=False)函数需要下面几个参数:
●model(模型): 要训练的神经网络模型(本例中为 Mamba);
●tokenizer: 处理输入数据的标记符;
●data_loader:数据加载器: 一个可迭代器,用于为训练提供成批数据;
●optimizer: 优化器: 用于更新模型权重的优化算法;
●criterion: 用于评估模型性能的损失函数;
●设备: 模型运行的设备(CPU 或 GPU);
●max_grad_norm: 用于梯度剪切的值,以防止梯度爆炸;
●DEBUGGING_IS_ON:启用调试信息的标志。
# 定义train函数 def train(model, tokenizer, data_loader, optimizer, criterion, device, max_grad_norm=1.0, DEBUGGING_IS_ON=False): model.train() total_loss = 0 for batch in data_loader: optimizer.zero_grad() input_data = batch['input_ids'].clone().to(device) attention_mask = batch['attention_mask'].clone().to(device) # 获取输入数据和标签 target = input_data[:, 1:] input_data = input_data[:, :-1] # 填充序列数据 input_data = pad_sequences_3d(input_data, pad_value=tokenizer.pad_token_id) target = pad_sequences_3d(target, max_len=input_data.size(1), pad_value=tokenizer.pad_token_id) if USE_MAMBA: output = model(input_data) loss = criterion(output, target) loss.backward(retain_graph=True) # 裁剪梯度 for name, param in model.named_parameters(): if 'out_proj.bias' not in name: # 裁剪梯度函数操作 torch.nn.utils.clip_grad_norm_(param, max_norm=max_grad_norm) if DEBUGGING_IS_ON: for name, parameter in model.named_parameters(): if parameter.grad is not None: print(f"{name} gradient: {parameter.grad.data.norm(2)}") else: print(f"{name} has no gradient") if USE_MAMBA and DIFFERENT_H_STATES_RECURRENT_UPDATE_MECHANISM: model.S6.h[:current_batch_size, ...].copy_(temp_buffer) optimizer.step() total_loss += loss.item() return total_loss / len(data_loader)
3.2 模型训练循环
# 输入预训练模型权重 encoded_inputs_file = 'encoded_inputs_mamba.pt' if os.path.exists(encoded_inputs_file): print("Loading pre-tokenized data...") encoded_inputs = torch.load(encoded_inputs_file) else: print("Tokenizing raw data...") enwiki8_data = load_enwiki8_dataset() encoded_inputs, attention_mask = encode_dataset(tokenizer, enwiki8_data) torch.save(encoded_inputs, encoded_inputs_file) print(f"finished tokenizing data") # 组合数据 data = { 'input_ids': encoded_inputs, 'attention_mask': attention_mask } # 分割训练和验证集 total_size = len(data['input_ids']) train_size = int(total_size * 0.8) train_data = {key: val[:train_size] for key, val in data.items()} val_data = {key: val[train_size:] for key, val in data.items()} train_dataset = Enwiki8Dataset(train_data) val_dataset = Enwiki8Dataset(val_data) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) # 初始化模型 model = Mamba(seq_len, d_model, state_size, device).to(device) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=5e-6) # 训练次数 num_epochs = 25for epoch in tqdm(range(num_epochs)): train_loss = train(model, tokenizer, train_loader, optimizer, criterion, device, max_grad_norm=10.0, DEBUGGING_IS_ON=False) val_loss = evaluate(model, val_loader, criterion, device) val_perplexity = calculate_perplexity(val_loss) print(f'Epoch: {epoch+1}, Training Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}, Validation Perplexity: {val_perplexity:.4f}')
上面代码是建立和训练 Mamba 模型的详细示例过程,包括数据集的组合和划分,模型的定义和初始化,以及损失函数和优化器的定义,最后则是设定训练循环的次数。
本文提供了从零开始构建 Mamba 的全部代码过程,读者们可以借助本文的讲解和代码,将Mamba算法模型从理论转化为具体实践。这一探索过程不仅可以巩固对 Mamba 内部工作原理的理解,而且还展示了新颖算法模型架构的实际设计步骤。通过这种实践方法,笔者发现了序列建模的细微差别以及 Mamba 在这一领域引入的效率。有了这些知识,笔者现在就可以在自己的项目中更好地尝试使用 Mamba,或更深入地开发新型的 AI 模型。
参考:
【1】https://arxiv.org/abs/2312.00752
【2】https://github.com/state-spaces/mamba
【3】https://pytorch.org/tutorials/beginner/basics/data_tutorial.html
【4】https://huggingface.co/datasets/enwik8
笔者|匡吉
审核|Los
移步公众号【深蓝AI】,第一时间获取自动驾驶、人工智能与机器人行业最新最前沿论文和科技动态。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。