嵌入层

嵌入层基础

嵌入(Embedding)是将离散符号映射到连续向量空间的过程。在大语言模型中，嵌入层实现了两个关键功能：

1. 语义表示: 将每个token映射为高维向量，捕获其语义特征
2. 维度转换: 将一维的token ID序列转换为模型所需的高维特征表示

嵌入层的数学表示

嵌入层本质上是一个查找表(Look-up Table)操作：

其中：

• 是词表大小
• 是嵌入维度

对于输入token序列 ，嵌入操作为：

基本实现

嵌入层在PyTorch中实现：

class Embedding(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        # 使用正态分布初始化嵌入权重
        nn.init.normal_(self.embedding.weight, mean=0, std=0.02)
        
    def forward(self, x):
        # x: [batch_size, seq_len]
        # 输出: [batch_size, seq_len, embedding_dim]
        return self.embedding(x)

大语言模型中的嵌入层设计

1. 嵌入维度选择

嵌入维度()是模型的关键超参数：

2. 参数共享策略

为了减少参数量，许多模型在嵌入层和输出层之间共享权重：

class TransformerLM(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        # ... 其他层 ...
        self.output_projection = nn.Linear(d_model, vocab_size, bias=False)
        
        # 共享权重
        self.output_projection.weight = self.token_embedding.weight

权重共享的好处：

• 减少模型参数量(~20-30%的参数减少)
• 在某些情况下提高性能
• 提供正则化效果

3. 嵌入层缩放

嵌入向量通常需要适当缩放以稳定训练：

def scaled_embedding(x, embedding_layer, scale_factor=None):
    # 默认使用嵌入维度的平方根缩放
    if scale_factor is None:
        scale_factor = embedding_layer.embedding_dim ** 0.5
        
    return embedding_layer(x) * scale_factor

不同模型使用的缩放策略：

• Transformer原论文:
• GPT系列: 自适应归一化
• T5: 无额外缩放

嵌入类型与变体

1. 词嵌入 vs 子词嵌入

现代大语言模型主要使用子词(subword)嵌入：

2. 上下文无关嵌入 vs 上下文敏感嵌入

传统嵌入方法与现代语言模型的区别：

# 传统静态嵌入 (Word2Vec, GloVe)
static_embedding = embedding_layer(token_ids)

# 上下文敏感嵌入 (BERT, GPT)
input_embedding = embedding_layer(token_ids)
contextual_embedding = transformer_layers(input_embedding)  # 包含上下文信息

3. 分解嵌入 (Factorized Embedding)

ALBERT等模型使用分解嵌入降低参数量：

class FactorizedEmbedding(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super().__init__()
        # 嵌入维度通常小于隐藏维度
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        # 投影层将嵌入映射到更大的隐藏空间
        self.projection = nn.Linear(embedding_dim, hidden_dim, bias=False)
        
    def forward(self, x):
        # x: [batch_size, seq_len]
        embedded = self.embedding(x)  # [batch_size, seq_len, embedding_dim]
        projected = self.projection(embedded)  # [batch_size, seq_len, hidden_dim]
        return projected

分解嵌入优势：

• 大幅减少参数量(尤其是大词表模型)
• 分离语义表示和转换空间
• 提供额外灵活性

专门化嵌入设计

1. 多模态嵌入

多模态模型(如CLIP、DALL-E)需要处理不同模态的嵌入：

class MultimodalEmbedding(nn.Module):
    def __init__(self, text_vocab_size, image_vocab_size, embedding_dim):
        super().__init__()
        # 文本嵌入
        self.text_embedding = nn.Embedding(text_vocab_size, embedding_dim)
        # 图像嵌入(例如将图像patch转换为嵌入)
        self.image_embedding = nn.Embedding(image_vocab_size, embedding_dim)
        # 模态类型嵌入
        self.modality_embedding = nn.Embedding(2, embedding_dim)  # 0=文本, 1=图像
        
    def forward(self, text_ids=None, image_ids=None):
        batch_size = text_ids.shape[0] if text_ids is not None else image_ids.shape[0]
        
        embeddings = []
        modality_ids = []
        
        if text_ids is not None:
            text_emb = self.text_embedding(text_ids)
            embeddings.append(text_emb)
            modality_ids.append(torch.zeros(text_ids.shape[0], text_ids.shape[1], 
                                          device=text_ids.device, dtype=torch.long))
            
        if image_ids is not None:
            img_emb = self.image_embedding(image_ids)
            embeddings.append(img_emb)
            modality_ids.append(torch.ones(image_ids.shape[0], image_ids.shape[1], 
                                         device=image_ids.device, dtype=torch.long))
            
        # 合并不同模态的嵌入
        combined_emb = torch.cat(embeddings, dim=1)
        combined_modality_ids = torch.cat(modality_ids, dim=1)
        
        # 添加模态类型嵌入
        modality_emb = self.modality_embedding(combined_modality_ids)
        
        return combined_emb + modality_emb

2. 旋转式嵌入 (RoFormer)

将旋转位置编码直接融入嵌入过程：

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        
        # 生成旋转角度的频率
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        self.register_buffer("inv_freq", inv_freq)
        
    def forward(self, positions):
        # positions: [batch_size, seq_len]
        freqs = torch.einsum("i,j->ij", positions.float(), self.inv_freq)
        # 计算旋转矩阵的元素
        emb = torch.cat((freqs, freqs), dim=-1)
        cos_pos = emb.cos()
        sin_pos = emb.sin()
        return cos_pos, sin_pos

嵌入层训练与优化

1. 初始化策略

不同初始化方法对嵌入质量的影响：

# 常见嵌入初始化方法
def initialize_embeddings(embedding_layer, method='normal'):
    if method == 'normal':
        # 正态分布初始化 (大多数LLM使用)
        nn.init.normal_(embedding_layer.weight, mean=0, std=0.02)
    elif method == 'uniform':
        # 均匀分布初始化
        nn.init.uniform_(embedding_layer.weight, a=-0.1, b=0.1)
    elif method == 'xavier':
        # Xavier/Glorot初始化
        nn.init.xavier_uniform_(embedding_layer.weight)
    elif method == 'pretrained':
        # 加载预训练嵌入
        pretrained = load_pretrained_embeddings()
        embedding_layer.weight.data.copy_(pretrained)

2. 冻结与微调策略

在预训练和微调阶段，嵌入层的处理策略：

# 冻结嵌入层
for param in model.embedding.parameters():
    param.requires_grad = False
    
# 嵌入层使用较小学习率
optimizer_grouped_parameters = [
    {
        "params": [p for n, p in model.named_parameters() if "embedding" in n],
        "lr": base_learning_rate * 0.1,
    },
    {
        "params": [p for n, p in model.named_parameters() if "embedding" not in n],
        "lr": base_learning_rate,
    },
]

3. 嵌入压缩技术

对于部署场景，嵌入层通常是模型大小的主要贡献者，可以采用以下压缩方法：

• 量化：降低精度(例如从FP32到INT8)

  # 8位嵌入量化
  quantized_embeddings = torch.quantize_per_tensor(
      embedding_layer.weight.data, scale=0.1, zero_point=0, dtype=torch.qint8
  )
  

• 剪枝：移除低频或不重要token的嵌入

  # 基于频率的嵌入剪枝
  token_frequencies = compute_token_frequencies(dataset)
  threshold = find_frequency_threshold(token_frequencies, coverage=0.995)
  
  # 只保留高频token的嵌入
  important_token_ids = torch.where(token_frequencies > threshold)[0]
  pruned_embeddings = embedding_layer.weight.data[important_token_ids]
  

• 低秩分解：使用矩阵分解降低参数量

  # 低秩分解嵌入
  U, S, V = torch.svd(embedding_layer.weight.data)
  k = int(min(U.shape) * 0.7)  # 保留70%的奇异值
  compressed_embeddings = (U[:, :k] @ torch.diag(S[:k]) @ V[:, :k].T)
  

嵌入层与词表设计

嵌入层的性能与词表设计密切相关：

1. 词表大小与质量权衡

def analyze_vocabulary_coverage(tokenizer, corpus, vocab_sizes=[10000, 20000, 30000, 50000]):
    """分析不同词表大小的覆盖率和压缩率"""
    results = {}
    
    for size in vocab_sizes:
        # 使用不同大小训练词表
        new_tokenizer = train_tokenizer(corpus, vocab_size=size)
        
        # 计算覆盖指标
        oov_rate = compute_oov_rate(new_tokenizer, test_corpus)
        compression_ratio = compute_compression_ratio(new_tokenizer, test_corpus)
        
        results[size] = {
            "oov_rate": oov_rate,
            "compression_ratio": compression_ratio,
            "params_in_embedding": size * embedding_dim
        }
    
    return results

2. 特殊token设计

大语言模型中的特殊token嵌入设计：

class SpecialTokenEmbedding(nn.Module):
    def __init__(self, vocab_size, embedding_dim, n_special_tokens):
        super().__init__()
        self.regular_embedding = nn.Embedding(vocab_size, embedding_dim)
        self.special_embedding = nn.Embedding(n_special_tokens, embedding_dim)
        
        # 正常初始化普通token嵌入
        nn.init.normal_(self.regular_embedding.weight, mean=0, std=0.02)
        
        # 特殊初始化特殊token嵌入(可能希望它们更显著)
        nn.init.normal_(self.special_embedding.weight, mean=0, std=0.04)
        
    def forward(self, x, is_special=None):
        # x: [batch_size, seq_len]
        # is_special: [batch_size, seq_len] 二值掩码指示特殊token
        
        if is_special is None:
            # 默认全部为普通token
            return self.regular_embedding(x)
            
        # 应用不同的嵌入层
        regular_emb = self.regular_embedding(x)
        special_emb = self.special_embedding(x)
        
        # 根据掩码选择相应的嵌入
        return torch.where(is_special.unsqueeze(-1), special_emb, regular_emb)

特殊token示例：

• [CLS], [SEP], [MASK] (BERT)
• <s>, </s>, <pad> (GPT系列)
• <|im_start|>, <|im_end|> (对话模型)

嵌入层与多语言支持

1. 多语言共享嵌入设计

class MultilingualEmbedding(nn.Module):
    def __init__(self, vocab_sizes, embedding_dim, shared_vocab_size=0):
        super().__init__()
        
        self.lang_embeddings = nn.ModuleDict()
        self.embedding_dim = embedding_dim
        
        if shared_vocab_size > 0:
            # 创建共享嵌入层(用于常见token)
            self.shared_embedding = nn.Embedding(shared_vocab_size, embedding_dim)
        else:
            self.shared_embedding = None
            
        # 为每种语言创建专用嵌入(用于特定语言的token)
        for lang, size in vocab_sizes.items():
            if self.shared_embedding is not None:
                # 减去共享词表大小
                size = size - shared_vocab_size
            self.lang_embeddings[lang] = nn.Embedding(size, embedding_dim)
            
    def forward(self, x, lang):
        """
        x: 输入token ids
        lang: 当前处理的语言
        """
        if self.shared_embedding is not None:
            # 区分共享和语言特定的tokens
            shared_mask = x < shared_vocab_size
            lang_specific_mask = ~shared_mask
            
            # 初始化输出嵌入
            embedding = torch.zeros(
                *x.shape, self.embedding_dim, 
                device=x.device, dtype=self.shared_embedding.weight.dtype
            )
            
            # 应用共享嵌入
            shared_indices = x * shared_mask
            embedding[shared_mask] = self.shared_embedding(shared_indices[shared_mask])
            
            # 应用语言特定嵌入
            lang_indices = (x - shared_vocab_size) * lang_specific_mask
            embedding[lang_specific_mask] = self.lang_embeddings[lang](lang_indices[lang_specific_mask])
            
            return embedding
        else:
            # 无共享词表，直接使用语言特定嵌入
            return self.lang_embeddings[lang](x)

小结

嵌入层是大语言模型的关键组件，为模型提供了将离散符号转换为连续表示的能力：

1. 基本原理：嵌入层是一个查找表操作，将token ID映射到高维向量
2. 设计考量：嵌入维度、参数共享、缩放因子等超参数影响模型性能
3. 嵌入变体：从基本词嵌入发展到分解嵌入、多模态嵌入等高级变种
4. 优化技术：初始化策略、冻结策略、压缩方法等提高嵌入层效率
5. 词表设计：词表大小和特殊token处理影响嵌入层的表现

随着大语言模型规模不断扩大，嵌入层的设计也在不断演进，以平衡表达能力、计算效率和部署需求。