Yongil Kim

Seoul National University

[ICML2023] Exploring the Benefits of Training Expert Language Models over Instruction Tuning

17 Dec 2023 » Transformer, PLM

[pdf] [github]

Joel Jang 1*, Seungone Kim 1*, Seonghyeon Ye 1*, Doyoung Kim 1*, Lajanugen Logeswaran 2, Moontae Lee 2,3, Kyungjae Lee 2, Minjoon Seo 1*
1 KAIST 2 LG AI Research 3 University of Illinois Chicago. Correspondence to: Joel Jang joeljang@kaist.ac.kr.  

image

Abstract

  • (Mutlitask prompting) LM 에 여러 가지 multitask 에 intruction tuning 을 진행하는 multitask prompted fine-tuning (MT) 을 useen task 에 대해서 좋은 능력을 보여왔다. 기존에, training task 의 수를 scaling 함으로써 성능 향상이 있다는 연구들이 많았다.
  • (Motivation) 저자들은 놀랍게도, 단 하나의 task 에 fine-tuned 된 Expert LM 이 300 개 이상의 태스크로 학습된 MT-LM 과 비교하여, BIG-benchmark 의 13개에 대해서 1.29%, 11개의 unseed dataset 에 대해서 3.20% 의 성능 우위가 있음을 발견하였다.
  • 이는 MT-LM 을 강력하게 하기 위해 task 의 수를 scaling 해야 한다는 기존의 연구에 의문점을 제시한다.
  • 이에 더해 single MT-LM 을 대신해 task 별로 seperate expert LM 을 학습시키는 것이 zero-shot inference 에 도움이 될 수 있음을 보인다. 이는 (1) instruction tuning 과정에서 종종 일어나는 negative task transfer 를 방지하고, (2) re-train 이나 catastrophic forgetting 없이 continual learning 을 가능하게하며, (3) 각각의 expert 를 혼합하였을 때 compositional capability 를 보인다.

Introduction

image

최근 Pretrained Language Model (PLM) 을 여러가지 task 에 instruction tuning 하는 MT-LM의 연구가 활발하다. 이는 성능이 매우 좋다고 알려져 있다. 그러나, 이 연구에서는 두 가지 파트로 나누어 MT-LM 의 current paradigm 에 의문점을 던진다.

Part1

기존에는 MT-LM 의 unseen task 에 대한 generalization 능력은 training 과정에서 배운 task 수에 scaling 한다는 연구가 많았다. 그러나 이 연구에서 우연히도, 단 하나의 task 를 배운 expert LM 이 300 개 이상의 task 를 배운 T0-3B 를 non-trivial margin 으로 이긴 것을 발견하였다.

이에 저자들은 T0-3B 를 학습시킨 296 개의 task 를 각각 하나씩만 배우게 expert LM 들을 학습시켰다. 이 296 개 중 7 개의 expert LM 이 T0-3B 의 unseen task 에 대해 더 높은 성능을 보인다(Figure 1). 이 7 개의 expert 로 부터 11개의 unseen task 를 측정했을 때 3.2%, Big bnech 에서는 1.29% 성능 우위를 보였다. 저자들은 또한 relevant expert 를 retrieve 하는 간단한 메커니즘을 통해 각각의 unseen task 에서 T0-3B 를 압도하는 성능을 얻을 수 있음을 보인다. 무려 12% 에 까까운 improvement 를 통해, 단순히 single MT-LM 을 나이브하게 학습시키는 것보다, 올바른 expert 를 choosing 하는 것이 더욱 효과적이고 효율적인 방법이라는 것을 보인다.

Part2

image

저자들은 위의 발견 외에도 RoE (Retrieval of Expert)가 MT-LM 보다 나은 세 가지 다른 advantage 를 발견한다.

  1. MT-LM 은 가끔 ‘seen’ task 에 대하여 negative task transfer 에 의해 sub-optimal performance 를 보인다. 이는 여러 가지 task 를 한 번에 배우는 것이, 오히려 특정 몇 개의 task 학습을 방해하는 것이다. 그러나 Expert LM 은 각각 의 task 를 독립적으로 학습하기 때문에 이러한 문제에서 자유롭다. 실험 결과, T0-3B 와 비교했을 떄, 36 개의 training task 에서 10.4% 에 해당하는 성능 우위를 보인다.

  2. MT-LM 은 catastrophic forgetting 문제가 있다. 그러나 RoE 방법은 이 문제가 전혀 없다. (absolutely no degredation)

  3. MT-LM 은 두 개의 task 를 composition 해야 할 경우 성능이 좋지 않은데 RoE 는 그렇지 않다. mT5-3B 두 개를 각각 summarization 과 translation expert 로 학습시킨 후, 이를 composition 했을 때, mT0-3B 과 summarization + translation 성능을 비교했을 때 우위를 보인다.

Expert Language Models

image

Training Experts

Training 과정에서는 Adapter 를 활용하여, parameter-efficient fine-tuning 을 진행한다. 이는 underlying LM 은 freeze 하고 adapter 부분만 학습하는 것이다. 위의 Figure 3 와 같이, 각각의 prompt 에 해당하는 task 를 배우는 Prompt Expert (PE) 들과, 각각의 Dataset 을 multiple training prompt 로 학습하는 Dataset Expert (DE) 로 나눌 수 있다. PE 를 학습할 때는 adapter 만 학습하고, DE 를 학습할 때는 전체 LM 을 다 학습한다.

Adapter 에 대해서 설명하면,

image

보통 Transformer 의 각 layer 에서는 위의 (1) 식의 hidden state 들을 (2) 식으로 self-attention 하는 과정들의 연속으로 이뤄진다. 이 때, Adapter 는

image

(1)식을 (3)식으로 바꿔서 hidden dimension e 로 보내는 FFN 을 추가한 뒤, 나머지는 freeze 하고 이 FFN 만 학습한다.

Retrieval-of-Experts (RoE)

이렇게 각각의 Expert 를 학습한 이후, Expert Library 를 구축 한 뒤, Dense retrieval 을 활용한다.

우선 저자들은 Expert Library 를 구성한다. 각각 Expert LM 이 학습한 S 개의 training instance 를 랜덤 샘플링 하여 library 를 구성한다. 따라서 expert library 의 크기는 [S X # of experts] 이다. 각각의 sentence 들로 부터 embedding 을 얻기 위해 Sentence Transformer 를 활용하였다.

이후 Retreival 과정에서는, 추론 과정의 target task 에서 $Q$ 개의 query 를 추출한다. 이후 $Q$ 개의 query 로 부터 MIPS (Maximum inner product search) 를 통해 Expert library 에서 $Q$ 개의 expert 를 가져온 후, 이 중 가장 많이 retreived 된 expert 를 선택한다.

마지막으로, Cold-fusion 이라는 연구에서, individually fine-tuned LM 을 merging 하는 것이 multitask fine-tuning 의 가능성을 보인다는 연구에 따라, Retrieved 된 expert LM 를 합쳐서 새로운 expert LM 을 제시한다. 합치는 방법은

image

위와 같으며, 여기서 $\tau$ 는 vanialla pre-trained LM 과 expert LM 과의 parameter difference $(\theta_{expert} - \theta_{vanilla})$ 이다. 따로 언급이 없으면 $\lambda$ 는 1/N 으로 expert LM 들을 uniformly merging 한다.

Experimental Setup

Training setup

T0 의 36 개 training dataset 을 활용한다. Prompt 는 T0 의 것을 활용하며, 296 개의 prompt (task) 가 있다. 따라서 296 개의 PE 와 36 개의 DE 가 생성된다. LM-adapted T5 model checkpoint 를 baseline 으로 활용하였다. Epoch 은 5 이며, lr 은 1e-4, expert library 를 위한 S=100 이다.

Evaluation setup

비교를 위한 MT-LM 은 T0-3B 와 T0-11B 이고, RoE 는 T5-3B + DE/PE 이다. T0 original paper 의 세팅 처럼, 11개의 unseen dataset 은 4 개의 category 가 되고, BIG-bench 로부터 13개의 dataset 을 활용한다. 추가적으로 T0 가 배우지 않은 8 개의 new generative task 를 활용한다. Inference 때, RoE 를 위한 $Q$ 는 32 로 고정한다.

Expert LMs Can Generalize to Unseen Tasks

Expert LM 이 새로운 패러다임이 될 수 있음을 실험적으로 검증한다. 아래의 Table 1 은 11 개의 unseen dataset 에 대한 결과, Table 2 는 BIG-Bench 13 개에 대한 결과, Table 3 는 8 개의 unseen genertavie task 에 대한 결과이다.

image

image

image

우선 Table 1 에서, T5(3B) + Cos PE (Cosmos-QA dataset 에 no_prompt_text) 가 T0-3B 를 11개 중 8개를 앞질렀다. 이는 MT-LM 의 scaling 에 대한 기존 연구 결과를 뒤집을 수 있는 결과이다. Table 2 에서 역시 Cos PE 가 가장 높은 mean acc 를 보인다. Table 3 에서도 T5 + Sam PE (Samsum dataset 에 ‘given the above dialog wirte a summary’ prompt) 가 T0-3B 를 8개 평균 6.83 점을 앞선다.

그리고 또, Table 1 에서 RoE 과정에서 Oracle 로 expert LM 을 가져와서 best performing 을 측정했을 때, T0-3B 뿐 아니라 더 큰 T0-11B, 심지어 GPT-3 보다도 각각 11.94%, 2.61%, 4.37% 증가한 것을 볼 수 있다. Table 3 에서 Oracle 은 13.69 점이나 증가하였다.

마지막으로, Oracle 이 아닌 RoE 방법을 통한 T5 + PE w/ RoE 는 11개 중 8 개의 unseen task 에서 T0-3B 를 앞질렀다. Oracle 과 비교했을 때, 여전히 성능 개선 여지가 충분하기 때문에 retriever side 에서 개선의 여지가 충분한 것도 볼 수 있다.

Merging of Experts

image

Table 4 에서 (Mer.) 표시는 Expert LM 들을 merging 한 것이다. 첫 세 줄은 PE LM과 merging LM 의 결과, 아래 세 줄은 DE LM 과 merging LM 의 결과이다. RoE 의 경우, merging 을 하더라도 COPA 등 몇 개의 경우 positive task transfer 가 있었지만, 대부분의 경우에서 negative task transfer 이 있었다.

이에 분석을 위해, Full LM training 을 하는 DE 를 merging 한 것이 아래의 세 줄의 결과이다. merging 을 한 것이 대부분의 결과에서 가장 좋거나, 두 번째로 좋은 결과를 내기 때문에, DE merging 은 negative task transfer 없이 composition ability 를 보인다고 주장할 수 있다.

Analysis of Experts

image

다시 Figure1 으로 돌아가서 저자들은 세 가지 측면의 분석을 제시한다.

첫 번째로, 8 개의 Training task category 중 유일하게 Multiple-Choice QA (MCQA) task 가 좋은 generalization 성능을 보인다. 이에 저자들은 11 개의 classification setting task 가 QA 형태를 instruction 에 필요로 하기 때문이라고 가정한다.

두 번째로, 36개 training dataset 에 대해, COSMOS-QA, SOCIAL-I_QA, DREAM 3 개의 training dataset 에 대해서만 consistently PE 든 DE 든 성능이 좋다. 이 세 데이터셋은 모두 commonsense reasoning dataset 이고, 이는 unseen task 에 대한 generalization 에서 필수불가결하다.

마지막으로, T5 + SAM PE 가 Table 3 에서 가장 좋은 성능을 보인다. SAM PE 는 SAMSUM dialog summarization dataset 에 대한 Expert LM 이다. 그러나 이 모델은 Table 1,2 dml classification setting 에서는 T0-3B 보다 10% 가까이 안좋아서 there’s no free lunch 를 보여준다.

Benefits of Expert LMS over MT LMs

Seen Task Performance

image

먼저, expert LM 이 negative task trasnfer 에 영향을 적게 받음을 보이기 위해, T5(3B) + PE W/ ROE의 성능을 36개의 Validation datset 에 대해, 두 MT LM 모델 T0-3B 및 T0-11B 과 비교한다. 위의 표에 나타난 대로, 각각 mean accuracy 에서 T0-3B 및 T0-11B보다 각각 +10.40% 및 +7.70% 더 높은 성과를 보인다.

이는 평가가 seen instruction 으로 이루어지기 때문에, 간단한 검색 메커니즘이 expert library 에서 best-performing expert 를 선택할 가능성이 높기 때문이다. 이는 T5(3B) + PE W/ ROE(ORC.)와 유사한 성능을 나타내는 것에서 알 수 있다. 실제로 T5(3B) + PE W/ ROE는 보이는 작업 중 296개에서 280개의 작업에서 동일한 Training 데이터셋에서 PE 를 검색하며, 296개의 작업 중 185개에서 동일한 prompt 와 dataset(oracle 에 해당하는)에서 PE를 검색한다.

Continual Learning of New Tasks

image

모델 배포 이후 추가 데이터셋에서 언어 모델을 미세 조정하고자 할 때, 미세 조정된 LM 을 continual learner 로 만드는 것이 중요하다 (Chakrabarty et al., 2022). 전체 original and additional task in each update 는 계산 부담이 크기 때문이다. 이전 연구는 Rehearsal-based 방법을 통해 이 문제를 해결하며, Fine-tuned LM을 original and additional dataset 에 conitnual learning 시킨다. (Chakrabarty et al., 2022). 그러나 이 접근 방식은 (1) 원본 데이터에 액세스할 수 있다고 가정하고 (2) instruction tuning 중 additional 샘플을 continual trainig 시키는 데 여전히 추가 계산 부담이 발생한다.

이 연구에서는 각각 별도의 Expert LM 을 각 additional task에 대해 training 시켜 전문가 라이브러리에 단순히 추가하는 distributed training 을 통해 original and additional dataset 에 액세스하지 않고도 동일한 결과를 얻을 수 있다는 것을 보여준다. 구체적으로, MT-LM (T0-3B) 을 continually training 시켜 CT0-3B 로 만드는 방법과 제안하는 distributed approach 간의 비교를 Table 6 에 제시했다.

표의 결과를 보면 제안하는 방식이 seen task 대한 성능 저하가 전혀 없고, unseen task 에 대한 경미한(-0.15%) 성능 저하를 보인다. 게다가, 평균적으로 +1.08 의 성능 우위가 MT LM 대조군에 비해 존재함을 보여준다. 이로써 original 데이터에 액세스하거나 무거운 계산 비용이 들지 않는 상태에서 distributed approach 는 대부분의 경우 원래 능력(seen task and unseen task)을 유지할 뿐만 아니라 target task 에서 CT0-3B를 능가한다.

Compositional Instructions

image

우리는 아래처럼 두 개의 instruction 을 합쳐서 줄 수 있다: “Write a summary of the following English text and translate the sentence into Korean.” where “Write a summary of the following English text.” and “Translate the sentence into Korean.” are two separate instructions seen during training.

이 compositional capability 를 테스트하기 위해, mT0-3B 를 MT-LM 으로 하고, 5 개의 summarization 과 translation 의 compositional task 를 학습시켰다. 이후, 제안하는 distributed approach 로 mT5-3B 두 개에 각각 summuarization 과 transaltion 을 학습시킨 후 ,Merging 을 했을 때, 5 개 중 4개의 task 에서 좋은 성능을 보였고, 한국어와 일본어 같은 low-resoruce language 에서는 더 큰 차이를 보였다. 왜냐하면, low-resource language 는 학습 과정에서 negative task transfer 에 의해 학습이 방해되기 때문이다. Table 8 에서는 cherry-picked 된 결과를 보여준다.

Conclusion


Expert language models trained on single tasks exhibit strong generalization to unseen tasks, surpassing multi-task language models by a significant margin, showcasing benefits in robustness, adaptability, and compositional instruction performance. The proposed distributed approach encourages exploration of collaborative expert training for potential future advantages in efficiency, privacy, and personalization, not explicitly covered in this paper (see limitations and discussion in Appendix D).