位置编码消融实验：Naive方案的得与失

"""位置编码消融实验对比脚本"""import numpy as npimport torchimport torch.nn as nnimport math  def naive_sinusoidal_positional_encoding(max_len, d_model):    """Naive sin/cos位置编码: 使用固定频率的周期函数"""    pe = torch.zeros(max_len, d_model)    positions = torch.arange(0, max_len).unsqueeze(1).float()    div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))    pe[:, 0::2] = torch.sin(positions * div_term)    pe[:, 1::2] = torch.cos(positions * div_term)    return pe.unsqueeze(0)  def relative_position_bias(max_len, num_heads, max_distance=64):    """相对位置偏置: 基于偏移量差值生成可学习偏置"""    # 生成所有位置对的距离偏移量 (i-j)    pos = torch.arange(max_len).unsqueeze(1)  # [max_len, 1]    offset = pos - pos.T  # [max_len, max_len] 相对偏移矩阵      # 截断超出预定义范围的偏移量    offset_clip = torch.clamp(offset, -max_distance, max_distance)      # 将偏移量统一映射到非负索引空间 (0 ~ 2*max_distance)    index = offset_clip + max_distance      # 为每个注意力头生成可学习的偏置表    bias_table = nn.Parameter(torch.randn(num_heads, 2 * max_distance + 1))    relative_bias = bias_table[:, index]  # [heads, max_len, max_len]    return relative_bias  class PositionEncodingAblationModel(nn.Module):    """用于消融实验的简化Transformer模型"""      def __init__(self, vocab_size, d_model=128, n_heads=4, max_len=64,                  encoding_type='naive'):        super().__init__()        self.encoding_type = encoding_type        self.token_embed = nn.Embedding(vocab_size, d_model)              if encoding_type == 'naive':            # Naive方案: 固定sin/cos编码, 不可训练            self.register_buffer('pos_enc', naive_sinusoidal_positional_encoding(max_len, d_model))        elif encoding_type == 'learned':            # 可学习方案: 位置嵌入作为可训练参数            self.pos_embed = nn.Parameter(torch.randn(1, max_len, d_model) * 0.1)        elif encoding_type == 'relative':            # 相对方案: 注意力层加入可学习的相对位置偏置            self.relative_bias = nn.Parameter(torch.randn(n_heads, max_len, max_len) * 0.1)              # 简化注意力: 仅使用一层自注意力（为演示代码精简）        self.attn_proj = nn.Linear(d_model, d_model)        self.output_proj = nn.Linear(d_model, vocab_size)        self.n_heads = n_heads        self.d_model = d_model          def forward(self, x):        batch, seq_len = x.shape        token_emb = self.token_embed(x)  # (batch, seq_len, d_model)              if self.encoding_type == 'naive':            # 引入位置信息: 词嵌入 + 位置编码            hidden = token_emb + self.pos_enc[:, :seq_len, :]        elif self.encoding_type == 'learned':            hidden = token_emb + self.pos_embed[:, :seq_len, :]        else:            hidden = token_emb  # 相对编码时位置信息已注入到注意力中              # 简化注意力（演示维度调整）        attn_out = self.attn_proj(hidden)              if self.encoding_type == 'relative':            # 将相对位置偏置作用到注意力打分上（此处用简化近似）            bias = self.relative_bias[:, :seq_len, :seq_len].mean(dim=0)  # (seq, seq)            attn_scores = torch.matmul(attn_out, attn_out.transpose(1, 2)) / math.sqrt(self.d_model)            attn_scores = attn_scores + bias.unsqueeze(0)            attn_scores = torch.softmax(attn_scores, dim=-1)            attn_out = torch.matmul(attn_scores, hidden)              logits = self.output_proj(attn_out)        return logits  def run_ablation_experiment():    """执行三种位置编码方案的消融对比实验"""    vocab_size = 200    batch_size = 4    seq_len = 32    epochs = 5      # 构造一份随机模拟训练数据——真实实验中需要替换为真实语料    torch.manual_seed(42)    data = torch.randint(0, vocab_size, (batch_size * 20, seq_len))      for enc_type in ['naive', 'learned', 'relative']:        model = PositionEncodingAblationModel(            vocab_size=vocab_size,             d_model=64,             n_heads=2,            max_len=64,            encoding_type=enc_type        )        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)        loss_fn = nn.CrossEntropyLoss()              total_loss = 0.0        for epoch in range(epochs):            for i in range(0, len(data), batch_size):                batch = data[i:i + batch_size]                logits = model(batch)                # 预测下一个词（语言建模任务）                loss = loss_fn(                    logits[:, :-1].reshape(-1, vocab_size),                    batch[:, 1:].reshape(-1)                )                optimizer.zero_grad()                loss.backward()                optimizer.step()                total_loss += loss.item()              print(f"{enc_type:>10s} | 平均损失: {total_loss / (epochs * len(data) / batch_size):.4f}")  if __name__ == '__main__':    run_ablation_experiment()