4. 可学习位置编码

任务说明

本案例实现可学习位置编码(Learnable Positional Encoding),将其添加到输入序列的嵌入上,并在训练过程中通过反向传播更新位置嵌入参数。同时实现一个固定正弦位置编码作为对照,在相同的小型文本分类任务上比较两者的收敛表现。

实现方案

构建一个仅含单层Transformer编码器的小型分类模型。输入序列先经过嵌入层,然后分别与可学习位置编码或固定正弦编码相加,再送入编码器。训练数据使用随机生成的整数序列(模拟词索引),标签为序列首尾元素之和的奇偶性。通过对比训练损失下降曲线来评估两类位置编码的效果。

可学习 vs 固定位置编码 — 训练对比流程 输入序列 x 嵌入层 可学习位置编码 固定正弦编码 相加 Transformer编码器 分类头 交叉熵损失 可学习路径 固定路径 两分支共享嵌入与编码器,仅位置编码不同

核心代码

import torch
import torch.nn as nn
import math

class LearnablePositionalEncoding(nn.Module):
    """可学习位置编码:参数随训练更新"""
    def __init__(self, d_model, max_len=128):
        super().__init__()
        self.pe = nn.Parameter(torch.randn(max_len, d_model) * 0.1)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        return x + self.pe[:x.size(1), :]

class FixedPositionalEncoding(nn.Module):
    """固定正弦位置编码:不参与训练"""
    def __init__(self, d_model, max_len=128):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)  # buffer不参与梯度更新

    def forward(self, x):
        return x + self.pe[:x.size(1), :]

class TinyTransformer(nn.Module):
    """单层编码器 + 分类头,用于位置编码对比"""
    def __init__(self, vocab_size, d_model=32, nhead=4, num_layers=1, max_len=32, pos_encoding='learnable'):
        super().__init__()
        self._emb = nn.Embedding(vocab_size, d_model)
        if pos_encoding == 'learnable':
            self.pos_enc = LearnablePositionalEncoding(d_model, max_len)
        else:
            self.pos_enc = FixedPositionalEncoding(d_model, max_len)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True, dim_feedforward=64)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.classifier = nn.Linear(d_model, 2)

    def forward(self, x):
        emb = self._emb(x)          # (batch, seq, d_model)
        emb = self.pos_enc(emb)          # 添加位置编码
        out = self.encoder(emb)          # (batch, seq, d_model)
        pooled = out.mean(dim=1)         # 全局平均池化
        return self.classifier(pooled)

def generate_data(num_samples, seq_len=16, vocab_size=100):
    """生成随机序列,标签为首尾元素之和的奇偶性"""
    x = torch.randint(0, vocab_size, (num_samples, seq_len))
    y = ((x[:, 0] + x[:, -1]) % 2).long()
    return x, y

def train_and_evaluate(pos_type, epochs=30, batch_size=64):
    """训练并返回损失历史"""
    torch.manual_seed(42)
    vocab_size, seq_len = 100, 16
    model = TinyTransformer(vocab_size, d_model=32, max_len=seq_len, pos_encoding=pos_type)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    x_train, y_train = generate_data(2000, seq_len, vocab_size)
    x_val, y_val = generate_data(400, seq_len, vocab_size)

    loss_history = []
    for epoch in range(epochs):
        model.train()
        perm = torch.randperm(x_train.size(0))
        total_loss = 0.0
        for i in range(0, x_train.size(0), batch_size):
            idx = perm[i:i+batch_size]
            xb, yb = x_train[idx], y_train[idx]
            optimizer.zero_grad()
            logits = model(xb)
            loss = criterion(logits, yb)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * xb.size(0)
        loss_history.append(total_loss / x_train.size(0))

        # 每5轮验证一次
        if (epoch + 1) % 5 == 0:
            model.eval()
            with torch.no_grad():
                logits = model(x_val)
                pred = logits.argmax(dim=1)
                acc = (pred == y_val).float().mean().item()
            print(f"[{pos_type}] Epoch {epoch+1:2d} | Train Loss: {loss_history[-1]:.4f} | Val Acc: {acc:.4f}")
    return loss_history

if __name__ == "__main__":
    print("=== 可学习位置编码训练 ===")
    learn_loss = train_and_evaluate('learnable')
    print("\n=== 固定位置编码训练 ===")
    fixed_loss = train_and_evaluate('fixed')

    # 输出最终对比
    print(f"\n可学习编码最终损失: {learn_loss[-1]:.4f}")
    print(f"固定编码最终损失:   {fixed_loss[-1]:.4f}")

运行与验证

依赖:Python 3.8+,PyTorch 1.12+(需包含torch.nn.TransformerEncoderLayer)。

运行方式:将上述代码保存为learnable_pos_encoding.py,在终端执行python learnable_pos_encoding.py。

验证方法:观察训练过程中每5轮打印的训练损失与验证准确率。判定标准:

  • 两类位置编码的验证准确率最终都应达到0.9以上(任务简单,模型容量足够);
  • 对比learn_loss与fixed_loss的最终值及下降曲线——可学习编码通常在训练后期损失更低或收敛更快,这体现了其参数随任务自适应调整的优势;
  • 若固定编码与可学习编码表现接近,说明该任务对位置编码形式不敏感,可进一步增大序列长度或改用更复杂任务来观察差异。

预期输出(基于代码逻辑):训练30轮后,两类编码的验证准确率均接近1.0,可学习编码的训练损失通常略低于固定编码(例如可学习约0.05~0.15,固定约0.10~0.25),具体数值因随机初始化而异。