AI嵌入技术：Word2Vec不够用了之后

Word2Vec Sentence-BERT 嵌入模型 NLP
作者: 醉月思 发布日期: 2021-05-13 分类: 技术实践

去年做一个新闻分类项目时，遇到的第一个问题传统的 TF-IDF 没法语义相似，词袋模型又丢失了上下文，后来研究了一圈嵌入技术，从 Word2Vec 开始，一路摸到 Sentence-BERT，。

为什么要用嵌入

最开始做文本分类时，我用的是 sklearn 的 TF-IDF 向量化器：

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB

vectorizer = TfidfVectorizer(max_features=5000)
X_train = vectorizer.fit_transform(train_texts)
X_test = vectorizer.transform(test_texts)

clf = MultinomialNB()
clf.fit(X_train, y_train)

这个方案在简单场景下能跑通，但很快发现问题：

  1. 同义词不识别："优秀" 和 "出色" 在向量空间里完全不同
  2. 上下文丢失："银行" 和 "河岸" 用同一个词向量表示
  3. 维度灾难：中文词表动辄几十万维度，模型训练慢

这时候才知道，问题的根源在于文本表示，不是模型不够强。嵌入技术的核心想法很简单：把词或句子映射到一个连续的低维向量空间里，在这个空间里，语义相近的词或句子距离更近。

Word2Vec 实践

Word2Vec 是 Google 团队在 2013 年提出的，核心是两个模型：CBOW（连续词袋模型）和 Skip-gram。CBOW 用上下文预测中心词，Skip-gram 用中心词预测上下文。我选择 Skip-gram，因为在稀疏词和长尾词的表现上更好。

环境准备

python -m venv venv
source venv/bin/activate
pip install gensim==4.3.2 numpy==1.24.3 jieba==0.42.1

注意 gensim 版本，4.0 以后 API 有变化，网上很多老教程会报错。我第一次用 3.8.3 版本，迁移到新环境时发现 model.wv 的用法全变了，折腾了半小时才搞明白。

数据准备与训练

我用的是中文新闻语料，需要先用 jieba 分词：

import jieba
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
import multiprocessing

# 分词预处理
def segment_text(input_file, output_file):
    with open(input_file, 'r', encoding='utf-8') as f_in, \
         open(output_file, 'w', encoding='utf-8') as f_out:
        for line in f_in:
            words = jieba.lcut(line.strip())
            f_out.write(' '.join(words) + '\n')

# 训练模型
def train_word2vec(corpus_file, model_path):
    sentences = LineSentence(corpus_file)

    model = Word2Vec(
        sentences,
        vector_size=300,          # 向量维度
        window=5,                 # 上下文窗口
        min_count=3,              # 最小词频
        workers=multiprocessing.cpu_count(),
        sg=1,                     # 1为Skip-gram，0为CBOW
        epochs=10,
        negative=10               # 负采样数量
    )

    model.save(model_path)
    return model

# 使用示例
model = Word2Vec.load("word2vec.model")
print(model.wv['深度学习'])  # 获取词向量
print(model.wv.most_similar('人工智能', topn=5))  # 找相似词

踩过的坑

坑 1：内存溢出

第一次训练时直接把所有文本加载到内存，结果 8GB 内存不够用。后来改用 LineSentence，逐行读取，内存占用直接降到 1GB 以内。

坑 2：中文分词错误

jieba 默认分词会把 "深度学习" 切成 "深度" 和 "学习"，导致这个词根本学不到。需要加载自定义词典：

jieba.load_userdict("custom_words.txt")  # 一行一词，格式：词语 词频 词性

坑 3：向量维度选择

一开始贪心选了 1000 维，训练时间和存储都爆炸。后来实测 300 维在大多数场景够用，维度再高收益递减。如果语料本身不大，200 维也可以试试。

词向量的局限

Word2Vec 跑通后，解决了词级别的语义相似问题，但遇到句子级别的任务还是乏力：

  1. 句子向量怎么算？简单平均会丢失词序信息
  2. 多义词问题："苹果" 在水果语境和公司语境应该有不同的向量
  3. 固定窗口：某些长距离依赖关系学不到

我当时用一个笨办法：把句子里的词向量取平均，再做分类。效果比 TF-IDF 好一些，但明显不够用。

Sentence-BERT 实践

后来在做一个语义相似度匹配任务时，发现 Word2Vec 完全不够用。调研了一圈后选择了 Sentence-BERT（SBERT），它在 BERT 的基础上加了一个 siamese 结构，专门用来计算句子级别的嵌入。

环境准备

pip install sentence-transformers==2.2.2 torch==2.0.1

注意 torch 版本，2.0 以后某些接口有变化。如果用 GPU，确保 CUDA 版本匹配，我第一次在 RTX 4090 上跑直接报错，查了半天发现是 CUDA 12.0 和 torch 2.0.0 不兼容。

快速上手

from sentence_transformers import SentenceTransformer

# 加载预训练模型
model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')

# 编码句子
sentences = [
    "深度学习是机器学习的一个分支",
    "神经网络是深度学习的基础技术",
    "今天天气不错"
]
embeddings = model.encode(sentences)

# 计算相似度
from sklearn.metrics.pairwise import cosine_similarity
similarity = cosine_similarity([embeddings[0]], [embeddings[1]])[0][0]
print(f"相似度: {similarity:.4f}")  # 输出 0.6xxx 左右

模型选择

Sentence-BERT 提供了很多预训练模型，我试过几个：

模型速度准确度适用场景
all-MiniLM-L6-v2最快中等实时匹配、大规模检索
paraphrase-multilingual-MiniLM-L12-v2中等较高中英文混合、多语言
stsb-roberta-large较慢最高精度要求高的场景

我的项目主要处理中文，但偶尔有英文，最后选了 paraphrase-multilingual-MiniLM-L12-v2，速度和精度比较平衡。

踩过的坑

坑 1：批量大小设置不当

第一次直接用 model.encode() 跑 10 万条数据，直接把 16GB 内存打满。后来设置合理的 batch_size：

embeddings = model.encode(sentences, batch_size=32, show_progress_bar=True)

GPU 环境可以开到 64 或更大，CPU 上 32 比较稳妥。

坑 2：输入长度限制

BERT 模型默认最大长度是 512 token，长文本会被截断。我用了一个简单的处理策略：截取前 500 个字符，或者分段后取平均向量。

def truncate_text(text, max_length=500):
    return text[:max_length]

坑 3：相似度阈值选择

做语义匹配时，一开始定死 0.8 作为阈值，结果召回率太低。后来用验证集做了一个阈值曲线，发现 0.65 左右效果最好，但也因业务场景而异。

微调实践

预训练模型在通用任务上够用，但我的项目是垂直领域新闻分类，效果不理想。决定做微调：

from sentence_transformers import InputExample, losses, models
from torch.utils.data import DataLoader

# 准备训练数据：输入对和标签（0=不相似，1=相似）
train_examples = [
    InputExample(texts=["深度学习是机器学习的一个分支", "神经网络是深度学习的基础技术"], label=1.0),
    InputExample(texts=["深度学习是机器学习的一个分支", "今天天气不错"], label=0.0),
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

# 选择损失函数
train_loss = losses.CosineSimilarityLoss(model=model)

# 微调
model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=100)

微调后，相似度计算的准确率从 82% 提升到了 91%。但要注意，微调数据需要足够多样，否则容易过拟合。我一开始只用 500 对数据微调，结果在测试集上反而下降了。

垂直领域微调前后的准确率对比如下，差距在通用预训练模型上很难靠调参弥补：

Sentence-BERT 垂直领域微调前后相似度计算准确率对比（82% vs 91%）

9 个百分点的提升说明，在领域语义匹配任务上，少量高质量微调数据比换更大的通用模型更直接。

实际应用场景

语义搜索

这是我用嵌入技术最多的场景。传统关键词搜索只能匹配字面，用嵌入后可以搜索语义相似的内容：

# 构建索引
corpus_embeddings = model.encode(corpus_texts)

# 搜索
query_embedding = model.encode([query])
cosine_scores = cosine_similarity(query_embedding, corpus_embeddings)[0]

# 取 Top-K
import numpy as np
top_k = np.argsort(cosine_scores)[::-1][:5]
for idx in top_k:
    print(f"{corpus_texts[idx]} (分数: {cosine_scores[idx]:.4f})")

文本去重

做内容聚合时，需要识别重复或高度相似的文章：

def is_duplicate(text1, text2, threshold=0.85):
    emb1 = model.encode([text1])[0]
    emb2 = model.encode([text2])[0]
    similarity = cosine_similarity([emb1], [emb2])[0][0]
    return similarity > threshold

阈值设多少要看业务，新闻类的我用 0.85，评论类用 0.9 更严格一些。

文本聚类

做话题发现时，用嵌入 + KMeans 比传统方法效果更好：

from sklearn.cluster import KMeans

embeddings = model.encode(texts)
kmeans = KMeans(n_clusters=5, random_state=42)
clusters = kmeans.fit_predict(embeddings)

经验判断

什么时候用 Word2Vec

  • 资源受限：CPU 环境、内存小、需要高并发
  • 词级别任务：关键词提取、词云生成、同义词扩展
  • 语料足够大：几百万级别以上，才能训练出高质量的词向量
  • 不需要句子级语义：比如简单的文本分类，词级别特征够用

什么时候用 Sentence-BERT

  • 句子级语义：相似度匹配、语义搜索、问答对匹配
  • 精度要求高：宁愿牺牲速度换取更好的语义表示
  • 资源足够：至少 8GB 内存，有 GPU 更好
  • 多语言场景：需要同时处理中文和英文

成本考虑

Word2Vec 训练一次大概需要 2-4 小时（取决于语料大小），但推理速度极快，单次编码几毫秒。Sentence-BERT 推理一次大概需要 50-100ms（CPU），GPU 上可以降到 10ms 以内。

如果业务对延迟敏感，可以考虑用 Word2Vec 做初步筛选，再用 Sentence-BERT 做精确匹配。

结语

嵌入技术说到底就是把文本变成数字，让机器能算。但怎么算、算什么，还是需要人的判断。从 Word2Vec 到 Sentence-BERT，看起来是技术的升级，更多是场景的适配。

这次折腾下来，最大的感受是：没有银弹。Word2Vec 够用的地方，不要上 BERT；Sentence-BERT 能解决的问题，不要硬 Word2Vec。技术选型的关键不是谁更先进，而是谁更贴合当前场景。

embedding 技术还在快速发展，OpenAI 的 text-embedding-ada-002、Cohere 的 embed 模型都在提供 API 级别的服务。如果算力不够、语料不足，直接调用这些服务可能比自己训练更划算。但这又是另一个故事了。

可用性说明：本文发布于 2021 年 5 月，距今已超过五年。文中涉及的软件版本、接口、下载地址、命令参数和操作界面可能已经发生变化，部分方案在当前环境下可能失效。请结合官方最新文档核对后再操作，生产环境使用前务必先行验证。

版权声明: 本文首发于 指尖魔法屋-AI嵌入技术：Word2Vec不够用了之后（https://blog.thinkmoon.cn/post/248-ai-embedding-word2vec-sentence-bert-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章