[Paper Review] When MOE meets LLMs: Parameter Efficient Fine-tuning for Multi-task Medical Applications

LLM에서 MOE 기법을 적용하여 Medical Domain에서 Task를 수행하는 Reference를 찾아보면서 논문을 읽은 내용을 정리해 보려고 합니다.

• 논문 원문 사이트

When MOE Meets LLMs: Parameter Efficient Fine-tuning for Multi-task Medical Applications

---

Abstract

최근 Large Language Models (LLMs)의 급격한 증가가 여러 분야에서 큰 주목을 받고 있습니다.

Fine-tuning은 웹 기반 healthcare system과 같은 특정 도메인에 일반 LLMs을 적용하기 위해 종종 필요합니다.

그러나 의료 응용 분야에서 LLMs를 fine-tuning하는 과정에서 두 가지 문제가 발생합니다.

첫 번째는 task variety 문제로, 이는 실제 의료 시나리오에서 다양한 작업이 포함된다는 것을 의미합니다.

• 이러한 다양성은 data imbalance 및 seesaw 문제로 인해 sub-optimal fine-tuning을 초래하는 경우가 많습니다.

두 번째로, LLMs의 대규모 매개변수는 fine-tuning에 많은 시간과 computation 자원을 요구합니다.

• 이러한 두 가지 문제를 해결하기 위해, 우리는 MOELoRA로 명명된 multi-task 의료 응용을 위한 새로운 parameter efficient fine-tuning framework를 제안합니다.

설계된 framework는 multi-task learning을 위한 mixture-of-expert (MOE)의 이점과 low-rank adaptation (LoRA) parameter efficient fine-tuning의 이점을 모두 흡수하도록 목표로 합니다.

 

MOE와 LoRA를 통합하기 위해, 우리는 trainable parameter로서 여러 expert를 고안했으며, 각 expert는 trainable parameter의 소형 크기를 유지하기 위해 low-rank matrix의 일부로 구성됩니다.

 

그런 다음, 모든 MOELoRA layer에 대해 task-motivated gate function을 제안하여 각 expert의 기여도를 제어하고 다양한 task를 위한 distinct parameter를 생성할 수 있습니다.

---

Introduction

따라서, 이 논문에서는 open-source LLMs와 의료 지식 및 임상 작업에 대한 fine-tuning에 초점을 맞춥니다.

의료 도메인에 대한 LLMs의 fine-tuning은 두 가지 주요 도전 과제를 포함합니다.

1. Task Variety Problem: 실제 클리닉에서 LLMs는 doctor recommendation, diagnosis prediction, medicine recommendation, medical named entity recognition , clinical report generation등의 다양한 작업에 적용될 수 있습니다.
2. High Tuning Cost: fine-tuning이 표준 접근 방식이었던 BERT 시대 동안에도, LLMs의 매개변수 수가 너무 많아서 도전 과제가 되었습니다.

Task variety 문제와 관련하여, 여러 multi-task 학습 프레임워크가 제안되었습니다.

그 중에서도 Mixture-of-Experts (MOE)는 standout 중 하나로, 이는 task-shared와 task-specific 지식을 학습하기 위해 전문가를 분리하여 사용하고, 전문가 기여도를 조정하는 gate function을 통합하여 task 간의 균형을 유지합니다.

 

최근 parameter efficient fine-tuning (PEFT) 방법론이 이러한 높은 fine-tuning 비용 문제에 대한 잠재적인 해결책을 제공했습니다. 그러나 이러한 방식은 시간 소모적이고 지식 공유의 문제가 있습니다. 또한, fine-tuning이 가능한 작은 파라미터 셋이라 할지라도, 데이터 불균형과 seesaw 문제로 인해 성능이 저하될 수 있습니다.

Figure 1: 데이터 불균형 문제에 대한 설명.

이를 설명하기 위해, 우리는 중국의 의료 데이터셋과 해당 데이터셋의 샘플 분포를 Figure 1에 분석하였습니다

따라서, unique training process를 통해 separate parameters를 사용하는 multi-task parameter efficient fine-tuning이 이러한 문제를 동시에 해결할 수 있습니다.

 

Task variety와 high tuning costs의 문제를 해결하기 위해, 우리는 여러 작업에 대해 별도의 파라미터 efficient fine-tuning 프레임워크를 제안합니다. 이 프레임워크는 LoRA의 기본적인 parameter efficiency scheme을 채택하였으며, 즉, dense layers의 parallel한 작은 parameter 셋만 fine-tuning합니다.

 

이때, 우리는 MOELoRA를 소개합니다. 이는 MOE와 LoRA의 강점을 결합한 새로운 multi-task PEFT framework입니다. 또한, 우리는 각 task에 대해 distinct parameters를 생성하기 위해 task-motivated gate function을 설계했습니다.

---

PRELIMINARY

LLMs for Medical Applications

지능형 의료 시스템은 현대의 웹 기반 의료 환경에서 점점 더 보편화되고 있습니다

많은 연구들은 일관된 입력 및 출력 패턴을 정의하여 의료 작업을 표준화하려고 노력하고 있으며, 이를 통해 모델 설계 과정을 간소화하고 있습니다.

 

의료 텍스트가 포함된 medical tasks를 처리하기 위해 LLMs를 사용하는 예시

위에 제시된 의료 named entity recognition (NER)의 예시와 같이, 전통적인 모델들은 일반적으로 의료 텍스트, 즉 IMI을 처리하여 head entities Ohead와 tail entities Otail를 생성합니다.

의료 작업을 LLMs에 적응시키기 위해, 우리는 입력과 출력 패턴 모두를 재구성해야 합니다.

• Input Modification: 우리는 LLMs가 의료 텍스트를 실행할 수 있도록 원래 의료 텍스트에 instruction templates를 통합합니다. Figure 2에서 예시된 바와 같이, template를 수정하여 다음과 같이 만듭니다:
• "Please recognize the medical entity in this sentence: [Medical Text]". 여기서 "[Medical Text]"는 새로운 의료 작업 IM을 위한 자리표시자로 사용됩니다.
• Output Modification: 기존 target 대신, 우리는 인식된 head entity Ohead와 tail entity Otail를 template로 통합합니다.
• "The medical text has the following pairs of entities: [head entity] is [head entity] and tail entity is [tail entity]". LLMs가 NER task를 수행하는 방법은 다음과 같이 설명할 수 있습니다.

---

Multi-task Fine-tuning

앞서 언급한 바와 같이, 의료 응용 프로그램은 name entity recognition, medical inquiry 등 다양한 작업을 포함할 수 있습니다.

우리의 목표는 이러한 task에 대해 LLMs를 fine-tune하여 각 task의 성능을 높이고, 동시에 전체 healthcare system에도 혜택을 줄 수 있도록 하는 것입니다.

 

multi-task fine-tuning을 위해, 주어진 structured data Dj를 고려해보면, 각 task Tj의 목표는 LLMs를 fine-tune하여 적합한 출력 및 입력 패턴을 template 형태로 유지하는 것입니다.

fine-tuning 도중 각 task에 대한 데이터 D가 주어진다면, multi-task fine-tuning 문제는 다음과 같이 명확히 할 수 있습니다.

---

Method

제안된 프레임워크에 대한 종합적인 설명을 제공합니다.

제안된 방법의 개요로 시작하여 MOELoRA와 task-motivated gate에 대해 설명합니다.

마지막으로 fine-tuning 및 inference 프로세스를 자세히 설명합니다.

Overview

MOELoRA를 사용하는 LLMs의 parameter efficient fine-tuning 및 inference 프로세스를 시각적으로 나타낸 것입니다.

 

parameter efficient fine-tuning framework에서, LoRA는 dense layers에서 low-rank matrices를 대체하는 방법으로서 개념을 도입합니다.

이를 바탕으로, 우리는 MOELoRA layers를 각 layer에 통합하여 keys, queries, 그리고 values의 학습을 지원합니다. 또한, feed-forward network (FFN)을 예시로 하여 설명합니다.

추가적으로, 각 MOELoRA layer는 다양한 작업에 대해 다양성을 포착하기 위해 multiple experts를 통합합니다.

다음, task-motivated gate function을 도입하여 각 MOELoRA layer에 있는 expert들이 해당 task에 맞는 parameters를 학습할 수 있도록 합니다.

이 gate function은 MOELoRA layers에서 expert 간의 기여도를 결정합니다. 따라서, MOELoRA는 각 작업에 대해 distinct fine-tuned weights를 생성할 수 있습니다.

---

MOELoRA

Low-rank Adaptation (LoRA)는 parameter efficient fine-tuning에서 그 효과와 효율성을 입증했습니다.

LoRA는 intrinsic dimensionality 현상에서 영감을 받았으며, LLMs의 parameter fine-tuning 문제를 low-rank decomposition으로 재구성합니다.

이 decomposition은 주어진 matrix A가 low-rank 및 trainable하게 하는 방식으로 이루어집니다.

설정에서 LoRA layer와 결합된 linear layer의 forward 과정은 다음과 같이 표현될 수 있습니다:

식 4.

여기서 x는 차원 din의 입력 벡터를 나타내며, hhh는 차원 dout의 출력 벡터를 나타냅니다.

훈련 가능한 low-rank 행렬의 rank는 r이며, 이는 훈련 가능한 파라미터 수를 결정합니다.

그러나, 원래 LoRA에서는 모든 작업에 대해 일괄적으로 파라미터가 fine-tuning되므로, 의료 지식의 다양한 측면을 학습하는 데 어려움이 있습니다.

여기서는 LoRA와 MOELoRA에 대한 훈련 가능한 파라미터 수를 논의합니다.

LoRA에서는 두 개의 low-rank 행렬 B∈Rdin×r 및 A∈Rr×dout이 모든 훈련 가능한 파라미터를 포함합니다.

---

Task-Motivated Gate Function

이 섹션에서는 task-motivated gate function의 세부 사항을 다룹니다. 각 expert의 기여도가 특정 작업에 맞춰져야 한다는 점을 강조했습니다. 이러한 기여도를 조절하기 위해 우리는 gate function을 도입합니다.

이들은 inherently task-specific이므로, 각 gate function은 작업의 정체성을 기반으로 설계되었습니다.

이 설계는 작업별로 파라미터를 복구할 수 없게 만들며, 다음과 같은 두 가지 주요 이점을 제공합니다:

1. Task별 맞춤화: 각 작업은 별도의 파라미터 셋으로 fine-tuning되며, 이는 데이터 불균형 문제를 완화합니다.
2. Inference 시 효율성: 복구된 fine-tuned LLMs는 감소된 inference latency를 가집니다. 이는 MOELoRA layer와 연관된 추가 forward 계산이 필요하지 않기 때문입니다.

---

Fine-tune and Inference

MOELoRA의 fine-tuning 및 inference 과정을 설명합니다. 가독성을 위해, 우리는 Algorithm 1에서 결론을 내립니다.

• Fine-tuning: 우리는 LLMs에서 지정된 layer에 대해 MOELoRA를 설정하고 여러 하이퍼파라미터를 설정합니다 (라인 1-3). 이후, fine-tuning 전 파라미터 efficient fine-tuning을 위해 모든 사전 학습된 파라미터를 동결하고, 각 샘플을 무작위로 선택한 배치에서 수행합니다 (라인 4-7).

Algorithm 1 MOELoRA의 Fine-tuning 및 Inference 과정

1. Fine-tuning이 필요한 LLMs와 레이어를 지정합니다.
2. Rank 값 r과 Scale 값 α를 지정합니다.
3. MOELoRA의 Expert 수 N을 지정합니다.

Fine-tuning 과정
4. 사전 학습된 LLMs의 모든 파라미터를 동결합니다. 예: Wq, Wk, Wv.
5. 데이터셋 D의 샘플 배치 B에 대해,
6. MOELoRA를 사용하여 LLMs에 대한 forward 과정을 수행합니다 (위의 식 4 참고).
7. 손실 함수 L를 계산합니다 (식 2 참고).
8. MOELoRA의 파라미터 {Ai,Bi}i=1N 및 Gate Function의 파라미터 {E,WT}를 업데이트합니다.
9. 종료.

Inference 과정
10. 모든 작업 Tj에 대해,
11. 각 Expert에 대해 기여 가중치 ωj를 계산합니다 (식 5 참고).
12. 각 작업에 대해 식 8을 사용하여 MOELoRA의 fine-tuned 파라미터를 복구합니다.
13. 종료.
14. 특정 작업 Tj에 대해, 해당 작업에 필요한 LLMs의 파라미터를 적용하여 예측을 수행합니다.

• LLMs의 파라미터 (라인 4)가 동결됩니다. Fine-tuning 동안, 우리는 동일한 작업의 샘플을 하나의 배치로 그룹화하는 대신, 모든 작업에서 데이터를 무작위로 샘플링합니다.
• 이는 일부 multi-task 연구에서 수행되는 방식입니다. 우리는 실험에서 성능 비교를 위해 배치에 대해 무작위 샘플링을 선택했습니다.
• 이 데이터 배치를 사용하여 forward 과정을 수행하고 fine-tuning을 위한 손실을 계산할 수 있습니다 (라인 6-7). 파라미터 업데이트를 위해, 우리는 MOELoRA의 파라미터와 task-motivated gate function, 즉 {Ai,Bi}i=1N 및 {E,WT}만 fine-tuning합니다.

Inference: 앞서 설명한 바와 같이, MOELoRA는 각 작업에 대해 fine-tuned 파라미터 행렬을 복구할 수 있습니다 (라인 8).

• 그런 다음, 각 작업에 대해 LLMs 파라미터와 함께 해당 작업을 수행하는 데 필요한 파라미터를 적용할 수 있습니다.

---

Experiment

이 섹션에서는 다음 연구 질문(RQ)에 대해 다루고자 합니다.

• RQ1: MOELoRA가 다른 parameter-efficient fine-tuning 전략 및 cross-task generalization 방법과 비교하여 성능 측면에서 어떤 차이가 있는가?
• RQ2: MOE 아키텍처와 gate function이 fine-tuning 과정에 미치는 영향은 무엇인가? 다양한 훈련 전략이 MOELoRA의 성능에 어떤 영향을 미치는가?
• RQ3: MOELoRA의 expert 수와 rank가 성능 결과에 어떤 영향을 미치는가?
• RQ4: 제안된 MOELoRA가 fine-tuning 및 inference 과정에서 효율적인가?

Table 1: PromptCBLUE 데이터셋의 간단한 설명 및 통계

Task	Description	# Train	# Validation	# Test
CMeIE	Name Entity Recognition	2,828	600	600
CHIP-CDN	Normalization	2,381	600	600
CHIP-CDEE	Attribute Extraction	1,562	600	600
CHIP-MDCFNPC	Clinic Entity Discovery	4,935	600	600
CHIP-CTC	Medical Text Classification	3,622	1,100	1,100
KUAKE-QIC	Query Intention	3,279	660	660
IMCS-V2-MRG	Report Generation	1,799	600	600
MedDG	Doctor Dialogue	4,964	600	600

• RQ5: 전문가들이 다양한 작업에서 지식을 포착하는 데 있어 특화되어 있는가?

---

Experimental Settings

Dataset

우리의 실험은 multi-task Chinese medical dataset인 PromptCBLUE에서 수행됩니다.

이 데이터셋은 다양한 의료 작업을 포함하며, 이는 LLMs와의 호환성을 높이기 위해 텍스트 형식으로 변환되었습니다.

우리의 지식에 따르면, PromptCBLUE는 LLMs에 맞춰진 유일한 multi-task 의료 데이터셋입니다. 우리는 computational constraints로 인해 8개의 작업을 무작위로 선택하여 실험에 사용했습니다.

원본 데이터셋의 중복된 샘플을 제거한 후, 우리는 학습 세트를 테스트 세트로 사용하지 않도록 설정했습니다. 이 데이터셋의 통계 정보는 Table 1에 요약되어 있습니다.

 

Baselines

우리의 실험에서는 다음과 같은 네 가지 종류의 baselines와 비교합니다.

• LLMs without Fine-tuning: LLMs가 다양한 작업을 수행할 수 있도록 ConText Learning을 적용합니다.
  • ChatGPT: ChatGPT는 가장 인기 있는 LLMs 중 하나입니다. 우리는 task-relevant ability를 평가하기 위해, 학습 데이터에서 3에서 10개의 입력-출력 쌍을 무작위로 선택하여, 입력과 동일한 task로 데모를 구성합니다.
  • Hautuo: Hautuo는 중국 의료 데이터셋에서 수집된 instruct dataset을 구축합니다. 우리는 in-context learning 방식으로 ChatGPT baseline과 공정하게 비교하기 위해 ChatGLM-6B의 version을 사용합니다.

• LLMs with Fine-tuning: 이 그룹은 fine-tuning 전략의 변종을 사용하여 파라미터를 조정합니다.
  • P-Tuning: P-Tuning은 프롬프트 벡터를 생성하는 프롬프트 인코더를 fine-tuning하여 입력 시퀀스에 삽입합니다.
  • LoRA (Full): LoRA는 dense layers의 low-rank matrices를 사용하여 모든 사전 학습된 파라미터를 동결합니다.
  • LoRA (Single): 우리는 LoRA (Single)를 task별로 따로 LoRA를 훈련하여 구현합니다.
  • LoRA (Full+TP): 우리는 입력 텍스트에 간단한 task demonstration을 추가하여 LLMs가 작업 간의 구분을 인식하도록 합니다. 구현 측면에서는 LoRA (Full)와 동일한 훈련 과정을 진행합니다.

Multi-task에 대한 추가 작업에 따라, 우리는 모든 task 벡터를 함께 더하고 validation 세트에서 scale factor를 조정했습니다.

 

Cross-task Generalization: Multi-task fine-tuning에 대한 cross-task generalization의 적용 가능성을 평가하기 위해, 우리는 최근의 두 가지 접근법인 LoRAHub와 MoLoRA를 평가합니다.

• LoRAHub: LoRAHub는 source tasks에서 fine-tuned 된 LoRA 파라미터를 조합하여, unseen target tasks에 대한 일반화를 모색하는 조립 방법을 제안합니다.
• 실험에서, 우리는 각 작업을 LoRA로 fine-tuning하고, 지정된 작업의 validation을 사용하여 조합 가중치를 학습하며 이 작업의 성능을 테스트합니다.
• MoLoRA: MoLoRA는 비교적 최근의 연구로, MOE 구조를 LoRA에 채택합니다.

그러나 MoLoRA의 gate는 token의 중간 embedding을 사용하여 expert weights를 도출합니다.

우리의 실험에서, 우리는 이를 multi-task 설정에 맞추어, 즉, 동일한 작업 세트에서 훈련 및 테스트하도록 조정했습니다.

Implementation Details

우리의 실험은 PyTorch 1.12.0 및 Python 3.9.5를 사용하여 Tesla V100 32G GPUs에서 시뮬레이션되었습니다.

LLM ChatGLM-6B는 Chinese language processing에 능숙한 것으로 인정받아 fine-tuning을 위한 기본 모델로 사용됩니다.

 

모든 LoRA fine-tuning baselines 및 제안된 MOELoRA에 대해, 우리는 훈련 가능한 레이어를 "query_key_value", "dense", "dense_h_to_4h", 및 "dense_4h_to_h"로 지정했습니다.

 

최대 입력 및 출력 길이는 각각 1,024 및 196으로 설정되었습니다. 우리는 배치 크기를 64로 설정하고 최대 8,000 training steps까지 진행했습니다. LoRA rank rrr는 16으로 고정되었으며, LoRA dropout α=0.1 α=0.1로 설정되었습니다.

 

MOELoRA의 경우, experts의 수는 8로 설정되었으며, sparse gate 버전의 MOELoRA에 대해 최적의 값을 찾기 위해 KKK 값을 1에서 7까지 검색했습니다.

 

테스트 중에는, 우리는 generation을 위한 온도를 0.95로 설정했습니다. 우리의 MOELoRA 구현은 PEFT 패키지와 호환되어, 제안된 MOELoRA의 더 쉬운 채택과 활용을 촉진할 수 있습니다.

Evaluation Metrics

평가를 위해, 우리는 각 작업의 특성에 맞는 다양한 메트릭을 사용합니다.

예를 들어, CMeIE는 named entity recognition (NER) 작업으로, 많은 entity 클래스가 있습니다 (CMeIE에는 1,262개의 클래스가 있음).

Table 2: PromptCBLUE에서 경쟁적인 베이스라인 및 MOELoRA의 전반적인 결과.

굵은 글씨는 최고 점수를 나타내며, 밑줄은 해당 방법의 최저 점수를 나타냅니다.

"**"는 통계적으로 유의미한 개선을 나타냅니다 (즉, 양측 t-test에서 (p < 0.05)로 최고 베이스라인과 비교).

 

따라서 우리는 이 작업에 대해 자주 사용되는 Micro-F1을 적용합니다. CHIP-CDN (579 클래스), CHIP-CDEE (998 클래스), CHIP-MDCFNPC (2,065 클래스)는 모든 task가 너무 많은 카테고리를 가지고 있으므로, Micro-F1이 평가에 사용됩니다.

 

비교적으로, CHIP-CTC (44개 클래스)와 QUAKE-QIC (7개 클래스) 작업은 더 적은 수의 클래스를 가지고 있으며, 이는 각 클래스의 동일한 중요성을 고려해야 하므로 Macro-F1을 적용합니다. 텍스트 생성 작업, 예를 들어 IMCS-V2-MRG 및 MedDG의 경우, Rouge-L 이 적용됩니다. 또한, 전체 작업에 걸친 평균 점수는 전체 성능을 평가하는 데 사용됩니다.

 

결과의 견고성과 재현성을 보장하기 위해, 테스트는 무작위 시드 {42, 43, 44}로 세 번 실행되며, 평균 점수가 다음의 실험 결과에 보고됩니다.

---

Overall Performance (RQ1)

MOELoRA와 경쟁적인 베이스라인들의 종합적인 실험 결과는 Table 2에 나와 있습니다.

MOELoRA(D)와 MOELoRA(S)는 각각 MOELoRA의 dense 및 sparse gate 디자인을 나타냅니다.

전체 평균 점수를 분석한 결과, MOELoRA(D)는 다른 모든 방법보다 일관되게 뛰어난 성능을 보여줍니다.

RQ1에 대한 응답을 위해, 세부 분석은 다음과 같습니다.

• LLMs without Fine-tuning: Fine-tuning이 없는 LLMs 그룹은 task-specific medical knowledge를 통합하기 위해 fine-tuning이 얼마나 중요한지를 보여주면서, 성능이 상당히 뒤쳐집니다.
• Parameter Efficient Fine-tuning Strategies: 파라미터 효율적인 fine-tuning 전략 중, LoRA 기반 메서드가 명확하게 P-Tuning을 능가합니다.
  • LoRA (Full) 및 LoRA (Full+TP)는 모든 작업에서 데이터를 사용하므로 LoRA (Single)보다 우수하지만, 일부 task에서는 underperform합니다. 이는 task 간의 지식 공유의 중요성을 강조합니다.
• Model Editing: Task-Arithmetic은 task 벡터를 생성하는 방식으로, 파라미터 효율적인 fine-tuning에는 적합하지 않습니다.
• Cross-task Generalization: Cross-task generalization 환경에서 두 가지 최근의 접근법을 평가하였으며, 이들은 multi-task 설정에서 어려움을 겪었습니다.
• Dense Gate vs. Sparse Gate: Table 2에서 볼 수 있듯이, sparse gate가 두 가지 task에서 더 나은 성능을 보였으나, multi-task 의료 응용 프로그램에서는 shared medical knowledge가 더 중요했습니다.
  • dense gate는 모든 expert를 활용하여 학습된 지식을 공유하는 데 도움이 되므로, 대부분의 작업에서 우수한 성능을 나타냅니다.
• Task-specific Observations: 성능의 변동은 작업 간에 명확하게 나타납니다.
  • 예를 들어, LoRA (Full) 및 LoRA (Full+TP)는 데이터셋이 큰 작업에서 우수한 성능을 보이며, LoRA (Single)는 데이터 불균형 문제를 강조하여 샘플이 적은 작업에서 빛을 발합니다.
  • MOELoRA는 대부분의 작업에서 일관되게 최적의 성능을 달성하며, 이러한 불균형 문제를 효과적으로 해결할 수 있음을 보여줍니다. MedDG 작업의 경우, ChatGPT와 Hautuo의 고유한 대화 기능이 다른 접근법에 비해 우위를 제공합니다.

---

Ablation Study (RQ2)

RQ2를 더 깊이 연구하고 각 구성 요소의 기여도를 이해하기 위해, 우리는 Table 3에서 제시된 ablation 연구 결과를 제공합니다.

w/o MOE (본질적으로 LoRA (Full)로 되돌아감) 변형은 MOE 아키텍처를 제외합니다.

이 변형은 완전한 MOELoRA와 비교하여 성능이 저하됨을 보여주며, MOE 아키텍처의 중요성을 강조합니다.

마찬가지로, gate function을 우회하여 균일한 expert weights를 사용하는 w/o gate 변형도 MOELoRA보다 성능이 뒤처지며, gate function의 효과를 강조합니다.

 

w multiple gate 변형은 각 MOELoRA 레이어에 대해 고유한 gate function을 사용합니다.

이 변형이 일부 작업에서는 유사한 성능을 달성하지만, 과도한 파라미터화로 인해 single gate function 설계에 비해 약간 뒤처집니다

게다가, 다중 gate function은 더 많은 훈련 가능한 파라미터를 초래하여, 효율성이 감소하는 결과를 낳습니다.

 

추가적으로, 우리는 다양한 훈련 전략이 미치는 영향을 분석합니다.

특히, w BT 방법 [36]은 동일한 작업에서 샘플을 하나의 배치로 통합합니다.

반면에, w RBT 접근법 [39]은 데이터 배치마다 무작위로 작업을 선택합니다. 이들 두 가지 방법 모두 MOELoRA에 덜 유리한 것으로 나타나며, 성능 저하를 초래합니다.

이 성능 비교는 특정 훈련 패턴의 영향력을 강조합니다.

 

제안된 MOELoRA의 견고성을 검증하기 위해, 우리는 attention 레이어에 LoRA 레이어만을 부과하는 실험을 수행했으며, 이를 LoRA (Full)-QKV 및 MOELoRA(D)-QKV로 명명하였습니다.

결과에서, 우리는 MOELoRA(D)-QKV가 대부분의 작업에서 LoRA (Full)-QKV를 능가할 수 있음을 발견했으며, 이는 Table 2에서 MOELoRA(D)와 LoRA (Full)의 성능 비교와 일치합니다.

또한, MOELoRA(D)는 MOELoRA(D)-QKV보다 우수하므로, 더 많은 MOELoRA 레이어가 지속적으로 fine-tuning 성능을 향상시킬 수 있음을 보여줍니다.

---

Hyper-parameter Analysis (RQ3)

Figure 4.

RQ3에 답하기 위해, 우리는 MOELoRA(D)의 성능에 대한 하이퍼파라미터의 영향을 더 깊이 탐구합니다. 특히, expert 수 N과 LoRA rank r의 변동이 결과에 미치는 영향을 분석합니다.

Figure 4에 나타난 바와 같이, 우리의 관찰은 N이 0에서 8로 증가함에 따라 MOELoRA의 성능이 개선된다는 것을 보여줍니다.

이 향상은 더 많은 수의 experts가 더 광범위한 지식 스펙트럼의 학습을 촉진할 수 있다는 사실에 기인할 수 있습니다.

그러나 N이 16으로 설정되면, 성능이 약간 하락하는 것을 관찰할 수 있습니다.

이는 각 expert에 대해 작은 LoRA rank가 설정되어, low-rank 행렬의 학습 능력을 저하시킬 수 있기 때문입니다. 따라서, 각 expert의 rank를 2로 설정했습니다.

Figure 4b에서 볼 수 있듯이, r의 증가가 일관되게 성능을 향상시키지만, 동시에 훈련 가능한 파라미터의 크기가 비례하여 증가합니다. 효율성과 성능 사이의 균형을 고려할 때, r의 실용적인 선택은 16이 될 것입니다.

---

Hyper-parameter Analysis (RQ4)

Figure 5.

훈련 및 inference 효율성을 평가하기 위해, 우리는 Figure 5에서 훈련 가능한 파라미터의 비율과 inference latency를 비교합니다.

Inference latency는 inference 샘플의 수에 대한 inference 시간의 평균을 통해 계산됩니다.

 

MOELoRA(M)는 task-motivated gate가 동반된 MOELoRA의 변형을 나타냅니다.

결과는 MOELoRA가 LoRA (Full)와 동일한 수준의 높은 훈련 및 inference 효율성을 달성함을 보여줍니다. 이는 LLMs 파라미터의 0.48% 이상을 fine-tuning할 필요 없이 리소스를 절약할 수 있음을 의미합니다.

 

MoLoRA 및 MOELoRA(M)는 각 훈련 가능한 low-rank 레이어에 대한 추가 gate를 설정하므로 더 많은 훈련 가능한 파라미터가 필요합니다.

Inference에 있어, 모든 모델은 동일한 inference latency를 필요로 하지만, MoLoRA는 fine-tuned 파라미터를 Equation (8)과 같이 복구할 수 없으므로, 샘플에서 expert weights를 추출해야 합니다.

 

따라서, MoLoRA는 inference 시 추가적인 forward 계산이 필요하여, 더 많은 inference latency를 초래합니다.

이 비교는 task-motivated gate 설계의 이점을 보여줍니다. RQ4에 대한 응답으로, 설계된 MOELoRA는 높은 훈련 및 inference 효율성을 달성하며, task-motivated gate에 의한 효율성 저하를 방지합니다.

---

Case Study (RQ5)

RQ4에 대해, 우리는 Figure 6에 있는 네 가지 작업에 대한 expert weights의 시각화를 제시합니다.
각 작업에서, 다른 색상의 막대 길이는 해당 expert의 weights를 나타냅니다.

Figure 6

Expert weights는 1로 정규화되었기 때문에, 각 작업의 막대 길이는 동일합니다.

매크로 수준에서 볼 때, 각 expert의 기여도가 상당히 다르며, 이는 서로 다른 expert가 의료 지식의 다양한 측면에서 특화된다는 개념을 강조합니다.

 

또한, 작업 간 weight의 현저한 차이는 의료 응용의 다양한 특성을 강조합니다.

CHIP-CDN 및 KUAKE-QIC 작업을 자세히 살펴보면, 해당 작업에서 expert weights가 대부분 일치하지만, experts 3과 4의 경우는 예외적으로 다릅니다.

 

진단용 단어 정규화가 질의 분류를 강화할 수 있다는 점을 고려할 때, expert weights의 유사성은 MOELoRA가 관련 작업에 도움이 되는 공유된 지식을 잘 활용할 수 있음을 시사합니다.

---

Related Work

LLM for Medical Applications

최근, LLMs의 강력한 기능이 많은 분야에서 입증되었으며, 의료 도메인을 포함하여 큰 주목을 받고 있습니다.

예를 들어, Med-PaLM은 새로운 벤치마크인 MultiMedQA를 제안했으며, 이는 임상 지식 평가를 위한 새로운 진단 질의 응답 과제를 추가했습니다.

 

Med-PaLM2는 새로운 prompting 전략과 ensemble 강화로 Med-PaLM을 향상시켰습니다. 이 전략은 MedQA 및 자기 일관성에 기반하여, MedQA에서 상당한 성과를 보였습니다. ChatDoctor는 100,000개의 환자-의사 대화를 사용하여 LLMs을 fine-tuning하였으며, 널리 사용되는 의료 상담 플랫폼에서 파생된 것입니다.

 

또한, HuaTuo는 CMeKG로 처음으로 중국어 의료 지식을 학습하고, 중국어 의료 텍스트에서 LLMs를 fine-tuning하였습니다. 보다 구체적인 의료 응용을 위해, Liu et al.는 의료 LLM 기반 모델의 환각을 해결하기 위한 모델 편집 방법을 설계하였으며, Xu et al. 는 의료 LLMs에서 발생하는 환각을 해결하기 위한 모델 편집 방법을 설계하였습니다.

 

그러나 대부분의 이전 작업은 특정 의료 작업에 초점을 맞추고 여러 중요한 작업을 동시에 다루는 것을 간과합니다. 또한, 이러한 접근법은 일반화 능력을 달성하기 위해 상당한 fine-tuning 비용이 필요합니다.

Parameter Efficient Fine-tuning

Parameter Efficient Fine-tuning (PEFT) 방법은 fine-tuning 파라미터의 수와 계산 복잡성을 최소화함으로써 LLMs의 새로운 작업에 대한 성능을 향상시키려 합니다. Adapter Tuning 은 가벼운 adapter 모듈을 처음 소개하였으며, 이는 소수의 훈련 가능한 파라미터만을 포함합니다.

 

Prefix-tuning과 P-Tuning은 훈련 가능한 continuous prompts 또는 embeddings을 원래 시퀀스에 추가하는 task-specific 가상 토큰을 구성합니다. 그러나 prompts를 사용하는 것은 시퀀스 길이 제한으로 인해 길이가 긴 입력에서 어려움을 초래할 수 있습니다.

LoRA는 각 dense layer에 대해 훈련 가능한 low-rank matrices 두 개를 도입하며, inference 중 추가 계산 없이 full fine-tuning과 유사한 성능을 달성한 것으로 입증되었습니다.

그러나 LoRA fine-tuning은 multi-task 의료 응용에서 여전히 최적화되어야 합니다. LoRA 기반 PEFT 방법은 여전히 발전 중이며, 우리는 그 첫 걸음을 내딛고 있습니다.

---

Conclusion

이 논문에서는 LLM-driven 의료 응용을 위한 multi-task parameter efficient fine-tuning의 첫 번째 단계를 탐구합니다. 효율성과 성능을 만족시키기 위해, 우리는 MOELoRA라는 새로운 multi-task fine-tuning 프레임워크를 제안합니다.

 

구체적으로, 우리는 여러 low-rank 행렬로 구성된 MOELoRA 아키텍처를 설계하여, trainable 파라미터를 task-specific 지식 및 높은 효율성으로 학습할 수 있도록 합니다. 또한, 각 작업에 대한 독특한 fine-tuned 파라미터를 생성할 수 있는 task-motivated gate function을 설계했습니다.

 

중국 의료 데이터셋에서 광범위한 실험을 통해, 제안된 MOELoRA의 효과를 검증했습니다. 향후에는 지식 그래프와 같은 복잡한 의료 지식을 LLMs와 결합하는 방법을 탐구할 예정입니다.