当前位置：首页 > news >正文

Lychee Rerank MM模型蒸馏：基于Qwen2.5-VL的小型化重排序模型训练思路

news 2026/5/12 18:05:00

Lychee Rerank MM模型蒸馏：基于Qwen2.5-VL的小型化重排序模型训练思路

1. 项目背景与需求分析

多模态检索系统在实际应用中面临着一个关键挑战：如何在保证精度的同时提升推理效率。Lychee Rerank MM基于Qwen2.5-VL-7B模型，虽然提供了卓越的重排序精度，但其计算资源需求限制了在实际生产环境中的部署范围。

核心痛点分析：

16-20GB的显存占用要求，需要高端GPU才能运行
批量处理时的推理速度较慢，影响用户体验
部署成本高昂，难以在资源受限环境中使用

模型蒸馏技术为解决这一问题提供了有效途径。通过将大型教师模型的知识转移到小型学生模型中，可以在保持较高精度的同时显著降低计算需求。

2. 蒸馏方案设计思路

2.1 整体架构设计

基于Qwen2.5-VL的Lychee Rerank MM蒸馏采用师生框架，其中：

教师模型：原始的Qwen2.5-VL-7B模型，提供高质量的重排序信号
学生模型：选择参数量更小的多模态模型作为基础（如1-3B参数规模）

2.2 知识转移策略

软标签蒸馏是核心方法之一。教师模型不仅输出最终的排序得分，还提供丰富的中间表示：

注意力权重的分布模式
隐藏层的激活模式
输出层中"yes"/"no"token的概率分布

对比学习蒸馏同时采用，让学生模型学会区分相关文档与不相关文档的相对排序关系，而不仅仅是绝对得分。

3. 具体实现步骤

3.1 数据准备与处理

蒸馏过程需要构建高质量的训练数据集：

def prepare_distillation_data(query_doc_pairs, teacher_model): """ 准备蒸馏训练数据 """ distillation_data = [] for query, document in query_doc_pairs: # 获取教师模型的完整输出 with torch.no_grad(): teacher_output = teacher_model(query, document) # 提取软标签和中间表示 soft_labels = teacher_output.logits_softmax attention_maps = teacher_output.attention_weights hidden_states = teacher_output.hidden_states distillation_data.append({ 'query': query, 'document': document, 'soft_labels': soft_labels, 'attention_maps': attention_maps, 'hidden_states': hidden_states }) return distillation_data

3.2 损失函数设计

蒸馏损失函数结合多个目标：

class DistillationLoss(nn.Module): def __init__(self, alpha=0.7, temperature=3.0): super().__init__() self.alpha = alpha # 软标签权重 self.temperature = temperature self.kl_div = nn.KLDivLoss(reduction='batchmean') self.mse_loss = nn.MSELoss() def forward(self, student_output, teacher_output, hard_labels): # 软标签蒸馏损失 soft_loss = self.kl_div( F.log_softmax(student_output.logits / self.temperature, dim=-1), F.softmax(teacher_output.logits / self.temperature, dim=-1) ) * (self.temperature ** 2) # 硬标签损失（真实标签） hard_loss = F.cross_entropy(student_output.logits, hard_labels) # 中间表示蒸馏损失 hidden_loss = self.mse_loss(student_output.hidden_states, teacher_output.hidden_states) # 组合损失 total_loss = (self.alpha * soft_loss + (1 - self.alpha) * hard_loss + 0.3 * hidden_loss) return total_loss

3.3 训练流程优化

训练过程采用分阶段策略：

初始化阶段：使用软标签进行初步知识转移
精调阶段：结合硬标签和软标签进行联合训练
对比学习阶段：引入排序对比损失，提升区分能力

def train_distillation(student_model, teacher_model, dataloader, optimizer): student_model.train() teacher_model.eval() for batch_idx, batch in enumerate(dataloader): queries, documents, hard_labels = batch # 教师模型前向传播（不计算梯度） with torch.no_grad(): teacher_outputs = teacher_model(queries, documents) # 学生模型前向传播 student_outputs = student_model(queries, documents) # 计算蒸馏损失 loss = distillation_loss( student_outputs, teacher_outputs, hard_labels ) # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step()

4. 模型压缩与优化

4.1 模型结构优化

在学生模型设计上，采用以下优化策略：

参数共享：在多模态融合层引入参数共享机制
注意力头剪枝：减少注意力头数量但保持表征能力
层数减少：使用更浅但更宽的网络结构

4.2 推理加速技术

量化压缩是重要的后续优化步骤：

# 训练后动态量化 quantized_model = torch.quantization.quantize_dynamic( student_model, # 原始模型 {torch.nn.Linear}, # 要量化的模块类型 dtype=torch.qint8 # 量化类型 ) # 保存量化模型 torch.save(quantized_model.state_dict(), 'distilled_rerank_quantized.pth')