深入解析BERT输入表示：词嵌入、位置嵌入与句段嵌入的融合机制

1次阅读

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

BERT 模型的强大能力很大程度上来源于其精心设计的输入表示机制。今天我们就来拆解这个 ” 三合一 ” 的嵌入系统，看看词嵌入 (token embeddings)、位置嵌入(position embeddings) 和句段嵌入 (segment embeddings) 是如何协同工作的。

词嵌入：将每个 token 映射到固定维度的向量空间，是语义表示的基础层。例如 ”bank” 在金融和河岸场景会有不同向量。
位置嵌入：为每个位置编号分配独特向量，解决 Transformer 无法天然感知顺序的问题。位置 3 和位置 7 的相同单词会获得不同表示。
句段嵌入：标识 token 属于句子 A 还是 B（主要用于问答和 NSP 任务），让模型理解句子边界。通常用 0 / 1 区分不同句子。

最终输入表示是三个嵌入的简单相加：

E = E_{token} + E_{position} + E_{segment}

假设嵌入维度为 d，则三者都是∈ℝ^(d)的向量。这个设计保证了：

维度一致性：所有嵌入必须同维度
信息平等：相加不会人为赋予某类嵌入更高权重
计算高效：比拼接 (concatenation) 节省 2 / 3 的后续计算量

import torch
import torch.nn as nn

class BERTEmbedding(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)  # (V, d)
        self.pos_embed = nn.Embedding(max_len, d_model)       # (L, d)
        self.seg_embed = nn.Embedding(n_segments, d_model)    # (2, d)

    def forward(self, x, seg_ids):
        # x: (batch_size, seq_len)
        # seg_ids: (batch_size, seq_len)
        batch_size, seq_len = x.shape

        # 生成位置编号 [0,1,2,...,seq_len-1]
        pos = torch.arange(seq_len).expand(batch_size, seq_len)  # (batch, seq)

        # 求和并做 LayerNorm（实际 BERT 会额外加）return self.token_embed(x) + \
               self.pos_embed(pos) + \
               self.seg_embed(seg_ids)

关键点说明：