[논문 리뷰] Knowledge-Augmented Reasoning Distillation for Small Language Models in Knowledge-Intensive Tasks

Minki Kang, Seanie Lee외 3명|arXiv (Cornell University)|2023. 05. 28.

Topic Modeling인용 수 21

한 줄 요약

KARD는 대형 LLM에서 추론을 소형 LMs로 증류하고, 외부 지식으로 보강되며 합리적 근거가 포함된 구절을 검색하기 위해 신경망 재정렬기를 사용하고, 지식 중심 QA 벤치마크에서 강력한 성능을 달성한다.

ABSTRACT

Large Language Models (LLMs) have shown promising performance in knowledge-intensive reasoning tasks that require a compound understanding of knowledge. However, deployment of the LLMs in real-world applications can be challenging due to their high computational requirements and concerns on data privacy. Previous studies have focused on building task-specific small Language Models (LMs) by fine-tuning them with labeled data or distilling LLMs. However, these approaches are ill-suited for knowledge-intensive reasoning tasks due to the limited capacity of small LMs in memorizing the knowledge required. Motivated by our theoretical analysis on memorization, we propose Knowledge-Augmented Reasoning Distillation (KARD), a novel method that fine-tunes small LMs to generate rationales obtained from LLMs with augmented knowledge retrieved from an external knowledge base. Moreover, we further propose a neural reranker to obtain documents relevant to rationale generation. We empirically show that KARD significantly improves the performance of small T5 and GPT models on the challenging knowledge-intensive reasoning datasets, namely MedQA-USMLE, StrategyQA, and OpenbookQA. Notably, our method makes the 250M T5 models achieve superior performance against the fine-tuned 3B models, having 12 times larger parameters, on both MedQA-USMLE and StrategyQA benchmarks.

연구 동기 및 목표

• 프라이버시와 계산 제약으로 인해 지식 집중형 작업에서 소형 LMs의 필요성을 동기 부여한다.
• 외부 지식 KB 구절로 보강하면서 LLM의 추론을 소형 LMs로 증류하기 위한 프레임워크를 제안한다.
• 추론 시 합리적 근거 생성을 위한 구절을 검색하기 위해 신경망 재정렬기(neural reranker)를 도입한다.
• 기준선과 비교하여 MedQA-USMLE, StrategyQA, OpenBookQA에서 KARD의 성능 향상을 보여준다.

제안 방법

• 학습 데이터를 위해 체인 오브 사유(chain-of-thought prompting)을 통해 합리적 근거를 생성하도록 LLMs를 활용한다.
• 질문에 조건지어 합리적 근거와 정답을 함께 생성하도록 소형 LM을 미세조정한다.
• 근거를 질의로 사용하여 검색된 KB 구절(LKB)로 소형 LM 학습을 보강한다.
• 합리적 근거 생성을 위해 검색된 구절의 순서를 재정렬하도록 신경망 재정렬기를 도입한다.
• KL-divergence 목적함수를 사용하여 합리적 근거에 관해 재검색기의 랭킹을 모방하도록 재정렬기를 학습한다.
• 추론 시 구절을 검색하고 재정렬하고 합리적 근거를 생성한 뒤 최종 답을 산출한다.

실험 결과

연구 질문

• RQ1 지식 보강 증류가 지식 집중형 작업에서 LLM의 추론을 소형 LMs로 효과적으로 전달할 수 있는가?
• RQ2 외부 지식과 재정렬기를 추가하는 것이 표준 추론 증류를 넘어서 소형 LM의 성능을 향상시키는가?
• RQ3 의료 및 다중 모달 추론 벤치마크에서 KARD가 baselines(적은 샘플, 파인튜닝, 표준 추론 distillation)와 어떻게 비교되는가?

주요 결과

• KARD는 모델 크기에 관계없이 MedQA-USMLE, StrategyQA, OpenBookQA에서 일관되게 baselines를 능가한다.
• 지식 보강은 소형 LMs의 기억화 필요성을 줄여 더 적은 매개변수로도 더 나은 성능을 가능하게 한다.
• 신경망 재정렬기는 합리적 근거 생성을 위한 구절의 관련성을 높여 BM25 검색보다 더 나은 다운스트림 정답을 산출한다.
• KARD는 더 작은 모델(예: 250M 매개변수)에서 강력한 이득을 제공하며 때로는 더 큰 파인튜닝 모델을 능가한다.
• DAPT는 KARD에 비해 한정된 이득을 제공하며, 추론 증류에서 지식 보강의 고유한 가치를 강조한다.