Introduction
- 作者提出了新的 LLM distillation 架构 Distilling Step-by-Step,通过 multi-task training 同时预测 CoT 和 label,相比传统的 SFT / KD 能在使用更少数据量/更小模型的条件下达到更高的精度;此外,Distilling Step-by-Step 在训练时甚至不需要加载 teacher model,极大地降低了训练开销
Method
- Standard finetuning and task distillation.
其中, f f f 为要训练的小模型, y ^ \hat y y^ 为 GT label (i.e., standard finetuning) 或 LLM 生成的 label (i.e., task distillation)
- Multi-task learning with rationales. 为了缩小 student 和 teacher 之间的 gap,可以在训练的时候给 student 输入由 teacher 生成的 Chain-of-Thought (CoT)
r
^
\hat r
r^ (we require users to produce a few example demonstrations (∼ 10-shot for all tasks) in order to use the few-shot CoT prompting mechanism)
但这种设计在 student 推理时仍然需要用 LLM 预先生成 CoT,这显然是不现实的;为此,作者改为在训练时让 student 去同时预测 y ^ \hat y y^ 和 r ^ \hat r r^,相当于是让 student 去拟合 LLM 的 CoT r ^ \hat r r^ 来得到最终的答案 y ^ \hat y y^,从而让 student 在训练时就学得了 CoT 的能力
需要注意的是,作者采用的是 multi-task training,通过在输入前加入不同的 “task prefixes” (
[label]
,[rationale]
) 来要求 student 生成 label / CoT,而不是做 single-task training 把 label 和 CoT 拼接到一起做预测,消融实验里有证明 multi-task training 的效果优于 single-task training (simply treating rationale and label predictions as a single joint task may harm the model’s performance on label prediction)
Experiments
- Setup. teacher – 540B PaLM model; student – T5 models.
- Reducing training data. (1) Distilling step-by-step outperforms standard finetuning with much less labeled examples.
(2) Distilling step-by-step outperforms standard distillation with much less unlabeled examples.
- Reducing model size. (1) Distilling step-by-step improves over standard baselines across varying model sizes used. (2) Distilling step-by-step outperforms LLMs by using much smaller task-specific models. (3) Unlabeled data augmentation further improves Distilling step-by-step. In Figure 7 on SVAMP, we augment the SVAMP training set with unlabeled examples from the AS-Div dataset.
- Ablation studies. (1) Distilling step-by-step works with different sizes of decently trained LLMs.
(2) Multi-task training is much more effective than single-task rationale and label joint prediction. 作者采用的 multi-task training 效果优于 single-task training (i.e., concatenate the rationale r ^ \hat r r^ and label y ^ \hat y y^ into a single sequence [ r ^ , y ^ ] [\hat r, \hat y] [r^,y^] and treat the entire sequence as the target output)