BERT句子嵌入表示与对比学习实战：从入门到生产环境部署

1次阅读

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

在 NLP 任务中，传统的句子嵌入方法如 TF-IDF 或 Word2Vec 平均池化存在明显缺陷：

无法处理一词多义（” 苹果 ” 公司 vs 水果）
忽略词序（” 猫追狗 ” 和 ” 狗追猫 ” 得到相同表示）
长文本信息稀释（重要信号被常见词淹没）

BERT 通过 Transformer 架构捕获上下文信息，而对比学习通过拉近相似样本、推远不相似样本来优化表示空间。二者的结合带来：

语义敏感度：BERT 的注意力机制识别关键词语义
空间一致性：对比学习使相似句子的嵌入距离更近
少样本适应：通过数据增强生成有效训练对

# Python 3.8+, PyTorch 1.12+
import torch
from transformers import BertModel, BertTokenizer
from torch import nn
import numpy as np

class BertContrastive(nn.Module):
    def __init__(self, model_name='bert-base-uncased'):
        super().__init__()
        self.bert = BertModel.from_pretrained(model_name)
        # 用[CLS] token 作为句子表示
        self.projection = nn.Linear(768, 256)  # 降维减少计算量

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids, attention_mask=attention_mask)
        cls_embedding = outputs.last_hidden_state[:, 0, :]  # 取 [CLS] 位置
        return self.projection(cls_embedding)

def contrastive_loss(embeddings, temperature=0.1):
    # embeddings: [batch_size, dim]
    sim_matrix = torch.matmul(embeddings, embeddings.T) / temperature
    exp_sim = torch.exp(sim_matrix)

    # 对角线是正样本对
    pos_pairs = torch.diag(exp_sim)
    # 每行的和减去自身是负样本对
    neg_pairs = exp_sim.sum(dim=1) - pos_pairs

    loss = -torch.log(pos_pairs / neg_pairs).mean()
    return loss

# 推荐配置
params = {
    'batch_size': 64,     # 太小影响对比学习效果
    'learning_rate': 3e-5,
    'temperature': 0.07,  # 控制相似度分布陡峭程度
    'max_length': 64      # 截断长文本
}

# 启用自动混合精度
from torch.cuda.amp import autocast

@torch.no_grad()
def batch_inference(texts):
    inputs = tokenizer(texts, padding=True, truncation=True, 
                      max_length=64, return_tensors="pt")
    with autocast():
        embeddings = model(**inputs)
    return embeddings.cpu().numpy()