[ACL 2023] Distilling Step-by-Step! Outperforming LLMs with Less Data and Smaller Model

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$ 为要训练的小模型，$\hat{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) $\hat{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 去同时预测 $\hat{y}$ 和 $\hat{r}$，相当于是让 student 去拟合 LLM 的 CoT $\hat{r}$ 来得到最终的答案 $\hat{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 $\hat{r}$ and label $\hat{y}$ into a single sequence $[\hat{r},\hat{y}]$ and treat the entire sequence as the target output)
  

References

• Hsieh, Cheng-Yu, et al. “Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes.” arXiv preprint arXiv:2305.02301 (2023).
• code: https://github.com/google-research/distilling-step-by-step