模型架构概述

MiniMind是一个轻量级开源语言模型项目，采用类似GPT的Transformer解码器架构。核心设计目标是保持模型性能的同时降低计算资源需求，适合研究者和开发者低成本部署实验。

核心组件

Transformer层堆叠 模型由N层相同的Transformer解码器块堆叠而成，每层包含多头自注意力机制和前馈神经网络。隐藏层维度通常设置为512或768，头数在8-16之间可配置。

注意力机制采用缩放点积形式： $$Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V$$

位置编码 使用可学习的位置编码替代原始Transformer的正弦函数，增强模型对位置信息的建模能力。位置嵌入维度与词嵌入维度相同，直接相加输入。

前馈网络 每层注意力后接两个全连接层，中间用GeLU激活： $$FFN(x)=W_2(GeLU(W_1x+b_1))+b_2$$

关键实现细节

权重共享 输出层的词嵌入矩阵与输入层的嵌入矩阵共享参数，显著减少模型参数量。这种技术被证明能提升小模型的语言建模性能。

层归一化 采用Pre-LN结构，在注意力层和前馈层前应用层归一化： $$x_{out}=x+Dropout(SubLayer(LayerNorm(x)))$$

稀疏注意力 可选配置局部窗口注意力或轴向注意力模式，减少长序列处理时的计算复杂度。窗口大小通常设置为128或256。

代码结构示例

模型主体实现通常包含以下PyTorch模块：

class TransformerLayer(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.1):
        super().__init__()
        self.self_attn = MultiheadAttention(d_model, nhead)
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)

class MiniMind(nn.Module):
    def __init__(self, vocab_size, n_layers=6, d_model=512):
        super().__init__()
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList(
            [TransformerLayer(d_model) for _ in range(n_layers)])

性能优化技术

梯度检查点 在训练阶段启用梯度检查点功能，以时间换空间，允许用更少的显存训练更深模型。典型配置是每2-4层设置一个检查点。

混合精度训练 默认使用FP16混合精度，通过自动梯度缩放维持训练稳定性。在支持Tensor Core的GPU上可获得2-3倍加速。

动态批处理 数据加载器实现动态批处理功能，根据序列长度自动调整batch size，最大化GPU利用率。典型设置是保持token数在4096左右。