GTE模型蒸馏教程：训练轻量级学生模型

如果你正在用GTE这类文本表示模型，可能会遇到一个头疼的问题：模型效果确实不错，但体积太大，推理速度慢，部署起来成本高。特别是当你想把它塞进手机应用或者边缘设备里的时候，简直无从下手。

别急，今天咱们就来聊聊怎么给GTE模型“瘦身”。用知识蒸馏技术，把大模型（老师）的“本事”教给一个小模型（学生），让小模型在保持大部分性能的前提下，变得又快又小。整个过程就像老师带学生，把精华传授过去。

我会手把手带你走一遍完整的流程，从理解蒸馏原理，到设计损失函数，再到实际的训练技巧和代码实现。就算你之前没接触过蒸馏，跟着做下来也能自己训练出一个可用的轻量级GTE模型。

1. 知识蒸馏到底在干什么？

在开始动手之前，咱们先得把“知识蒸馏”这事儿想明白。别被这个词吓到，它的核心思想其实很直观。

想象一下，你是个经验丰富的老师（大模型），面对一个复杂的问题，你不仅能给出最终答案，还能解释为什么选这个答案，甚至能指出其他选项为什么不对。现在，你要教一个新手学生（小模型）。如果只告诉学生最终答案，他可能学得慢，还容易错。但如果你把解题思路、关键步骤、容易踩的坑都讲清楚，学生就能更快地掌握精髓。

知识蒸馏干的就是这个事儿。传统的模型训练，小模型只看标准答案（标签）来学习。而蒸馏呢，让小模型同时看标准答案，更重要的是看大模型给出的“软标签”。

什么是软标签？举个例子，大模型判断一张图片是“猫”的概率是0.9，是“狗”的概率是0.09，是“兔子”的概率是0.01。这个概率分布（[0.9, 0.09, 0.01]）就是软标签。它比硬标签（[1, 0, 0]）包含了更多信息，比如“这张图有点像狗，但非常确定是猫”。

对于GTE这样的文本表示模型，它的“知识”体现在它生成的文本向量里。同一个句子，大模型生成的向量和小模型生成的向量应该尽可能相似。同时，对于一对句子（比如一个问题和一段文本），大模型判断的相关性分数，也应该成为小模型学习的目标。

所以，我们蒸馏GTE模型，主要目标是让学生模型的输出向量，在方向和数值上都尽可能靠近老师模型的输出向量。这样，学生就学到了老师对文本语义的理解能力。

2. 动手前的准备工作

理论明白了，咱们就准备开干。你需要准备好环境和数据。

2.1 环境搭建

建议使用Python 3.8以上版本，并创建一个独立的虚拟环境。核心依赖库如下：

# 创建并激活虚拟环境（可选，但推荐）
python -m venv gte_distill_env
source gte_distill_env/bin/activate  # Linux/Mac
# gte_distill_env\Scripts\activate  # Windows

# 安装核心库
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118  # 根据你的CUDA版本调整
pip install transformers
pip install datasets
pip install sentence-transformers  # 方便加载和使用GTE模型
pip install tqdm  # 进度条
pip install scikit-learn  # 用于评估

如果你打算用阿里ModelScope的GTE模型，也可以安装：

pip install modelscope

2.2 数据准备

蒸馏需要训练数据。理想情况下，你需要一个文本对数据集，比如(查询, 相关文档)对。如果没有，用通用的文本语料库也行，比如维基百科文章、新闻文本等，我们可以从中构造对比学习任务。

这里我提供一个简单的示例，使用datasets库加载一个现成的句子对数据集（如STS-B），并演示如何构造用于蒸馏的数据批次。

from datasets import load_dataset
import torch
from torch.utils.data import DataLoader

# 加载一个公开的语义文本相似度数据集
dataset = load_dataset("stsb_multi_mt", "en", split="train")
print(f"数据集示例: {dataset[0]}")

# 假设我们只取前1000条做演示
train_texts_a = dataset["sentence1"][:1000]
train_texts_b = dataset["sentence2"][:1000]
# 原始相似度分数在0-5之间，我们可能需要归一化或直接使用
train_scores = torch.tensor(dataset["similarity_score"][:1000], dtype=torch.float)

# 定义一个简单的数据集类
class TextPairDataset(torch.utils.data.Dataset):
    def __init__(self, texts_a, texts_b, scores=None):
        self.texts_a = texts_a
        self.texts_b = texts_b
        self.scores = scores

    def __len__(self):
        return len(self.texts_a)

    def __getitem__(self, idx):
        item = {
            "text_a": self.texts_a[idx],
            "text_b": self.texts_b[idx]
        }
        if self.scores is not None:
            item["score"] = self.scores[idx]
        return item

# 创建数据集和数据加载器
train_dataset = TextPairDataset(train_texts_a, train_texts_b, train_scores)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

# 看看一个批次长什么样
batch = next(iter(train_loader))
print(f"批次大小: {len(batch['text_a'])}")
print(f"文本A示例: {batch['text_a'][0]}")
print(f"文本B示例: {batch['text_b'][0]}")
print(f"分数示例: {batch['score'][0]}")

在实际项目中，你可能需要更大、更多样化的数据。可以考虑MS MARCO、NQ等检索数据集，或者混合多个来源的文本。

2.3 加载老师模型（GTE）

接下来，我们把“老师”请出来。这里以Hugging Face上的Alibaba-NLP/gte-multilingual-base为例。这是一个多语言的GTE基础模型，效果不错，参数量也适中。

from transformers import AutoModel, AutoTokenizer
import torch.nn.functional as F

teacher_model_name = "Alibaba-NLP/gte-multilingual-base"
print(f"正在加载老师模型: {teacher_model_name}")

# 加载分词器和模型
teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)
teacher_model = AutoModel.from_pretrained(teacher_model_name, trust_remote_code=True)

# 将模型设置为评估模式，并放到GPU上（如果有）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
teacher_model.to(device)
teacher_model.eval()
print(f"老师模型已加载到 {device}")

# 写一个辅助函数来获取句子的向量表示
def get_embedding(model, tokenizer, texts, device):
    """获取一批文本的向量表示（取[CLS]位置并归一化）"""
    # 分词
    batch_dict = tokenizer(texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
    batch_dict = {k: v.to(device) for k, v in batch_dict.items()}
    
    # 前向传播，不计算梯度
    with torch.no_grad():
        outputs = model(**batch_dict)
        # 取[CLS]位置的向量作为句子表示
        embeddings = outputs.last_hidden_state[:, 0]
        # 对向量进行L2归一化，方便计算余弦相似度
        embeddings = F.normalize(embeddings, p=2, dim=1)
    return embeddings.cpu()  # 移回CPU以备后用

# 测试一下老师模型
test_texts = ["Knowledge distillation is a technique.", "It transfers knowledge from a large model to a small one."]
teacher_embeddings = get_embedding(teacher_model, teacher_tokenizer, test_texts, device)
print(f"老师模型生成的向量形状: {teacher_embeddings.shape}")  # 应该是 [2, 768]
# 计算两个句子的余弦相似度
cos_sim = torch.mm(teacher_embeddings, teacher_embeddings.T)
print(f"两句子的余弦相似度: {cos_sim[0, 1].item():.4f}")

老师模型准备好了，它能给任何句子生成一个768维的向量。我们的目标就是让后面要定义的学生模型，也能生成非常接近的向量。

3. 设计学生模型与蒸馏损失

现在轮到“学生”出场了。学生模型应该比老师模型小很多。我们可以选择一个更小的预训练模型架构作为起点，比如bert-base-uncased（1.1亿参数）甚至bert-tiny（400万参数），具体取决于你对速度和模型大小的要求。

3.1 定义学生模型

我们以bert-base-uncased为例，它比gte-multilingual-base小，但结构相似，蒸馏起来比较方便。

from transformers import BertModel, BertTokenizer

student_model_name = "bert-base-uncased"
print(f"正在初始化学生模型: {student_model_name}")

student_tokenizer = BertTokenizer.from_pretrained(student_model_name)
student_model = BertModel.from_pretrained(student_model_name)

student_model.to(device)
# 注意：学生模型开始需要训练，所以是train模式
student_model.train()
print(f"学生模型已加载到 {device}")

3.2 核心：设计蒸馏损失函数

这是最关键的一步。我们要让学生向老师学习两个东西：

1. 表示学习：学生模型自己生成的向量，应该和老师模型生成的向量尽可能相似。
2. 关系学习：对于一对句子，学生模型计算出的相似度分数，应该和老师模型计算出的分数接近。

因此，我们的损失函数通常由两部分组成：

def distillation_loss(student_emb_a, student_emb_b,
                      teacher_emb_a, teacher_emb_b,
                      labels=None, alpha=0.5, temperature=2.0):
    """
    计算蒸馏损失。
    参数:
        student_emb_a/b: 学生模型对句子A/B生成的向量
        teacher_emb_a/b: 老师模型对句子A/B生成的向量
        labels: 真实的相似度标签（可选）
        alpha: 控制蒸馏损失和真实标签损失的权重（0-1）
        temperature: 温度参数，用于软化概率分布
    """
    # 1. 表示损失 (Representation Loss)
    # 让学生模型的每个输出向量直接逼近老师模型的向量
    # 使用均方误差 (MSE) 或余弦相似度损失都可以
    mse_loss = torch.nn.MSELoss()
    rep_loss = (mse_loss(student_emb_a, teacher_emb_a) + mse_loss(student_emb_b, teacher_emb_b)) / 2.0

    # 2. 关系损失 (Relation Loss)
    # 计算学生模型预测的句子对相似度
    student_sim = torch.cosine_similarity(student_emb_a, student_emb_b, dim=1)  # 形状 [batch_size]
    # 计算老师模型预测的句子对相似度
    teacher_sim = torch.cosine_similarity(teacher_emb_a, teacher_emb_b, dim=1)  # 形状 [batch_size]
    
    # 使用MSE让学生的相似度预测接近老师的预测
    relation_loss = mse_loss(student_sim, teacher_sim)

    # 3. 可选：结合真实标签的损失 (如果labels存在)
    if labels is not None:
        # 将真实标签归一化到0-1之间（假设原始标签是0-5）
        normalized_labels = labels / 5.0
        label_loss = mse_loss(student_sim, normalized_labels.to(device))
        # 总损失是蒸馏损失和真实标签损失的加权和
        total_loss = alpha * (rep_loss + relation_loss) + (1 - alpha) * label_loss
    else:
        # 如果没有真实标签，只使用蒸馏损失
        total_loss = rep_loss + relation_loss
    
    return total_loss, rep_loss, relation_loss

参数解释：

• alpha: 这个参数很重要。如果设为1，学生只跟老师学；如果设为0，学生只跟真实标签学（这就变成普通训练了）。通常可以设一个中间值，比如0.7，让学生既学老师的“经验”，也兼顾真实数据。
• temperature: 在分类任务的蒸馏中常用，用于“软化”老师输出的概率分布，让标签更平滑。在我们这个回归任务（输出向量）中，有时可以不用，或者用一个较小的值。上面的代码暂时没体现，你可以尝试在计算相似度前对向量除以temperature。

这个损失函数是蒸馏的核心，它同时约束了学生模型的“内在表示”和“推理关系”都要像老师。

4. 完整的训练流程

万事俱备，只欠训练。我们把数据加载、模型前向、损失计算、反向传播的流程串起来。

import torch.optim as optim
from tqdm import tqdm

# 初始化优化器
optimizer = optim.AdamW(student_model.parameters(), lr=2e-5, weight_decay=0.01)
# 学习率调度器（可选）
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader)*3) # 假设训练3个epoch

num_epochs = 3
log_interval = 10  # 每10个批次打印一次日志

student_model.train()

for epoch in range(num_epochs):
    epoch_loss = 0.0
    epoch_rep_loss = 0.0
    epoch_rel_loss = 0.0
    
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
    
    for batch_idx, batch in enumerate(progress_bar):
        texts_a = batch['text_a']
        texts_b = batch['text_b']
        # 注意：STS-B的分数是0-5，我们需要它
        scores = batch['score'].to(device) if 'score' in batch else None
        
        # 1. 用老师模型获取“软标签”（向量和相似度）
        with torch.no_grad():
            teacher_emb_a = get_embedding(teacher_model, teacher_tokenizer, texts_a, device)
            teacher_emb_b = get_embedding(teacher_model, teacher_tokenizer, texts_b, device)
        
        # 2. 用学生模型前向传播
        # 对学生模型的输入进行分词
        inputs_a = student_tokenizer(texts_a, max_length=128, padding=True, truncation=True, return_tensors='pt')
        inputs_a = {k: v.to(device) for k, v in inputs_a.items()}
        student_outputs_a = student_model(**inputs_a)
        student_emb_a = F.normalize(student_outputs_a.last_hidden_state[:, 0], p=2, dim=1)
        
        inputs_b = student_tokenizer(texts_b, max_length=128, padding=True, truncation=True, return_tensors='pt')
        inputs_b = {k: v.to(device) for k, v in inputs_b.items()}
        student_outputs_b = student_model(**inputs_b)
        student_emb_b = F.normalize(student_outputs_b.last_hidden_state[:, 0], p=2, dim=1)
        
        # 3. 计算损失
        loss, rep_loss, rel_loss = distillation_loss(
            student_emb_a, student_emb_b,
            teacher_emb_a.to(device), teacher_emb_b.to(device),
            labels=scores, alpha=0.7
        )
        
        # 4. 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        # 可以添加梯度裁剪，防止梯度爆炸
        torch.nn.utils.clip_grad_norm_(student_model.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()  # 更新学习率
        
        # 累计损失用于日志
        epoch_loss += loss.item()
        epoch_rep_loss += rep_loss.item()
        epoch_rel_loss += rel_loss.item()
        
        # 更新进度条描述
        if batch_idx % log_interval == 0:
            progress_bar.set_postfix({
                'Loss': loss.item(),
                'Rep_L': rep_loss.item(),
                'Rel_L': rel_loss.item()
            })
    
    # 打印每个epoch的平均损失
    avg_loss = epoch_loss / len(train_loader)
    avg_rep = epoch_rep_loss / len(train_loader)
    avg_rel = epoch_rel_loss / len(train_loader)
    print(f"Epoch {epoch+1} 结束: 平均总损失={avg_loss:.4f}, 表示损失={avg_rep:.4f}, 关系损失={avg_rel:.4f}")

print("训练完成！")

训练过程会持续几十分钟到几小时，取决于数据量、模型大小和你的硬件。你会看到损失值逐渐下降，这意味着学生模型正在慢慢吸收老师模型的知识。

5. 效果评估与保存模型

训练完了，咱们得看看这个“学生”学得到底怎么样。

5.1 在测试集上评估

我们用一个留出的测试集，或者公开的基准（如STS-B的测试集），来比较学生模型和老师模型的表现。

from sklearn.metrics import mean_squared_error, mean_absolute_error
import numpy as np

# 切换到评估模式
student_model.eval()

# 假设我们有一个测试集 test_loader
# 这里我们用STS-B的测试部分做演示（实际应使用未参与训练的数据）
test_dataset = load_dataset("stsb_multi_mt", "en", split="test")
test_texts_a = test_dataset["sentence1"][:200]  # 取200条测试
test_texts_b = test_dataset["sentence2"][:200]
true_scores = test_dataset["similarity_score"][:200] / 5.0  # 归一化到0-1

all_student_sims = []
all_teacher_sims = []

with torch.no_grad():
    for i in range(0, len(test_texts_a), 16):  # 分批处理
        batch_a = test_texts_a[i:i+16]
        batch_b = test_texts_b[i:i+16]
        
        # 老师预测
        teacher_emb_a = get_embedding(teacher_model, teacher_tokenizer, batch_a, device)
        teacher_emb_b = get_embedding(teacher_model, teacher_tokenizer, batch_b, device)
        teacher_sim_batch = torch.cosine_similarity(teacher_emb_a, teacher_emb_b, dim=1).cpu().numpy()
        all_teacher_sims.extend(teacher_sim_batch)
        
        # 学生预测
        inputs_a = student_tokenizer(batch_a, max_length=128, padding=True, truncation=True, return_tensors='pt')
        inputs_a = {k: v.to(device) for k, v in inputs_a.items()}
        student_emb_a = F.normalize(student_model(**inputs_a).last_hidden_state[:, 0], p=2, dim=1)
        
        inputs_b = student_tokenizer(batch_b, max_length=128, padding=True, truncation=True, return_tensors='pt')
        inputs_b = {k: v.to(device) for k, v in inputs_b.items()}
        student_emb_b = F.normalize(student_model(**inputs_b).last_hidden_state[:, 0], p=2, dim=1)
        
        student_sim_batch = torch.cosine_similarity(student_emb_a, student_emb_b, dim=1).cpu().numpy()
        all_student_sims.extend(student_sim_batch)

# 计算评估指标
all_student_sims = np.array(all_student_sims)
all_teacher_sims = np.array(all_teacher_sims)

# 1. 学生与老师的相似度预测有多接近？（蒸馏效果）
mse_student_teacher = mean_squared_error(all_student_sims, all_teacher_sims)
mae_student_teacher = mean_absolute_error(all_student_sims, all_teacher_sims)
print(f"学生 vs 老师预测 - MSE: {mse_student_teacher:.4f}, MAE: {mae_student_teacher:.4f}")

# 2. 学生与真实标签的对比（下游任务效果）
# 注意：学生可能没直接学真实标签（如果alpha=1），但通过学老师间接学到了
mse_student_true = mean_squared_error(all_student_sims, true_scores[:len(all_student_sims)])
# 计算皮尔逊相关系数（更常用于STS任务）
corr_student = np.corrcoef(all_student_sims, true_scores[:len(all_student_sims)])[0, 1]
corr_teacher = np.corrcoef(all_teacher_sims, true_scores[:len(all_teacher_sims)])[0, 1]
print(f"学生 vs 真实标签 - MSE: {mse_student_true:.4f}, 相关系数: {corr_student:.4f}")
print(f"老师 vs 真实标签 - 相关系数: {corr_teacher:.4f}")
print(f"老师模型大小（参数量）: {sum(p.numel() for p in teacher_model.parameters()):,}")
print(f"学生模型大小（参数量）: {sum(p.numel() for p in student_model.parameters()):,}")

理想情况下，你会看到学生模型的预测和老师模型的预测非常接近（MSE和MAE很小），并且学生模型在下游任务（如语义相似度计算）上的表现，虽然可能略低于老师，但远好于一个从头训练的小模型，同时模型体积和推理速度却有巨大优势。

5.2 保存与使用蒸馏后的模型

评估满意后，就可以保存你的劳动成果了。

# 保存整个学生模型（包括结构和参数）
save_path = "./gte_distilled_student"
student_model.save_pretrained(save_path)
student_tokenizer.save_pretrained(save_path)
print(f"模型已保存到 {save_path}")

# 保存为PyTorch的state_dict（更轻量）
torch.save(student_model.state_dict(), "./gte_distilled_student.pth")

# 如何加载并使用
from transformers import BertModel, BertTokenizer
loaded_model = BertModel.from_pretrained(save_path)
loaded_tokenizer = BertTokenizer.from_pretrained(save_path)
loaded_model.eval()
# ... 使用方式与之前相同

现在，你就拥有了一个轻量版的GTE模型。你可以把它部署到资源受限的环境中，用于语义搜索、文本聚类、检索增强生成（RAG）等任务，享受更快的推理速度和更低的资源消耗。

6. 一些实用的训练技巧与注意事项

第一次跑可能不会那么完美，这里分享几个我实践中觉得好用的技巧：

1. 数据质量是关键：蒸馏效果很大程度上取决于老师模型在训练数据上的表现。如果某些数据老师都搞不定，学生更难学好。确保你的训练数据相对干净，并且覆盖了目标应用场景。
2. 损失函数可以更丰富：除了MSE，可以尝试余弦嵌入损失（torch.nn.CosineEmbeddingLoss）来直接约束向量方向，或者结合KL散度来匹配向量分布。有时组合多个简单的损失比一个复杂的损失更有效。
3. 渐进式蒸馏：如果学生和老师差距太大（比如从7B模型蒸馏到100M模型），直接蒸馏可能困难。可以尝试先蒸馏到一个中等模型，再用这个中等模型作为老师去蒸馏更小的学生。
4. 关注中间层：除了最终输出向量，老师模型中间层的特征图也包含丰富信息。可以考虑让学生模型的某些中间层输出也去匹配老师的对应层（这称为“隐层蒸馏”），但实现起来更复杂。
5. 温度参数调优：如果你在损失函数中引入了温度来软化目标，记得调整它。温度太高，所有目标变得模糊；温度太低，接近原始硬目标。通常需要在0.5到10之间尝试。
6. 验证集早停：一定要留一个验证集，监控模型在未见数据上的表现。当验证集损失不再下降时，就停止训练，防止过拟合到训练数据的噪声上。
7. 学生模型架构选择：不一定非要和老师架构相同。有时一个更小但更现代的架构（如DistilBERT、TinyBERT）作为学生，可能比同参数量但架构老旧的模型效果更好。

7. 总结

走完这一趟，你应该对如何给GTE模型做知识蒸馏有了清晰的把握。从准备数据和环境，到设计融合了表示学习和关系学习的损失函数，再到完整的训练循环和效果评估，每一步都是为了把大模型的“内力”有效地传递给小模型。

实际用下来，这套方法在大多数场景下都能得到不错的结果。学生模型通常能保留老师70%-90%的性能，而模型大小和推理速度却能改善好几倍。当然，你也会遇到一些挑战，比如如何选择最合适的损失权重alpha，如何处理特别难的数据等，这都需要根据你的具体任务进行调整。

如果你刚接触模型蒸馏，建议先从一个小规模实验开始，比如用几百条数据跑一个epoch，快速验证流程是否通畅。然后再逐步扩大数据量，精细调整超参数。这个过程本身也是学习模型压缩和优化思想的好机会。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。