浅谈 Transformer 模型的输入嵌入环节

如下图所示，笔者最近学习 Transformer 模型的架构，学习到了输入嵌入这一章节。

本文是笔者的学习笔记。

输入嵌入的定义

输入嵌入（Input Embedding）是将离散的符号（如单词、字符）转换为连续的向量表示的一种方法。神经网络在处理输入数据时，无法直接处理离散符号，因此需要将这些符号映射到一个高维的实数向量空间中。这样，输入嵌入使得网络能够捕捉到符号之间的语义关系，并使模型的训练更加高效。

例如，在自然语言处理中，单词是离散符号，传统的机器学习模型无法理解单词本身的语义信息。通过将单词转换为向量（如 Word2Vec 或 GloVe 的方式），每个单词可以被表示为一个高维向量，从而能够体现语义相似性，例如 king 和 queen 具有相近的向量。

Transformer 模型中的输入嵌入

Transformer 模型使用输入嵌入来表示序列中的每个元素。由于 Transformer 模型的核心组件是自注意力机制（Self-Attention），需要将序列中的每个符号表示为向量，这些向量不仅要捕捉到符号本身的信息，还需要携带位置信息，以便模型在全局上下文中理解每个符号的意义。

常见的输入嵌入方法

Transformer 模型中最常用的输入嵌入方法主要包括以下几种：

1. 词嵌入（Word Embedding）
2. 位置嵌入（Positional Encoding）
3. 特定任务的额外嵌入（例如，分段嵌入 Segment Embedding）

每种方法都有其独特的作用和特点，下面将逐一详细讨论。

1. 词嵌入（Word Embedding）

词嵌入是最基础的输入嵌入，用于将每个单词或词片段（token）转换为向量表示。在 Transformer 模型中，词嵌入层通常是一个可训练的矩阵。对于输入序列中的每个单词，通过查表操作获得对应的词向量。

例如，考虑句子 I love AI。假设词汇表大小为 10,000，每个单词用一个 512 维的向量表示。在这种情况下，嵌入层的权重矩阵是一个大小为 (10000, 512) 的矩阵。输入句子 I love AI 会被首先标记化为 [I, love, AI]，然后通过查找嵌入矩阵得到 [E_I, E_love, E_AI]，其中每个 E_x 是一个 512 维的向量。

一个常用的实现如下：

import torch
import torch.nn as nn

# 定义词嵌入的参数
vocab_size = 10000  # 词汇表大小
embedding_dim = 512  # 嵌入维度

# 创建嵌入层
embedding = nn.Embedding(vocab_size, embedding_dim)

# 示例输入，假设输入为词汇表索引
input_tokens = torch.tensor([1, 23, 456])  # "I love AI" 在词汇表中的索引

# 获取嵌入
embedded_tokens = embedding(input_tokens)
print(embedded_tokens.shape)  # 输出形状为 (3, 512)

这里，input_tokens 是句子中每个词对应的词汇表索引，通过嵌入层后得到了它们的连续向量表示。

2. 位置嵌入（Positional Encoding）

词嵌入本身无法体现单词在句子中的位置信息。而 Transformer 模型由于完全依赖于注意力机制，没有像 RNN 那样的顺序信息，因此需要加入位置嵌入来让模型了解符号的相对位置。位置嵌入通过为每个输入标记加入一个位置向量，使得模型能够感知每个单词的位置信息。

位置嵌入可以通过多种方式实现，其中一种最常用的是基于正弦和余弦函数的编码方法。以下公式给出了位置嵌入的计算方法：

• 对于位置 pos 和嵌入维度 i，定义位置嵌入向量为：

  PE(pos, 2i) = sin(pos / 10000^(2i / d))
  PE(pos, 2i+1) = cos(pos / 10000^(2i / d))
  

其中，pos 是词在序列中的位置，i 是嵌入向量的维度索引，d 是嵌入向量的总维度。这种编码方式的设计使得位置编码之间存在特定的关系，便于模型通过自注意力机制捕获位置依赖。

位置嵌入的代码实现如下：

import numpy as np
import torch

def positional_encoding(max_len, d_model):
    pos_encoding = np.zeros((max_len, d_model))
    
    for pos in range(max_len):
        for i in range(0, d_model, 2):
            pos_encoding[pos, i] = np.sin(pos / (10000 ** ((2 * i) / d_model)))
            pos_encoding[pos, i + 1] = np.cos(pos / (10000 ** ((2 * i) / d_model)))
    
    return torch.tensor(pos_encoding, dtype=torch.float)

# 定义位置嵌入的参数
max_len = 50  # 最大序列长度
d_model = 512  # 嵌入维度

# 获取位置嵌入
pos_encoding = positional_encoding(max_len, d_model)
print(pos_encoding.shape)  # 输出形状为 (50, 512)

在 Transformer 中，词嵌入和位置嵌入会相加后作为模型的输入。位置嵌入帮助模型理解词与词之间的顺序关系，这对自然语言处理任务中的语义理解至关重要。

3. 特定任务的额外嵌入

Transformer 模型在一些特定任务（如 BERT）中，还会加入额外的嵌入信息。例如，BERT 通过引入分段嵌入（Segment Embedding）来区分句子 A 和句子 B，以便在输入中明确标识不同的上下文段。分段嵌入是一个可训练的向量，用来区分输入的不同部分。

假设我们有两个句子：

• The cat is sleeping.
• The dog is barking.

为了让模型能够理解这是两个句子，通过使用分段嵌入，句子 A 的每个单词对应的分段标记可以是 0，句子 B 的每个单词对应的分段标记可以是 1，这样模型可以知道输入中的哪些部分是属于同一个上下文段的。

综合应用与实例分析

在 Transformer 模型中，这些嵌入方法结合起来，用于为模型提供尽可能多的信息。这些嵌入方法共同作用，使 Transformer 模型不仅能了解每个词的具体语义，还能意识到它们之间的顺序和结构。

下面我们通过一个代码实例，将词嵌入、位置嵌入和分段嵌入结合起来，展示 Transformer 中输入嵌入的整体实现过程：

import torch
import torch.nn as nn

class TransformerInputEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, max_len, num_segments):
        super(TransformerInputEmbedding, self).__init__()
        self.word_embedding = nn.Embedding(vocab_size, d_model)
        self.segment_embedding = nn.Embedding(num_segments, d_model)
        
        pos_encoding = positional_encoding(max_len, d_model)
        self.register_buffer('pos_encoding', pos_encoding)
    
    def forward(self, input_ids, segment_ids):
        word_embeds = self.word_embedding(input_ids)  # (batch_size, seq_len, d_model)
        segment_embeds = self.segment_embedding(segment_ids)  # (batch_size, seq_len, d_model)
        
        # 获取位置嵌入
        pos_embeds = self.pos_encoding[:input_ids.size(1), :].unsqueeze(0)  # (1, seq_len, d_model)
        
        # 相加得到最终嵌入
        embeddings = word_embeds + segment_embeds + pos_embeds
        return embeddings

# 参数定义
vocab_size = 10000
d_model = 512
max_len = 50
num_segments = 2

# 创建输入嵌入层
embedding_layer = TransformerInputEmbedding(vocab_size, d_model, max_len, num_segments)

# 输入示例
input_ids = torch.tensor([[1, 23, 456, 0, 0]])  # 示例句子
segment_ids = torch.tensor([[0, 0, 0, 1, 1]])  # 分段信息

# 获取最终嵌入
output_embeddings = embedding_layer(input_ids, segment_ids)
print(output_embeddings.shape)  # 输出形状为 (1, 5, 512)

这个代码实例展示了如何将词嵌入、位置嵌入和分段嵌入相结合，最终得到用于 Transformer 模型的输入嵌入。input_ids 表示输入的句子，其中每个值为词汇表索引，segment_ids 则用来区分句子中的不同部分。

输入嵌入的实际应用案例

为了帮助理解这些嵌入的实际作用，以下是一些典型的场景：

机器翻译

在机器翻译任务中，输入嵌入的应用非常关键。例如，假设我们需要将英文句子翻译为法文，输入模型的英文句子需要转换为连续的词嵌入，位置嵌入提供每个词的位置信息，分段嵌入则可能区分输入句子和生成句子。

通过这种方式，模型不仅知道输入的词汇是什么，还知道它们在句子中的顺序，从而能够更好地理解上下文并生成合理的翻译结果。

自然语言理解任务

在自然语言理解任务（如 BERT）中，输入嵌入层在模型的预训练和下游任务中起着至关重要的作用。例如，在问答系统中，输入可以包含一个问题和一个上下文段落，分段嵌入用来区分问题和段落，使得模型在理解问题与段落之间的关系时更具针对性。

输入嵌入的优化与改进

虽然标准的词嵌入和位置嵌入已经非常有效，但近年来仍有很多研究致力于改进输入嵌入的方法。

• 动态位置嵌入（Dynamic Positional Encoding）：传统的位置编码是固定的，但在某些应用中，动态位置编码可以使模型根据不同任务的需求调整位置信息。

• 词片段嵌入（Subword Embedding）：在 GPT-3 等模型中，使用词片段而不是完整单词进行嵌入，这样可以有效减少词汇表大小并处理未知词汇。词片段嵌入通过 BPE（Byte Pair Encoding） 等技术实现，将复杂词汇分解为更小的可处理单元。

• 知识嵌入（Knowledge Embedding）：有些模型通过引入外部知识库的嵌入来增强模型的知识能力，使模型在面对开放领域问题时表现更好。

这些改进的方向大大丰富了输入嵌入的应用场景，并且提升了 Transformer 模型在不同任务中的表现能力。

省流版

输入嵌入是 Transformer 模型中非常重要的一部分，它将离散符号映射到连续向量空间，使模型能够理解输入数据的语义关系。常见的输入嵌入方法包括词嵌入、位置嵌入以及特定任务的额外嵌入，这些嵌入共同作用，使得 Transformer 模型能够高效处理各种自然语言任务。