Distillation and Refinement of Reasoning in Small Language Models for Document Re-ranking

本文提出InteRank方法，通过知识蒸馏和强化学习训练一个3B参数小型语言模型，在推理密集型文档重排序任务中生成解释并实现与70B+参数模型相当的性能，在BRIGHT基准上位列第三。

Knowledge Distillation, Reinforcement Learning, Reasoning, Document Ranking, Small Language Model, Interpretability

Chris Samarinas, Hamed Zamani

University of Massachusetts Amherst

Generated by grok-3

Background Problem

搜索引擎和检索增强生成系统面临越来越多的复杂查询，这些查询需要多步骤推理和综合分析来理解查询意图与文档内容之间的深层联系。传统排名模型依赖昂贵的人工标注数据，而现有神经排名模型在推理密集型任务（如BRIGHT基准）上表现不佳，nDCG@10仅约18%。同时，大型语言模型（LLM）虽在零样本重排序中展现出强大推理能力，但其计算成本和延迟限制了大规模部署。小型模型虽然更高效，但缺乏复杂的推理能力。因此，本研究旨在通过知识蒸馏和强化学习，训练一个紧凑的3B参数模型，在推理密集型文档重排序任务中实现与70B+参数模型相当的性能，同时提供可解释的排名决策。

Method

本文提出了一种名为InteRank的训练方法，用于在小型语言模型（3B参数）中提炼推理能力，应用于推理密集型文档重排序。具体方法如下：

核心思想：通过将排名过程分解为显式推理步骤，并在推理时投入计算资源进行逐步分析，提升小型模型的推理能力，而非直接预测相关性分数。
两阶段架构：第一阶段使用轻量级检索模型（如BM25或密集嵌入模型）获取前100个候选文档；第二阶段使用语言模型对查询-文档对生成推理解释，并输出离散相关性标签（0-非相关，1-部分相关，2-高度相关），最终结合检索分数和重排序分数（通过超参数 $\alpha=100$ 加权）进行排名。
训练流程：
1. 合成数据生成：从StackExchange等社交网站收集20K个问答对，提取链接文档作为正样本，并通过教师模型（70B参数Llama 3.3）生成相关查询和解释，构建包含查询、文档、解释和标签的训练数据集。
2. 知识蒸馏：通过监督微调，将教师模型的推理和标签能力转移到学生模型（3B参数Llama 3.2），优化目标为最大化教师生成输出的对数似然： $\theta^1 = \underset{\theta}{\text{arg}\,\text{max}} \,\mathbb{E}_{(q,d,e,l)\sim D_{\text{synth}}} \left[ \log p_{\theta}(e,l|q,d) \right]$ 。
3. 强化学习（RL）：通过采样多个输出（ $k=8$ ），使用奖励模型（8B参数Llama 3.1）评估解释质量，并对奖励值进行归一化和阈值筛选（ $\tau=0.85$ ），以更新模型参数： $\theta^{t+1} = \underset{\theta}{\text{arg}\, \text{max}} \, \mathbb{E}_{(q, d, \hat{y}, \overline{\mathcal{R}}) \sim D_t} \left[ (\overline{\mathcal{R}}(q, d, \hat{y}))^m \log p_\theta(\hat{y}|q, d) \right]$ ，其中 $m=3$ 为奖励缩放因子。
关键点：方法无需人工标注数据，强调生成解释对推理能力的重要性，并通过RL探索多样化的推理路径。

Experiment

实验在BRIGHT基准上进行，该基准包含1,384个推理密集型查询，覆盖StackExchange的7个领域（生物学、地球科学等）和5个编码与数学推理数据集，平均每个查询有6.37个正文档。实验设置如下：

模型与训练：学生模型为Llama 3.2 3B，教师模型为Llama 3.3 70B，使用QLoRA进行参数高效微调（4位量化，rank-64适配器），上下文长度限制为4K token，在单A100 GPU上训练（批量大小16，学习率2e-4）。RL阶段采样温度为1.0，奖励阈值为0.85，进行两轮训练。
基线对比：包括传统稀疏检索（BM25）、多种规模的密集检索模型（TAS-B 66M至Qwen1.5 7B）、交叉编码器重排序模型（MiniLM、ModernBERT-L）以及零样本LLM重排序器（Llama 3.2 3B和Llama 3.3 70B）。
结果分析：
1. InteRank（3B）在BRIGHT基准上平均nDCG@10达到27.4%（结合Qwen1.5检索器和领域特定指令），在排行榜上位列第三，仅次于JudgeRank（8B+70B+405B集成）和使用GPT-4查询重写的70B模型，优于所有其他基线，包括Reason-to-Rank（8B模型）。
2. 解释对性能至关重要，去除解释组件后性能从21.5%降至14.4%（nDCG@10），表明推理步骤对提升排名效果的贡献。
3. RL优化效果因领域而异，第一轮RL平均提升1.1%（nDCG@10），在数学和编码领域有显著改进（如TheoQ提升4.6%），但第二轮在部分简单领域（如生物学）出现下降，提示可能过拟合。
评价：实验设置较为全面，覆盖多种检索器和重排序器组合，BRIGHT基准的选择也合理，针对推理密集型任务。然而，RL阶段的效果不稳定，未探讨奖励模型设计对结果的影响，且缺乏对合成数据质量的深入分析，可能导致结果的高估。此外，实验未与其他小型模型的RL优化方法直接对比，方法的独特性可能被夸大。

Further Thoughts

本文提出的通过生成解释来提升小型模型推理能力的思路具有启发性，但也引发了一些值得深入探讨的问题。首先，合成数据依赖教师模型的输出质量，而教师模型本身可能存在推理偏差或领域局限性，未来可以探索多教师模型集成或数据多样性增强策略来缓解这一问题。其次，强化学习阶段的奖励模型设计对结果影响显著，但论文未详细讨论奖励模型的选择标准或潜在缺陷，是否可以引入更复杂的奖励机制（如结合人类反馈或多维度评估）来进一步优化推理路径？此外，本方法在推理密集型任务上的成功是否可以推广到其他领域，如对话系统或知识图谱推理，特别是在资源受限的边缘设备上部署小型模型时，是否需要额外的模型压缩技术（如量化或剪枝）来平衡性能和效率？最后，与近期一些工作（如DeepSeek R1在推理数据蒸馏中的泛化能力研究）结合，或许可以探索推理数据的领域适应性对不同模型蒸馏效果的影响，为小型模型的推理能力提升提供更广泛的应用场景。