如何利用大模型蒸馏出小模型实现降本

如何让小模型的推理效果在某些领域比 ChatGPT 这样的大模型还要更强？这篇论文提供了一个思路：https://arxiv.org/abs/2212.10071，借助思维链（CoT）逐步解决复杂推理任务的能力，可以使用大模型作为推理教师，针对一批数据集，让大模型给出详尽的解答思路，然后把问题和解题过程交给学生模型进行 Fine-tuning。 这个解决思路也有一个专有名词，叫做“模型蒸馏”，其效果还是非常亮眼的，在保持同样的推理能力，甚至超越大模型的情况下，模型的大小降低到原来的 1/500~1/25。这可以帮助很多特定场景降低成本，例如使用自建的蒸馏小模型替代直接调用 ChatGPT，很多简单场景都适用，如果蒸馏出来的模型足够小，还可以直接在端侧（移动设备或嵌入式系统）部署，在用户本地完成推理，进一步降低服务成本。

最近的研究表明，思维链（Chain-of-Thought，CoT）提示可以引导语言模型逐步解决复杂的推理任务。然而，基于提示的CoT方法依赖于如GPT-3 175B等非常庞大的模型，这在大规模部署上是不可行的。在本文中，我们利用这些大型模型作为推理教师，以实现较小模型的复杂推理并将模型大小要求降低数个数量级。我们提出了Fine-tune-CoT，一种从非常大的教师模型中生成推理样本以微调较小模型的方法。我们在各种公共模型和复杂任务上评估了我们的方法。我们发现Fine-tune-CoT使得较小模型具有显著的推理能力，远远超过基于提示的基线甚至在许多任务中超过教师模型。此外，我们通过利用教师模型生成每个原始样本的多个不同解释的能力来扩展我们的方法。用这样多样化的推理丰富微调数据，可以在各种数据集上实现显著的性能提升，即使对于非常小的模型也是如此。我们进行了消融实验和样本研究，以了解学生模型推理能力的出现。

图1：Fine-tune-CoT利用由教师生成的推理来教授学生。我们通过零-shot思维链推理提示一个非常大的教师模型，如GPT-3 175B，来解决复杂问题。然后，我们使用推理样本来微调一个更小的学生模型。详情请参见图2。

图2：我们提出的Fine-tune-CoT方法的详细概述。步骤1：通过生成多步推理解释（绿色），提示一个非常大的教师模型来解决复杂问题（黄色）。步骤2：根据最终预测的正确性对完成进行过滤（红色）。问题、推理和答案用于构成推理样本，包括提示和多步解决方案。步骤3：策划过的推理样本用于微调一个小型、轻量级的学生以展示推理能力。基于LM的教师的应用实现了多样化的推理——为每个原始样本生成多个不同的解释，以丰富微调数据。这提升了学生模型的性能，而无需额外的人工注释。

实操代码

OpenAI API实验 OpenAI API实验已在oai模块中实现。请参考notebooks/example_oai_finetune_cot.ipynb，了解如何从头到尾运行Fine-tune-CoT。自定义实验（在GPU上） 自定义实验基于PyTorch Lightning实现，代码位于custom模块中。请参考custom_train.py和scripts/custom/*.sh，了解如何使用Fine-tune-CoT对模型（如T5、Flan-T5和GPT-2）进行微调。

开始搭建

pip install -r requirements.txtpython setup.py develop

环境

该代码已在Python<=3.10、PyTorch Lightning<=1.9、PyTorch>=2.0下进行测试。

数据

我们很自豪地分享我们的所有原始实验数据！所有数据都以json或jsonl格式组织，以供您使用 :)云存储文件夹链接：https://www.dropbox.com/sh/hwcncpyomx87h20/AACqgVdd-ZzBQ3ncJcKqw0cVa?dl=0https://drive.google.com/drive/folders/1C6kah3WV36N8omlUl-TeU9tsJADZNaJV

文件列表：

ldataset.tar.gz：以统一的json格式编制的12个任务数据集 放置在PROJECT/data/dataset/目录下lcompletion_data.tar.gz：所有教师和学生的完成数据，即所有实验的推理数据，解压后约8GB 放置在PROJECT/saved/completion_data/目录下lteacher_completion_data.tar.gz：Zero-shot-CoT（带有多样化推理）在默认教师模型text-davinci-002上使用OpenAI API的完成数据。是completion_data.tar.gz的子集，价值约$1000+，由KAIST的OSI实验室以为您提供。 放置在PROJECT/saved/completion_data/目录下lfinetune_data.tar.gz：用于通过微调API微调OpenAI学生的所有数据，以jsonl格式提供。这些数据源自教师完成数据，并可通过我们的代码生成。 放置在PROJECT/saved/finetune_data/目录下

生成论文结果

在下载完整的completion_data.tar.gz后，您可以运行notebooks/results.ipynb来生成我们论文中的所有结果表格和图表。该代码将（重新）评估包含在完成数据中的所有原始文本模型输出。

补充资源

基于模板的划分（论文附录 E.3） MultiArith和Date Understanding的基于模板的划分保存在/data/splits/*__template.json中。少量提示 从Wei 2022中适应的少量提示保存在/data/few_shot_cot_prompts.json中。

数据结构

data.dataset.Dataset

{"metadata": {"dataset_key": "multiarith"},"data": [{"sample_index": 0,"question": "string","answer": "string","rationale": "string?"}]}data.completion.CompletionDataset{"metadata": {"dataset_key": "multiarith","base_model": "curie","finetune_key": "zs_cot_multiarith","train_key": "ft_cot","prediction_template": "ft_cot_token",},"data": {"": [{"sample_index": 0,"completion_index": 0,"question": "string","answer": "string","prompt": "string","completion": "string","finish_reason": "string","reasoning_prompt": "string?","reasoning_completion": "string?","reasoning_finish_reason": "string?",}]}}

数据组织·

= B__T_

文件组织模式

saved/|–– completion_data/|–– B___C_/|-- D_.json # base model inference|-- F___D_.json # default fine-tuned model inference|-- F___T___D_.json # custom fine-tuned model inference|–– finetune_data/|–– P_/|–– F_{.*|/}|–– model_metadata/|–– B_|–– F___T_.json文件组织示例saved/|–– completion_data/|–– B_text-davinci-002__C_zs_cot/|–– B_text-davinci-002__C_zs_cot_long/|–– B_text-davinci-002__C_fs_cot/|–– B_curie__C_zs_cot/|–– B_curie__C_fs_cot/|–– B_curie__C_zs/|–– B_curie__C_ft_cot/|–– finetune_data/|–– F_zs_cot_multiarith/ # text-davinci-002_zs_cot|–– F_zs_cot_long_multiarith/|–– model_metadata/|–– B_curie/|–– F_zs_cot_multiarith.json