当前位置：首页 > news >正文

MaxStateSuper模型详解与实现

news 2025/7/7 18:01:20

模型架构概述

MaxStateSuper模型的核心由三个关键组件构成：

MaxStateSuper模块：创新性的多头注意力机制变体
门控前馈网络(FeedForward)：增强非线性表达能力
解码器层(DecoderLayer)：整合注意力和前馈网络

关键组件详解

1. MaxStateSuper模块

class MaxStateSuper(torch.nn.Module):def __init__(self, dim_size, heads):super(MaxStateSuper, self).__init__()self.heads = headsassert dim_size % heads == 0, "Dimension size must be divisible by head size."# 合并三个线性层为一个，减少参数self.combined = nn.Linear(dim_size, 4 * dim_size, bias=False)# 可学习的加权参数self.alpha1 = torch.nn.Parameter(torch.tensor(0.5))self.alpha2 = torch.nn.Parameter(torch.tensor(0.5))self.alpha3 = torch.nn.Parameter(torch.tensor(0.5))self.alpha4 = torch.nn.Parameter(torch.tensor(0.5))self.alpha5 = torch.nn.Parameter(torch.tensor(0.5))self.alpha6 = torch.nn.Parameter(torch.tensor(0.5))

核心创新点：

合并线性变换：将传统Transformer中独立的Q、K、V变换合并为单一线性层，减少参数数量
累积最大值操作：引入torch.cummax沿序列维度计算累积最大值
参数化加权：使用可学习参数灵活组合不同表示

def gen_model(self, a, b, c, d, e):# 可学习参数的加权组合x = self.alpha1 * b + self.alpha2 * d + \self.alpha3 * a + self.alpha4 * c + \self.alpha5 * e# 与累积最大值表示交互x = a * e + xx = b * e + xx = c * e + xx = d * e + xx = e * self.alpha6 + xreturn x

2. 门控前馈网络

class FeedForward(torch.nn.Module):def __init__(self, hidden_size):super(FeedForward, self).__init__()self.ffn1 = torch.nn.Linear(hidden_size, hidden_size)self.ffn2 = torch.nn.Linear(hidden_size, hidden_size)self.gate = torch.nn.Linear(hidden_size, hidden_size)self.relu = torch.nn.ReLU()

门控机制数学表示：
$\cdot \text{ReLU}(gate(x)))$

3. 解码器层

class DecoderLayer(torch.nn.Module):def __init__(self, hidden_size, num_heads):super(DecoderLayer, self).__init__()self.self_attention = MaxStateSuper(hidden_size, num_heads)self.ffn = FeedForward(hidden_size)self.layer_norm = torch.nn.LayerNorm(hidden_size)# 可学习的残差连接权重self.alpha = torch.nn.Parameter(torch.tensor(0.5))

残差连接公式：
$\text{LayerNorm}(\alpha \cdot FFN(x) + (1-\alpha) \cdot \text{Attention}(x))$

完整模型架构

class SamOut(torch.nn.Module):def __init__(self, voc_size, hidden_size, num_heads, num_layers):super(SamOut, self).__init__()self.em = torch.nn.Embedding(voc_size, hidden_size, padding_idx=259)# 堆叠多个解码器层self.decoder_layers = torch.nn.ModuleList([DecoderLayer(hidden_size, num_heads) for _ in range(num_layers)])self.head = nn.Linear(hidden_size, voc_size, bias=False)

训练过程分析

# 定义超参数
voc_size = 8192 + 268
num_layers = 8
hidden_size = 2 ** 6 * num_heads  # 层数自适应设计
learning_rate = 0.001
batch_size = 32
num_epochs = 1000# 模型参数量计算
params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total parameters: {params:,}")

训练流程：

随机生成训练数据（实际应用中替换为真实数据集）
前向传播计算预测结果
使用交叉熵损失函数
反向传播优化参数

for epoch in range(num_epochs):# 数据准备data = torch.randint(low=0, high=voc_size, size=(batch_size, 50))input_tensor = data[:, :-1]target_tensor = data[:, 1:]# 前向传播output, _ = model(input_tensor)# 计算损失output = output.reshape(-1, voc_size)target = target_tensor.reshape(-1)loss = criterion(output, target)# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

完整代码

import timeimport pandas as pd
import torch
from torch import nn, optimclass MaxStateSuper(torch.nn.Module):def __init__(self, dim_size, heads):super(MaxStateSuper, self).__init__()self.heads = headsassert dim_size % heads == 0, "Dimension size must be divisible by head size."# 合并三个线性层为一个self.combined = nn.Linear(dim_size, 4 * dim_size, bias=False)# self.out_proj = nn.Linear(dim_size, dim_size)# self.layer_norm = torch.nn.LayerNorm(5)self.alpha1 = torch.nn.Parameter(torch.tensor(0.5))# self.alpha1 = torch.nn.Parameter(torch.tensor([[0.05] * 5]))# self.alpha2 = torch.nn.Parameter(torch.tensor([[0.05]] * 5))#self.alpha2 = torch.nn.Parameter(torch.tensor(0.5))#self.alpha3 = torch.nn.Parameter(torch.tensor(0.5))self.alpha4 = torch.nn.Parameter(torch.tensor(0.5))self.alpha5 = torch.nn.Parameter(torch.tensor(0.5))self.alpha6 = torch.nn.Parameter(torch.tensor(0.5))# self.alpha7 = torch.nn.Parameter(torch.tensor(0.5))# self.alpha8 = torch.nn.Parameter(torch.tensor(0.5))def forward(self, x, state=None):b, s, d = x.shape# 合并后的线性变换并分割combined = self.combined(x).chunk(4, dim=-1)out, out1, out2, out3 = combined# 调整张量形状，使用view优化out = out.view(b, s, self.heads, -1).permute(0, 2, 1, 3)out1 = out1.view(b, s, self.heads, -1).permute(0, 2, 1, 3)out2 = out2.view(b, s, self.heads, -1).permute(0, 2, 1, 3)out3 = out3.view(b, s, self.heads, -1).permute(0, 2, 1, 3)# out4 = out4.view(b, s, self.heads, -1).permute(0, 2, 1, 3)out4 = torch.cummax(out2, dim=2)[0]out = self.gen_model(out, out1, out2, out3, out4)# 恢复形状out = out.permute(0, 2, 1, 3).contiguous().view(b, s, d)return out, statedef gen_model(self, a, b, c, d, e):x = self.alpha1 * b + self.alpha2 * d + self.alpha3 * a + self.alpha4 * c + self.alpha5 * ex = a * e + xx = b * e + xx = c * e + xx = d * e + xx = e * self.alpha6 + xreturn xclass FeedForward(torch.nn.Module):def __init__(self, hidden_size):super(FeedForward, self).__init__()self.ffn1 = torch.nn.Linear(hidden_size, hidden_size)self.ffn2 = torch.nn.Linear(hidden_size, hidden_size)self.gate = torch.nn.Linear(hidden_size, hidden_size)self.relu = torch.nn.ReLU()# self.gr = torch.nn.Dropout(0.02)def forward(self, x):x1 = self.ffn1(x)x2 = self.relu(self.gate(x))xx = x1 * x2x = self.ffn2(xx)return xclass DecoderLayer(torch.nn.Module):def __init__(self, hidden_size, num_heads):super(DecoderLayer, self).__init__()self.self_attention = MaxStateSuper(hidden_size, num_heads)# self.self_attention = MaxState(hidden_size, num_heads)self.ffn = FeedForward(hidden_size)self.layer_norm = torch.nn.LayerNorm(hidden_size)self.alpha = torch.nn.Parameter(torch.tensor(0.5))def forward(self, x, state=None, ):x1, state = self.self_attention(x, state)x = self.layer_norm(self.alpha * self.ffn(x1) + (1 - self.alpha) * x)return x, stateclass SamOut(torch.nn.Module):def __init__(self, voc_size, hidden_size, num_heads, num_layers):super(SamOut, self).__init__()self.em = torch.nn.Embedding(voc_size, hidden_size, padding_idx=259)self.decoder_layers = torch.nn.ModuleList([DecoderLayer(hidden_size, num_heads) for _ in range(num_layers)])self.head = nn.Linear(hidden_size, voc_size, bias=False)# self.alpha = [torch.nn.Parameter(torch.tensor(0.5)) for i in range(num_layers)]# self.layer_norm = torch.nn.LayerNorm(hidden_size)def forward(self, x, state=None):x = self.em(x)if state is None:state = [None] * len(self.decoder_layers)i = 0for ii, decoder_layer in enumerate(self.decoder_layers):x1, state[i] = decoder_layer(x, state[i])x = x1 + xi += 1x = self.head(x)return x, stateif __name__ == '__main__':# 这里假设 DecoderLayer 已经定义好了，具体实现可以参考之前提供的代码或根据需要自定义# 定义超参数voc_size = 8192 + 268num_layers = 8hidden_size = 2 ** 6 * num_layersnum_heads = num_layerslearning_rate = 0.001batch_size = 32num_epochs = 1000# 初始化模型model = SamOut(voc_size=voc_size, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers)params = 0# [i.shape[0]  and len(i.shape) == 1 elif i.shape[1] * i.shape[0]for i in model.parameters():if i.shape != torch.Size([]):params += i.numel()print(params)# 定义损失函数和优化器criterion = nn.CrossEntropyLoss(ignore_index=3)  # 忽略填充标记的损失计算optimizer = optim.Adam(model.parameters(), lr=learning_rate)# 模拟一些训练数据（实际应用中应该使用真实的数据集）# 训练循环start_time = time.time()for epoch in range(num_epochs):data = torch.randint(low=0, high=voc_size, size=(batch_size, 50))  # 输入序列长度为50input_tensor = data[:, :-1]target_tensor = data[:, 1:]# 前向传播output, _ = model(input_tensor)# 将输出reshape以适应 CrossEntropyLoss 的输入要求output = output.reshape(-1, voc_size)target_tensor = target_tensor.reshape(-1)# 计算损失loss = criterion(output, target_tensor)# output_mean = (torch.nn.functional.softmax(output, -1)-1).mean()**2# c = loss.item() / 50# loss = loss - output_mean# loss = losoptimizer.zero_grad()  # 清除梯度# 反向传播和优化loss.backward()optimizer.step()print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}--')print("Training complete.{}".format(time.time() - start_time))
# Epoch [1/1], Loss: 4.0645,idx -142800:   1%|▏         | 239/16667 [02:11<2:50:52,  1.60it/s]
# Epoch [1/1], Loss: 4.0816,idx -145200:   1%|▏         | 243/16667 [02:21<2:55:54,  1.56it/s]