嵌入层优化：从查找到计算

嵌入层：被忽视的显存大户

在LLM中，嵌入层（Embedding Layer）往往是被忽视的显存消耗者。以Llama-3为例，词表大小128K、隐藏维度4096、FP16精度下，嵌入层参数量为 128K × 4096 × 2 = 1GB。对于更大的词表（如Qwen的152K），嵌入层可达1.5GB以上。

在70B模型中，嵌入层占总参数量的约2%，但在推理时它还需要一个等大的输出投影层（LM Head），两者合计占4%——这还不算小。

标准嵌入层

查找表

标准嵌入层是一个简单的查找表——词表中的每个token对应一个可学习的向量：

class StandardEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(vocab_size, d_model))
    
    def forward(self, input_ids):
        # 简单的索引查找
        return self.weight[input_ids]  # [batch, seq_len, d_model]

权重共享

许多LLM共享嵌入层和输出层的权重（Weight Tying）：

class SharedEmbeddingModel(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        # 输出层直接使用嵌入矩阵的转置
        # 不需要额外的参数
    
    def forward(self, input_ids):
        x = self.embedding(input_ids)
        x = self.transformer(x)
        # 输出logits = x @ embedding.weight.T
        logits = F.linear(x, self.embedding.weight)
        return logits

权重共享的好处：

减少参数量（节省1GB+显存）
输入和输出在同一个语义空间中
训练时梯度信号更强

但也有研究指出，权重共享可能导致冲突——输入嵌入需要捕捉"词的语义"，输出层需要捕捉"词的预测概率分布"，两者的最优表示可能不同。

嵌入层压缩

矩阵分解

将大的嵌入矩阵分解为两个小矩阵的乘积：

class FactorizedEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, d_factor):
        """
        d_factor: 分解后的中间维度
        """
        super().__init__()
        # V × D → V × d_factor × d_factor × D
        self.factor1 = nn.Parameter(torch.randn(vocab_size, d_factor))
        self.factor2 = nn.Parameter(torch.randn(d_factor, d_model))
    
    def forward(self, input_ids):
        # 查找低维因子，然后投影到高维
        low_dim = self.factor1[input_ids]  # [batch, seq, d_factor]
        return low_dim @ self.factor2      # [batch, seq, d_model]

参数量从 V×D 降低到 V×d + d×D。当 D=4096、d=256、V=128K 时，参数量减少约16倍。

ALBERT式分解

ALBERT使用了这种分解策略，将嵌入维度固定为一个较小值（如128），然后通过投影矩阵扩展到隐藏维度（如4096）：

class ALBERTEmbedding(nn.Module):
    def __init__(self, vocab_size, d_embedding=128, d_model=4096):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_embedding)
        self.projection = nn.Linear(d_embedding, d_model)
    
    def forward(self, input_ids):
        return self.projection(self.embedding(input_ids))

量化嵌入

将嵌入矩阵量化到低精度：

class QuantizedEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, n_bits=4):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_bits = n_bits
        
        # 量化的权重和缩放因子
        self.weight_quantized = nn.Parameter(
            torch.zeros(vocab_size, d_model, dtype=torch.int8 if n_bits == 8 else None)
        )
        self.scales = nn.Parameter(torch.ones(vocab_size, 1))
    
    def forward(self, input_ids):
        # 查找并反量化
        quantized = self.weight_quantized[input_ids]
        return quantized.float() * self.scales[input_ids]

动态嵌入

哈希嵌入

不存储完整的嵌入矩阵，而是通过哈希函数动态生成嵌入：

class HashEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, n_buckets=4096, n_hash=4):
        super().__init__()
        self.n_hash = n_hash
        self.n_buckets = n_buckets
        
        # 每个哈希函数对应一个小嵌入表
        self.embeddings = nn.ModuleList([
            nn.Embedding(n_buckets, d_model // n_hash)
            for _ in range(n_hash)
        ])
        
        # 哈希函数
        self.hash_seeds = [random.randint(0, 2**32) for _ in range(n_hash)]
    
    def forward(self, input_ids):
        outputs = []
        for i, seed in enumerate(self.hash_seeds):
            # 哈希到桶
            bucket_ids = (input_ids * seed) % self.n_buckets
            outputs.append(self.embeddings[i](bucket_ids))
        
        return torch.cat(outputs, dim=-1)  # [batch, seq, d_model]

哈希嵌入的参数量从 V×D 降低到 K×B×(D/K)，其中B是桶数（通常远小于V）。

乘性嵌入

通过乘法运算生成嵌入，完全不需要查找表：

class MultiplicativeEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, d_seed=32):
        super().__init__()
        # 使用随机投影矩阵
        self.projection = nn.Parameter(torch.randn(d_seed, d_model))
        # 将token ID编码为种子向量
        self.token_encoder = nn.Embedding(vocab_size, d_seed)
    
    def forward(self, input_ids):
        seed = self.token_encoder(input_ids)  # [batch, seq, d_seed]
        return seed @ self.projection  # [batch, seq, d_model]

嵌入层的梯度优化

稀疏梯度

嵌入层的梯度是稀疏的——每次只有被访问的token对应的行获得梯度。PyTorch的nn.Embedding支持稀疏梯度：

# 启用稀疏梯度
embedding = nn.Embedding(vocab_size, d_model, sparse=True)

# 使用SparseAdam优化器
optimizer = torch.optim.SparseAdam(embedding.parameters(), lr=0.001)

稀疏梯度的优势：

减少梯度计算量（只更新被访问的行）
减少梯度通信量（分布式训练中）
支持更大的词表

梯度裁剪

嵌入层的梯度可能非常大（特别是高频词），需要特别的梯度裁剪：

def embedding_gradient_clipping(embedding, max_norm=1.0):
    """嵌入层专用梯度裁剪"""
    if embedding.weight.grad is not None:
        # 只对有梯度的行裁剪
        grad = embedding.weight.grad
        if grad.is_sparse:
            grad_values = grad._values()
            grad_norm = grad_values.norm(dim=-1, keepdim=True)
            scale = (max_norm / (grad_norm + 1e-6)).clamp(max=1.0)
            grad_values.mul_(scale)
        else:
            torch.nn.utils.clip_grad_norm_(embedding, max_norm)

2026年嵌入层前沿

上下文感知嵌入

传统嵌入是静态的——同一个token在任何上下文中都有相同的嵌入。上下文感知嵌入根据上下文动态调整：

class ContextualEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, d_context=64):
        super().__init__()
        self.base_embedding = nn.Embedding(vocab_size, d_model)
        self.context_encoder = nn.Linear(d_model, d_context)
        self.context_to_offset = nn.Linear(d_context, d_model)
    
    def forward(self, input_ids):
        base = self.base_embedding(input_ids)
        
        # 计算上下文向量（窗口平均）
        context = base.mean(dim=1, keepdim=True)  # [batch, 1, d_model]
        context_offset = self.context_to_offset(
            self.context_encoder(context)
        )
        
        # 基础嵌入 + 上下文偏移
        return base + context_offset

多粒度嵌入

不同层使用不同粒度的嵌入——浅层使用字符级嵌入，深层使用词级嵌入：

class MultiGranularityEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.char_embedding = nn.Embedding(vocab_size, d_model)
        self.word_embedding = nn.Embedding(vocab_size, d_model)
        self.granularity_gate = nn.Linear(d_model, 1)
    
    def forward(self, input_ids, layer_idx, n_layers):
        char_emb = self.char_embedding(input_ids)
        word_emb = self.word_embedding(input_ids)
        
        # 随着层数加深，从字符嵌入过渡到词嵌入
        alpha = layer_idx / n_layers
        gate = torch.sigmoid(self.granularity_gate(char_emb))
        
        return (1 - alpha * gate) * char_emb + alpha * gate * word_emb

语义嵌入聚类

将语义相似的token的嵌入聚类在一起，减少有效参数量：

class ClusteredEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model, n_clusters=1024):
        super().__init__()
        self.n_clusters = n_clusters
        self.cluster_centers = nn.Parameter(torch.randn(n_clusters, d_model))
        self.token_to_cluster = nn.Parameter(torch.randint(0, n_clusters, (vocab_size,)))
        self.token_offsets = nn.Parameter(torch.zeros(vocab_size, d_model))
    
    def forward(self, input_ids):
        cluster_ids = self.token_to_cluster[input_ids]
        centers = self.cluster_centers[cluster_ids]
        offsets = self.token_offsets[input_ids]
        return centers + offsets

实践建议

权重共享：默认启用输入输出权重共享，节省大量参数
稀疏梯度：大词表（>50K）务必使用稀疏梯度优化
初始化：嵌入层使用正态分布初始化，均值0，标准差1/sqrt(d_model)
冻结策略：微调时可以冻结嵌入层（特别是小数据集场景）
量化：推理时嵌入层可以量化到INT4/INT8，对精度影响小

结语

嵌入层是LLM与文本世界的接口，其效率直接影响模型的训练和推理性能。从查找表到矩阵分解、从静态到动态嵌入、从密集到稀疏梯度，嵌入层优化是一个看似简单但充满深度的领域。随着词表规模的持续增长和多语言需求的提升，嵌入层优化将越来越重要。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。

嵌入层：被忽视的显存大户#

标准嵌入层#

查找表#

权重共享#

嵌入层压缩#

矩阵分解#

ALBERT式分解#

量化嵌入#

动态嵌入#

哈希嵌入#

乘性嵌入#

嵌入层的梯度优化#

稀疏梯度#

梯度裁剪#

2026年嵌入层前沿#

上下文感知嵌入#

多粒度嵌入#

语义嵌入聚类#

实践建议#

结语#

加入讨论#

📌 相关推荐