Unsupervised Information Refinement Training of LLMs for RAG

Unsupervised Information Refinement Training of LLMs for RAG

소개

본 논문은 ACL 2024 메인 컨퍼런스에서 발표된 "Unsupervised Information Refinement Training of Large Language Models for Retrieval-Augmented Generation"에 대한 리뷰입니다. 이 연구는 Retrieval-Augmented Generation (RAG) 상황에서 다양한 검색된 텍스트에 대처할 수 있도록 언어 모델의 파인튜닝 방법을 제안합니다.

RAG의 개요

RAG는 사용자의 쿼리에 적합한 출력을 생성하기 위해 다음과 같은 3단계 파이프라인을 구성합니다:

1. 검색 (Retrieval): 쿼리가 주어졌을 때 지식 베이스에서 관련 패시지를 추출합니다.
2. 증강 (Augmentation): 원본 쿼리와 추출된 패시지를 바탕으로 증강된 쿼리를 생성합니다.
3. 생성 (Generation): 증강된 쿼리를 바탕으로 최종 답변을 생성합니다.

RAG의 한계점

RAG의 주요 한계점은 검색 단계에서 추출된 모든 텍스트가 유익하지 않을 수 있다는 것입니다. 이는 다음과 같은 문제를 야기할 수 있습니다:

• 잘못된 정보 포함
• 정답 생성에 필요한 정보 부족

이러한 문제를 해결하기 위해 두 가지 접근 방식이 있었습니다:

1. 프롬프트 기반 방식: 잘못된 패시지가 입력되었을 때 프롬프트를 통해 올바른 정보로 변환합니다. 하지만 이 방식은 모델 내부 파라미터가 업데이트되지 않아 근본적인 한계가 있습니다.
2. 모델 파라미터 업데이트 방식: 특정 태스크에 대한 RAG 데이터를 구축하여 언어 모델을 학습시킵니다. 그러나 이 방식은 일반화 성능을 잃을 수 있고, 데이터 구축에 많은 비용이 발생합니다.

제안된 방법: InfoRAG

저자들은 이러한 한계점을 해결하기 위해 "InfoRAG"라는 비지도 학습 방법을 제안합니다. 이 방법은 기존의 위키피디아와 같은 텍스트 데이터를 활용하여 효율적으로 모델을 학습시킵니다.

InfoRAG의 주요 특징

1. 모델 내부 파라미터 업데이트 가능
2. 기존 데이터를 입력과 출력 형태로 변환하여 저비용으로 다양한 태스크에 대한 일반화 성능 확보

InfoRAG의 세 가지 시나리오

1. 모든 필요한 지식이 검색된 텍스트에 있는 경우
   • 목적: 복잡하고 긴 검색 텍스트에서 관련 없는 정보를 배제하고 직접적이고 간결한 텍스트만 생성
2. 검색된 텍스트에 부정확한 지식이 포함된 경우
   • 목적: 언어 모델의 내부 지식을 활용하여 잘못된 지식을 수정하고 누락된 지식을 완성
3. 검색된 텍스트에 문제 해결을 위한 정답이 없는 경우
   • 목적: 언어 모델의 내부 지식을 활용하여 질문 해결

방법론

데이터 수집 및 전처리

1. 위키피디아 데이터에서 K개의 연속된 문장을 추출 (K = 15)
2. 각 문장을 prefix와 target으로 분리하여 언어 모델의 입력과 출력으로 구성

시나리오별 학습 방법

1. 시나리오 1
   • 입력: K개의 연속된 문장과 i번째 문장의 prefix
   • 출력: i번째 문장의 target
2. 시나리오 2
   • 입력: K개의 연속된 문장(각 문장에 informative token 추출 및 랜덤 마스킹/대체 적용)과 i번째 문장의 prefix
   • 출력: i번째 문장의 target
3. 시나리오 3
   • 입력: i번째와 i+1번째 문장을 제외한 K개의 연속된 문장과 i번째 문장의 prefix
   • 출력: i번째 문장의 target

Informative Token 선택 방법

1. 각 레이어마다 next word distribution 산출
2. 최종 레이어와 가장 차이가 큰 중간 레이어 선택 (Jensen-Shannon Divergence 사용)
3. 선택된 레이어에서 JSD 값이 가장 큰 상위 50% 토큰을 informative token으로 선택

학습 과정

• 세 가지 시나리오에 대한 데이터를 2:4:4 비율로 혼합하여 언어 모델 학습

실험 결과

1. 제로샷 설정에서의 성능 향상
   • 다양한 태스크에서 로버스트하게 성능 향상
2. Few-shot 학습에서의 성능
   • 예시를 추가할수록 베이스라인 대비 높은 성능 향상
3. 멀티스텝 추론에서의 성능
   • 각 추론 단계의 성능 향상
4. 다양한 검색 텍스트 시나리오에서의 성능
   • 모든 시나리오에서 효과적인 성능 향상
5. Ablation Study
   • 단순 위키피디아 추가 학습보다 InfoRAG 방식이 더 효과적
6. 각 시나리오별 학습 방법의 효과
   • 시나리오 2와 3은 단일 및 다중 학습에서 모두 성능 향상
   • 시나리오 1은 다른 시나리오와 결합 시 성능 향상
7. Catastrophic forgetting 방지
   • 기존 모델의 언어 생성 능력을 유지하면서 RAG 태스크에서 성능 향상

결론

InfoRAG는 다양한 검색 텍스트 시나리오에 대응할 수 있는 효과적인 언어 모델 학습 방법을 제안합니다. 실험 결과를 통해 다양한 상황에서의 로버스트한 성능 향상을 입증하였으며, 각 시나리오에 적용한 학습 방법의 적합성을 검증하였습니다.

InfoRAG 사용 방법 (30단계)

1. 위키피디아 데이터셋 준비
2. 데이터 전처리: K개의 연속된 문장 추출 (K = 15)
3. 각 문장을 prefix와 target으로 분리
4. 베이스 언어 모델 선택 (예: LLaMA 7B, 13B)
5. 시나리오 1 데이터 생성
6. 시나리오 2 데이터 생성을 위한 informative token 추출 준비
7. 각 레이어의 next word distribution 계산
8. Jensen-Shannon Divergence를 사용하여 최종 레이어와 중간 레이어 간 차이 계산
9. 가장 큰 차이를 보이는 중간 레이어 선택
10. 선택된 레이어에서 JSD 값이 가장 큰 상위 50% 토큰을 informative token으로 선택
11. 선택된 informative token에 대해 마스킹, 대체, 유지 작업 수행
12. 시나리오 2 데이터 생성 완료
13. 시나리오 3 데이터 생성
14. 세 가지 시나리오 데이터를 2:4:4 비율로 혼합
15. 학습 하이퍼파라미터 설정 (learning rate, batch size 등)
16. LoRA 방식으로 파인튜닝 준비
17. 모델 학습 시작
18. 주기적으로 검증 세트에서 성능 평가
19. 최적의 체크포인트 선택
20. 테스트 데이터셋 준비
21. 제로샷 설정에서 성능 평가
22. Few-shot 설정에서 성능 평가
23. 멀티스텝 추론 태스크에서 성능 평가
24. 다양한 검색 텍스트 시나리오에서 성능 평가
25. 기존 언어 생성 능력 유지 여부 확인 (예: MLM 태스크)
26. 결과 분석 및 해석
27. 모델 최적화 (필요시 하이퍼파라미터 조정 및 재학습)
28. 최종 모델 선택
29. 실제 RAG 시스템에 통합
30. 지속적인 모니터링 및 업데이트

이 30단계를 따라 InfoRAG를 구현하고 사용할 수 있습니다. 각 단계는 세부적인 구현 details가 필요할 수 있으며, 실제 적용 시에는 사용 가능한 컴퓨팅 리소스와 데이터셋의 특성에 따라 조정이 필요할 수 있습니다.