Zhu, Yutao, et al. "One token can help! learning scalable and pluggable virtual tokens for retrieval-augmented large language models." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 39. No. 24. 2025.
https://ojs.aaai.org/index.php/AAAI/article/view/34813
Abstract
현존 RAG 방법들은 프롬프트 기반으로 LLM이 검색된 정보를 활용하게 하거나, LLM을 RAG 용도로 파인튜닝한다. 파인튜닝 방법은 성능이 좋지만 파라미터를 변경하므로 LLM의 범용적인 생성 능력이 약해진다. 이러한 한계는 하나의 LLM이 여러 목적으로 사용중일 때 실용성을 떨어트린다.
본 논문에서는 확장성이 있으며 탈부착가능한 RAG용 가상 토큰을 학습시키는 SPRING 방식을 제안한다. LLM 파라미터는 보존하고 별도 토큰의 임베딩만 파인튜닝하므로, LLM의 성능은 높이고 일반화 능력은 보존된다. 또한 학습전략에 따라 이 방식은 뛰어난 확장성, 유연성, 일반성을 가진다.
Methodology
프롬프트 방식과 파인튜닝 방식의 효율성 모두 얻기 위한 방법으로 본 논문에서 제시하는 SPRING을 사용한다.
공식
언어 모델은 자연어 문장의 확률 분포를 계산한다. 여기서 사용되는 auto-regressive 모델은 다음 토큰을 예측한다.
$p_{\text{LM}}=\prod_{i=1}^{m}p_{\theta}(x_{i}|x_{<i})$
- $x_{<i}$ : $x_i$ 이전의 토큰 문장
- $\theta$ : LLM 파라미터
RAG 시스템에서는 검색 코퍼스($\mathcal {D}$)와 검색기($M$)가 존재한다. RAG에서의 생성 과정은 사용자 입력($x_{<i}$)과 검색된 결과($R = M_{\mathcal {D}}(x_{<i})$)를 활용해 수식화된다.
$p_{\text{RAG}}=\prod_{i=1}^{n}p_{\theta}(x_{i}|R;x_{<i})$
RAG를 통한 QA 시스템에서는 여기에 질문($Q$)가 추가된다.
$p_{\text{RAG-QA}}=\prod_{i=1}^{m}p_{\theta}(a_{i}|R;Q;a_{<i})$
QA 작업의 학습 목표는 질문 $Q$와 현재까지 생성한 토큰 $a_{<i}$를 통해 정답 $A=\{a_i\}^m_{i=1}$를 생성하는 것이다. 검색기는 여러 개의 문서를 가져올 수 있으며, `\n\n`와 같은 문서 구분자를 사용해서 만든 긴 텍스트 시퀀스가 $R$이 된다. $R$ -> $Q$ 순서로 사용한다. (간결함을 위해 수식에서 복잡한 프롬프트는 생략됨)
확장 가능하고 탈부착가능한 RAG용 가상 토큰
SPRING은 RAG 용도로 학습 가능한 가상 토큰을 LLM에 입력한다. 앞선 $p_{\text{RAG-QA}}$의 검색 결과 $R$과 질문 $Q$ 사이에 n개의 토큰 $T=[t_1,t_2, ..., t_n]$을 추가한다. 생성 과정은 다음과 같다.
$p_{\text{SPRING}}=\prod_{i=1}^{m}p_{\theta,\delta}(a_{i}|R;[t_{1},t_{2},\cdot\cdot\cdot,t_{n}];Q;a_{<i})$
- $\delta \in \mathbb{R}^{n \times d}$ : 가상 토큰의 학습가능한 파라미터
- $d$ : LLM 임베딩 크기
- $\theta$ : 백본 LLM의 파라미터 (학습 시 고정)
이러한 별도 파라미터 추가 방식은 매우 효율적인데, $|\delta| \ll |\theta|$ 이기 때문이다. 예시로 Mistral-7b 모델의 임베딩 크기는 4096인데 가상 토큰 50개 추가하는 경우, 50 x 4096 = 0.2M 파라미터만 학습시키면 된다. 이는 전체 모델의 0.003%이므로 파인튜닝 대비 매우 효율적이다.
가상 토큰 T를 검색 결과 R과 질문 Q 사이에 배치하는 이유는 두가지이다.
- LLM은 Causal 어텐션을 사용하므로 검색 결과 뒤에 가상 토큰이 있어야 검색 결과를 활용할 수 있다.
- 최근 LLM이 입력의 끝부분에 특히 예민하다는 것이 밝혀졌다. LLM이 질문을 잘 이해하는 것이 가장 중요하므로 가상 토큰을 중간에 배치한다.
Scalable
LLM은 최대 입력 길이로 자주 제한되므로, 검색된 많은 토큰을 활용하지 못한다. SPRING은 검색된 문서의 길이에 따라 가상 토큰 개수를 조절하면서 RAG 추론 성능을 높일 수 있다. 이를 위해 전체 토큰 n개 중 랜덤하게 k개를 선택해서 사용하는 학습 시퀀스 $[R, t_1, t_2, ..., t_k;Q]$을 구축한다. 가상 토큰의 수가 고정되지 않으므로, 추론 시목적에 따라 가상 토큰 수를 조절할 수 있다.
Pluggable
SPRING에서 유저의 입력이 외부 지식을 요구하면 RAG를 위한 가상 토큰이 추가될 수 있다. 반대로 외부 지식이 불필요한 경우 가상 토큰 없이 바로 LLM을 사용할 수 있다. LLM 파라미터를 변경하지 않기 때문에 LLM 본연의 능력은 그대로 사용할 수 있다.
Inference
학습을 통해 가상 토큰 임베딩은 RAG용으로 최적화된다. 하지만 이 토큰은 모델의 어휘 공간에 존재하지 않는다. 별도의 스페셜 토큰(ex. [r1], …, [r50])을 어휘 공간에 추가하며, 앞서 학습된 임베딩을 스페셜 토큰 임베딩으로 사용한다. 추론 시에는 검색된 결과 R 뒤에 스페셜 토큰(ex. [r1], …, [rk])를 추가하고 질문 Q까지 포함해서 LLM에 전달한다. 이를 통해 SPRING은 효과적으로 검색결과와 유저 입력의 갭을 메꾸며, 외부 지식을 사용하는 LLM의 능력을 상당히 개선시킨다.
Experiment
- 학습 데이터셋 : 8개의 qa 데이터셋 (TriviaQA, NQ, HQA 등)
- 유의사항 : 정답과 함께 주어지는 완벽한 출처(golden passage)를 사용하지 않고, 검색된 문서만 사용해서 실용성을 높였다.
- 평가 데이터셋 : PopQA, Fermi, Musique, Bamboogle
- 일반화 능력을 확인하기 위한 데이터셋
- 평가 시 test set 사용, 없는 경우 dev set 사용
- 평가기준 : EM(Exact Match), F1 score
- 검색 문서 : Wikipedia & MS MARCO
- Wikipedia에는 높은 퀄리티의 지식이 있으며, MS MARCO는 방대한 크기의 웹페이지 문서를 포함한다.
- 검색기 : E5-large (BM25와 비교 시 고성능)
Baseline method
사전학습 모델과 지시어 파인튜닝된 모델을 모두 백본으로 사용해서 비교한다. LLM은 Mistral-7b, LLaMA-2-7b, LLaMA-2-13b를 사용한다. 아래 방식들과 SPRING을 비교한다.
- Concat : 검색 결과와 질문을 단순히 합침
- Prompt : 검색 정보를 활용할 수 있게 직접 만든 프롬프트 사용
- Prefix-tuning : LLM 첫 입력에 50개의 prefix를 추가함
- LoRA : 백본을 효율적으로 파인튜닝함
- LoRA 적용 후 SPRING까지 적용한 모델을 $\text{SPRING}^+$이라고 함
구현 디테일
- Pytorch, HuggingFace 가속 라이브러리 사용
- 가상 토큰 임베딩은 아래와 같은 프롬프트의 임베딩으로 초기화
- “According to the previous relevant passages, please answer the following question. Only return the answer without any other words”
- `if (토큰 개수 > 프롬프트)` : 위 문장을 반복
- `if (토큰 개수 < 프롬프트)` : 프롬프트를 분할
- 가상 토큰 임베딩을 랜덤하게 초기화하는 경우, 프롬프트 초기화보다 성능이 떨어짐
실험 결과
RAG, non-RAG 상황에 대해 테스트를 진행했다. SPRING 추론 시 가상 토큰은 50개만 사용했다.
- SPRING은 프롬프트 방식에 비해 RAG 성능이 상당히 좋다. 특정 데이터(TriviaQA, CoQA)에서는 LoRA 파인튜닝 방법도 넘어선다. 높은 성능임에도 0.2M의 파라미터만 학습하면 되므로 효율적이기도 하다.
- LoRA도 상당히 성능이 좋지만 non-RAG 상황에서는 파인튜닝을 하지 않은 모델보다 성능이 떨어진다.
- non-RAG 상황에서는 SPRING, SPRING+만 프롬프트 방식보다 높은 성능을 보인다. 검색을 하지 않더라도 가상 토큰이 효과적임을 알 수 있다. 가상 토큰이 LLM이 목적과 포맷을 이해하게 돕기 때문으로 추측된다.
- LoRA의 체크포인트를 사용한 SPRING+가 가장 높은 성능을 보인다. 또한 모든 백본 LLM은 SPRING을 추가했을 때 성능이 좋아진다. 따라서 SPRING은 다양한 LLM에 대해 RAG 시나리오에서 성능을 높이는데 적합하다.
- 프롬프트 방식은 RAG에서 LLM 성능을 높이나, 학습하지 않는 방법이므로 성능 개선에는 한계가 있다.
- SPRING은 훈련에 사용되지 않은 데이터셋에 대해 강건한 성능을 보이므로 일반화 능력이 입증된다.
- prefix-tuning 기법은 SPRING보다 RAG 성능이 떨어진다. SPRING처럼 가상 토큰을 검색된 문서 뒤에 사용하는 것이 더 합리적이고 효과적임을 알 수 있다.
추가 분석
추가적으로 다양한 세팅에 대해 SPRING을 평가한다. Mistral-7b-instruct 모델로 실험을 진행한다.
non-RAG 에서의 성능
파인튜닝 방식이 LLM 본연의 능력에 미치는 영향을 확인한다. LoRA와 SPRING에 의해 파인튜닝된 모델을 non-RAG 작업에서 평가한다.
- 평가 데이터셋 : BoolQ, CommonsenseQA, GSM8K, MMLU (수학 능력, 지식 이해)
- SPRING은 가상 토큰을 사용하지 않고 LLM만 사용할 수 있으므로 LLM 고유 능력이 보존된다. LoRA는 RAG를 위해 LLM의 파라미터를 변경하므로, non-RAG 작업에서는 성능이 떨어진다.
- few-shot 평가에서 성능이 저하되므로, LoRA 파인튜닝은 in-context 학습 능력을 악화시킨다. RAG를 위한 입력 포맷에 모델이 오버피팅하게 되므로, 일반적인 생성 능력(지시를 따르는 능력)은 떨어지게 된다.
토큰 퀄리티의 영향
추론 시 사용할 가상 토큰의 개수는 임의로 조절할 수 있다. 이러한 scalable 장점을 확인하기 위해 다양한 가상 토큰 개수로 학습한 SPRING과, 고정된 가상 토큰 개수(k=50)으로 학습된 SPRING을 비교한다.
일반적으로 SPRING의 성능은 가상 토큰 수에 따라 같이 증가한다. 놀라운 점은 하나의 가상 토큰만 사용하더라도 RAG 작업에서 성능이 개선된다는 것이다. 고정된 수의 토큰으로 학습하는 경우, 동일한 토큰 수에서 추론 시에만 성능이 나오므로 유연성이 떨어진다.
리트리버 종류에 따른 효과
SPRING은 E5-large에 의해 검색된 문서로 파인튜닝된다. 다른 리트리버 사용했을 때의 성능을 조사하기 위해, BM25, BGE-base, E5-large로 검색된 문서로 학습을 진행한다.
전체 리트리버에 대한 실험에서 SPRING은 프롬프트 방식보다 높은 성능을 보인다. 또한 성능 차이(Variance)가 적으므로 검색기의 품질 변화에 덜 민감하고 강건하다. 그리고 고성능 리트리버로 학습된 SPRING을 추론 시, 저성능 리트리버와 사용해도 여전히 성능을 유지한다. 따라서 SPRING은 다양한 검색 퀄리티에 적응하며, 리트리버를 변경하더라도 가상 토큰을 재학습시킬 필요없는 이점이 있다.
검색된 문서의 영향도
검색된 문서 중 top-m($m \in [1,5]$)개를 랜덤하게 선택해서 SPRING 학습 샘플을 만든다. 랜덤한 개수는 SPRING의 적응력을 높인다. 이런 학습 전략의 효과를 평가하기 위해, 0~5 문서에 대한 SPRING 성능을 비교한다.
SPRING은 더 많은 문서(m=1~4)를 사용할수록 QA 성능이 개선된다. 하지만 4개의 문서에서 최대 성능을 보이고 이를 넘어가면 성능이 떨어진다. 검색된 지식에 노이즈가 축적되어서로 추측된다. 그럼에도 불구하고 검색된 문구가 없는(m=0) 상황보다는 성능이 개선된다.
Cross-Dataset 일반화 능력
- 학습 데이터셋 :
- 8개의 QA 데이터셋(유사한 LLM 능력 요구)를 섞은 Mix 데이터셋
- NQ, TQA 각 데이터셋
Mix 데이터셋을 학습한 모델은 전반적으로 모든 데이터셋에 대한 성능이 높으므로, multi-task 학습의 장점이 증명된다. 단일 데이터셋(ex. NQ)에서 학습할때는 해당 테스트셋에서만 높은 성능을 보이고, 다른 데이터셋에서는 성능이 떨어진다. NQ만 학습하는 경우 MS MARCO 데이터셋에서 SPRING이 프롬프트 방식보다 성능이 낮다. 이는 다양한 데이터셋을 SPRING 학습에 사용하는 것이 일반화 능력에 유리함을 보인다.
Contribution & Limitation
- SPRING은 문서와 질문 사이에 들어갈 가상 토큰을 학습해서 RAG 작업에 높은 성능을 보인다.
- stage가 추가되었는데 Prompt, LoRA 방법과 비교할 수 있는 추론 속도 관련 자료가 없음
- 가상 토큰 탈부착이 가능하므로, 다른 파인튜닝 기법과 달리 LLM 능력을 유지할 수 있어서 non-RAG에서도 좋은 성능을 보인다.
- non-RAG에서 가상 토큰을 빼면 LoRA보다 유리한 결과가 나오는데, 이 실험 환경이 공정한지 의문
- RAG/non-RAG를 구분해서 가상 토큰 사용 여부를 시스템적으로 구현해서 비교하는 것이 더 공정할 것으로 보임