[논문리뷰]One token can help! learning scalable and pluggable virtual tokens for retrieval-augmented large language models (AAAI, 2025)

Date: 2026.02.13 Updated: 2026.02.13

카테고리: NR

Yutao Zhu, Zhaoheng Huang, Zhicheng Dou, and Ji-Rong Wen. 2025. One token can help! learning scalable and pluggable virtual tokens for retrieval-augmented large language models. In Proceedings of the Thirty-Ninth AAAI Conference on Artificial Intelligence and Thirty-Seventh Conference on Innovative Applications of Artificial Intelligence and Fifteenth Symposium on Educational Advances in Artificial Intelligence (AAAI’25/IAAI’25/EAAI’25). AAAI Press, Article 2915, 9 pages.

1. Problem Statement

이 논문이 풀고자 하는 핵심 task는 Retrieval-Augmented Generation(RAG), 그 중에서도 RAG기반 질의 응답 (RAG-based Question Answering, RAG-QA)에서 LLM이 검색 문서를 더 잘 활용하다록 만들면서도, None-RAG 상황에서의 일반 생성 능력을 훼손하지 않는 것이다. 기존에는 (i) Prompt Engineering을 통해 검색 결과를 단순 활용하거나, (ii) 모델을 RAG에 적응하도록 학습시키는 방식이 주류였다. (ii)의 경우 성능은 좋아도 LLM의 파라미터를 바꾸면서 일반 능력이 손상되는 문제가 발생해 사용(이미 배포된 LLM)에서 치명적이라고 주장한다.

SPRING은RAG 성능을 올리되, 배포된 LLM의 기존 기능(비-RAG 일반 능력)을 유지하면서, 필요할 때만 붙였다 떼는 plug-and-play 형태로 만드는 한 편, 실무 제약(최대 입력 길이) 때문에 검색 컨텍스트가 길어질수록 붙일 수 있는 토큰 수가 달라질 수 있어, 임의 개수의 가상 토큰을 추론 시점에 선택적으로 쓸 수 있는 scalable 메커니즘도 함께한다.

2. Limitations of Existing Works

프롬프트 기반 RAG의프롬프트 의존성] 프롬프트 기반 방법은 파라미터 튜닝 없이 어떤 LLM에도 적용 가능하지만, 효과가 사람의 프롬프트 설계 숙련도와 LLM의 instruction 이해 능력에 크게 좌우된다. 즉 같은 검색결과를 주어도어떻게 말하느냐에 성능이 흔들리는 구조적 한계가 있다.
[튜닝 기반 RAG의일반 능력 훼손] RAG 성능을 올리기 위해 LLM 파라미터(θ)를 직접 최적화하는 계열(사전학습/파인튜닝, LoRA 포함)은 보통 프롬프트 기반보다 성능이 높지만, 그 대가로 비-retrieval 시나리오에서의 일반 능력(commonsense reasoning, in-context learning 등)을 저하시킬 수 있음을 명시한다. 특히 이미 다목적으로 배포된 LLM에서는 파라미터 변경이 기존 기능에 영향을 주므로 적용 장벽이 크다.
[Retriever/Passage 조건 변화에 대한재학습 부담] RAG 파이프라인은 retriever가 교체되거나(top-k 정책 포함) 검색 패시지 수가 바뀌는 일이 비일비재하다. 이때 특정 retriever/특정 passage 수에 과적합된 방식은 시스템 업데이트마다 재학습이 필요해지며, 논문은 이를 실용성 관점의 핵심 병목으로 지적한다.

3. Methodology

3.1. Model Overview

SPRING은 retrieval로 검색된 문서 패세지 $R$과 질문 $Q$ 사이에 학습 가능한 가상 토큰 $T=[t_1, \dots, t_n]$을 삽입하여 RAG-QA의 입력을 $[R; t_1, \dots, t_n; Q]$ 형태로 재구성한다. 이 때, 백본 LLM의 파라미터 $\theta$는 freeze하고, 각 토큰의 임베딩 $\delta$만 학습한다. 학습 시에는 총 $n$개 중 임의의 $k(k \leq n)$를 뽑아 처음 $k$개 토큰 $t_{1:k}$만 쓰는 spring 최적화로 추론에서 토큰 개수를 가변적으로 써도 되게 만드는 scalable 구조를 가진다.

또한 오른쪽은 추론 단계에서 이 가상 토큰이 vocab에 없다는 문제를 해결하기 위해 $[r1],…,[r50][r1],\dots,[r50][r1],…,[r50]$와 같은 special token을 vocab에 추가하고 학습된 임베딩을 merge한 뒤, 실제 실행에서는 Retriever로 $R$을 얻고 $[R [r1]⋯[rk] Q][R\;[r1]\cdots[rk]\;Q][R[r1]⋯[rk]Q]$를 LLM에 넣어 생성하는 plug-and-play 가능한 모델이다.

3.2. Problem Formulation

RAG에서 corpus 집합을 $D$, retrieval $M$이라 할 때, 검색 결과를 $R = M_D(x_{<i})$로 정의했을 때, RAG-QA에서 토큰 생성 확률은 다음과 같다.

$$p_{\text{RAG-QA}} = \prod_{i=1}^m p_\theta(x_i \vert R;Q;x_{<i})$$

3.3. Scalable & Pluggable Virtual Tokens

SPRING은 학습 가능한 가상 토큰 $T=[t_1, \dots, t_n]$을 검색된 문서 패세지 $R$과 질문 $Q$ 사이에 삽입한다. 이때 생성 확률은 다음과 같이 정의된다.

\[p_{\text{SPRING}} = \prod_{i=1}^m p_\theta(x_i \vert R;[t_1, \dots, t_n ];Q;x_{<i})\]

LoRA 튜닝을 하는 것이 아닌, 가상 토큰들의 임베딩 $\delta \in \mathbb{R}^{n\times d}$만 학습에 사용한다. 이를 학습이 끝난 직후 추론시 모델에 삽입하여 vocabulary set에 학습된 가상 토큰들을 추가한다.

검색된 문서와 질문 사이에 문서를 삽입하는 이유는 두 가지이다. (i) auto-regressive에서 토큰이 $R$ 이후에 위치하며 $R$에 attention을 걸 수 있어 정보검색 이해를 돕고, (ii) LLM이 입력의 끝부분에 민감하므로, 매 샘플에서 일관되게 질문 직전 위치를 고정해 질문 이해에 대한 부작용을 줄이려는 목적이다.

3.4. Scalable Training: 임의 개수 토큰을 추론에서 쓰기 위함

SPRING은 학습 때부터 토큰 개수 변동을 노출시키는 방식으로 구현된다. 각 학습 샘플을 $\{ R,Q,V\}$에 대해 총 $n$개의 토큰이 있을 때, 무작위로 $k$개를 뽑고 이 가상 토큰들만 입력에 추가한다.

4. Experiments

4.1. Main Results

Table 1의 핵심 결론은 with retrieval에서 SPRING이 Prompt 대비 크게 향상되고, without retrieval에서 LoRA가 붕괴하는 반면 SPRING은 성능을 유지하거나 오히려 높인다는 점이다. with retrieval 평균 EM/F1은 Prompt 19.80/51.30, SPRING 28.40/60.15, LoRA 30.78/61.91, SPRING+ 32.17/62.96으로 정리되며, SPRING은 소수 파라미터로도 큰 개선을 만들고 SPRING+는 최고 성능을 보인다. without retrieval 평균에서는 Prompt 12.59/46.19 대비 LoRA가 EM 0.02로 급락하는 반면, SPRING은 16.32/52.24, SPRING+는 17.96/53.74로 유지·개선된다.

4.2. Impact of Number of Virtual Tokens

Figure 3은 $k$를 늘릴수록 평균 EM/F1이 전반적으로 증가하고, $k=1$만으로도 RAG 성능이 유의미하게 개선된다는 점을 보여준다. 반대로 고정 $k=50$만으로 학습한 변형(Fix 50)은 추론에서도 동일한 $k$에 묶여 가변 $k$를 지원하는 목적(scalable)과 맞지 않다는 비교가 함께 제시된다.

5. Conclusion

Contribution

[파라미터 보존형 RAG 적응(Plug-and-Play) 제안]: LLM 파라미터 $\theta$를 고정한 채, 검색된 문서와 질문 사이에 학습 가능한 virtual tokens만 삽입해 RAG 성능을 올리고 non-RAG에서는 토큰을 제거해 일반 능력을 유지하는 설계를 제안
[실사용 변화에 대한 일반화(robustness) 실증]: retriever가 바뀌거나(passages 수 포함) 환경이 변해도 성능이 안정적으로 유지됨을 분석 실험으로 보여, retriever 업데이트가 잦은 RAG 파이프라인에서 재학습 부담을 줄일 수 있음을 실증

Limitations

[최대 입력 길이 제약의 잔존]: scalable 설계의 동기 자체가 LLM의 최대 입력 길이 제약에 기반하므로, retrieval 컨텍스트가 길어질수록 사용할 수 있는 토큰/컨텍스트가 제한되는 시스템 제약은 그대로 남음

Meaningful