BERT输入表示解析：如何高效组合token、position和segment embeddings

共计 2073 个字符，预计需要花费 6 分钟才能阅读完成。

三种 Embeddings 的基础作用

在 BERT 模型中，输入表示由三种不同类型的嵌入组合而成，每种嵌入都承担着特定的功能：

• Token Embeddings：将词汇表中的每个单词映射到一个固定维度的向量空间。例如，对于单词 ”apple”，会对应一个 768 维的浮点数向量（以 BERT-base 为例）。数学上可表示为：$E_{token} \in \mathbb{R}^{n \times d}$，其中 n 为序列长度，d 为隐藏层维度。

• Position Embeddings：为序列中的每个位置赋予独特的向量表示，使模型感知词序信息。位置编码公式为：$PE_{(pos,2i)} = sin(pos/10000^{2i/d})$，$PE_{(pos,2i+1)} = cos(pos/10000^{2i/d})$。

• Segment Embeddings：用于区分不同句子（如问答任务中的问题和答案）。通常用 0 表示第一个句子，1 表示第二个句子，对应学习两个特定的嵌入向量。

简单相加的潜在问题

虽然直接将三种嵌入相加是最常见的做法，但这种操作可能引发以下问题：

1. 信息混淆：不同嵌入空间的向量直接相加可能导致语义信息相互干扰
2. 数值范围差异：未经归一化的各嵌入可能具有不同的数值分布范围
3. 维度不对齐：当各嵌入维度不一致时（如某些定制化模型）会导致运算错误

PyTorch 实现代码

import torch
import torch.nn as nn

class BERTEmbeddingCombiner(nn.Module):
    def __init__(self, vocab_size, max_len, d_model, n_segments=2):
        super().__init__()
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(max_len, d_model)
        self.seg_embed = nn.Embedding(n_segments, d_model)

        # 初始化位置 ID（实际使用中应根据真实序列长度生成）self.register_buffer('position_ids', torch.arange(max_len).expand((1, -1)))

    def forward(self, input_ids, segment_ids=None):
        """
        参数:
            input_ids: [batch_size, seq_len]
            segment_ids: [batch_size, seq_len]，若不提供则默认为 0
        返回:
            组合后的嵌入: [batch_size, seq_len, d_model]
        """
        seq_len = input_ids.size(1)

        # 生成 token 嵌入
        token_embeddings = self.token_embed(input_ids)

        # 生成位置嵌入（截取实际序列长度）position_embeddings = self.pos_embed(self.position_ids[:, :seq_len])

        # 生成 segment 嵌入（默认为 0）if segment_ids is None:
            segment_ids = torch.zeros_like(input_ids)
        segment_embeddings = self.seg_embed(segment_ids)

        # 组合三种嵌入（关键步骤）combined = token_embeddings + position_embeddings + segment_embeddings

        # 层归一化（可选）return combined

不同组合策略对比

通过实验对比多种组合方式在 GLUE 基准测试上的表现：

关键发现：
– 简单相加在速度和内存上最优，精度损失可接受
– 复杂组合方式在小数据集上可能过拟合
– 门控机制在长文本任务中表现突出

生产环境最佳实践

1. 维度对齐检查 ：在模型初始化时验证三种嵌入的维度一致性

assert token_embed.weight.shape == pos_embed.weight.shape == seg_embed.weight.shape

1. 归一化处理 ：对相加后的嵌入应用 LayerNorm

2. 分段掩码 ：处理单句输入时显式设置 segment_ids 为全 0

3. 位置扩展 ：当序列超长时，采用循环位置编码或外推方法

4. 混合精度训练 ：使用 AMP 自动管理嵌入层的数值范围

开放性问题

1. 是否存在比简单相加更有效的非线性组合方式？
2. 如何设计自适应权重的 embedding 组合机制？
3. 对于多模态场景，如何扩展当前的 embedding 组合范式？
4. 能否通过知识蒸馏将复杂组合方式压缩到简单相加的模型中？

在实际应用中，embedding 组合策略的选择需要根据具体任务需求、硬件条件和实时性要求进行权衡。建议先采用标准实现进行基线测试，再针对性地尝试优化方案。