让小模型学会何时求助大模型,提升推理效率与准确率
Knowledge Distillation with Training Wheels
- 将知识蒸馏建模为带熵正则的值优化问题,用策略学习实现师生协作
- 测试时按规则主动请求教师帮助,准确率提升且教师使用更智能
- 适合需要高效推理的场景,尤其在翻译与摘要任务中表现突出
知识蒸馏常用于生成式语言建模,通过大模型指导小模型训练以提升其能力。本文提出更通用的蒸馏框架:学生模型在训练中学习教师知识,并在测试时根据预设规则主动请求教师帮助。我们首先将知识蒸馏建模为熵正则化的值优化问题,采用路径一致性学习求解,得到一种基于同策略与异策略示范的新算法。进一步结合约束强化学习,构建可测试时调用教师模型的框架。此时,模型不仅学习内容,还需判断各部分难度以决定是否求助。实验在机器翻译与文本摘要任务中验证有效性,观察到准确率与教师使用率的趋势变化,表明该方法实现了传统推测解码无法达到的性能-效率平衡点。
原文摘要 · Abstract (English)
Knowledge distillation is used, in generative language modeling, to train a smaller student model using the help of a larger teacher model, resulting in improved capabilities for the student model. In this paper, we formulate a more general framework for knowledge distillation where the student learns from the teacher during training, and also learns to ask for the teacher's help at test-time following rules specifying test-time restrictions. Towards this, we first formulate knowledge distillation as an entropy-regularized value optimization problem. Adopting Path Consistency Learning to solve this, leads to a new knowledge distillation algorithm using on-policy and off-policy demonstrations. We extend this using constrained reinforcement learning to a framework that incorporates the use of the teacher model as a test-time reference, within constraints. In this situation, akin to a human learner, the model needs to learn not only the learning material, but also the relative difficulty of different sections to prioritize for seeking teacher help. We examine the efficacy of our method through experiments in translation and summarization tasks, observing trends in accuracy and teacher use, noting that our approach unlocks operating points not available to the popular Speculative Decoding approach.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。