[ICLR 2025 Under Review] Optimizing Inference-Time Reasoning in LLMs via Retrieval-Augmented Reflection

최근 inference time을 통해 모델이 reasoning을 거쳐 답변을 내놓게 하는 방식이 많이 사용되고 있다고 생각했다.

이에 따라, RAG에도 그런 방법론을 적용한 논문이 있다가 찾아보다가 리뷰

 

1. Introduction

• LLM의 등장과 CoT (chain-of-thought)와 같은 프롬프트 방법론의 발전으로 인해 많은 자연어 태스크의 발전을 이루었음
• 하지만 모델이 생성하는 텍스트가 항상 사실을 기반으로 생성된다는 보장을 하긴 어려웠음 (ex, hallucination)
• RAG (Retrieval Augmented Generation) 방법론은 검색 결과를 통해 LLM이 사실 관계에 기반하여 생성할 수 있도록 유도
• 본 논문에서는 CoT를 사용할 시에 reasoning step에서 발생할 수 있는 hallucination을 RAG를 통해 해결하는 것을 목표

 

2. Method

2.1 Retrieval Augmented Reflection (RaR)

• 가장 먼저 사용자 질의 x를 LLM에 입력하여 초기 답변인 y_raw를 생성

• y_raw의 생성 결과를 발전시키기 위해서 생성된 x와 y_raw를 다시 LLM에 입력하여 검색에 활용할 쿼리 q 생성

• 생성한 쿼리 q를 이용하여 검색기를 통해 연관된 k개의 document 검색 (검색은 벡터의 코사인 유사도 활용)

• 사용자 질의 x, 초기 답변 y_raw, 관련 문서 V_k를 LLM에 입력하여 개선된 답변 y_reflection 생성

• RAG는 사용자 질의 x와 검색된 문서 V_k를 통해서만 생성하지만, RaR는 모델이 초기에 생성한 답변인 y_raw를 이용한다는 점에서 차이가 있다. 
• RaR는 같은 방법을 여러 번 반복하여 수행한다면 더 성능을 끌어올릴 수 있다. (y_reflection을 다시 y_raw로 취급하고 시작)

2.2 Iterative Retrieval Augmented Reflection (Iterative RaR)

2.2.1 Reflection on Intermediate Reasoning Steps

• 이 과정에서는 각 reasoning step에 대한 개선 과정이 추가
• 질의 x를 통해 j개의 reasoing step y_thought(y_thought_1, ..., y_thought_j) 생성
• 생성한 y_thought_i 와 질의 x를 이용하여 검색에 활용할 쿼리 q_i 생성
• 생성한 각 쿼리를 통해 관련 문서를 검색하고, 질의 x, y_thought_i, q_i를 이용하여 개선된 답변 y_RaR_i 생성

• 이 단계를 통해서 각 reasoning step에서 발생할 수 있는 오류를 제거할 수 있다

2.2.2 Refinement of Overall Response

• 앞 단계에서 생성한 y_RaR와 질의 x를 통해 다시 쿼리 생성
• 다시 관련 문서 V_k를 검색
• 질의 x, y_RaR, 관련 문서 V_k를 통해 더욱 정제된 답변 y_RaR' 생성
• 이 과정은 y_RaR'를 다시 y_RaR로 사용하여 반복할 수 있다. 

Iterative RaR 알고리즘

2.3 Optimizing RaR by Scaling Inference-time Computation

• 모델을 scaling하는 방법은 크게 2가지 방법이 있다.
• 단순 모델의 파라미터 수를 늘리는 방법
• 모델의 input(RAG의 경우, 입력되는 문서의 수 조절) / output token (생성할 때 더 많은 reasoning step을 거치게 하는 것)을 늘리는 방법
• Iterative RaR은 모델의 input/output 토큰을 늘리는 방법이라고 볼 수 있음

 

3. Experiments

3.1 Benchmarks

• Code Generation: ClassEval, HumanEval, HumanEval+, MBPP, MBPP+
• Mathematical Reasoning: GSM8K, GSM-HARD
• Task Planning: 마인크래프트 환경을 통한 평가 (?)
• Question Answering: TriviaQA

3.2 Main Results

• 전체적으로 좋은 성능을 보였으며, 모델의 크기와 상관 없이 RaR를 적용했을 때 언어 모델의 성능을 더 끌어올릴 수 있었다
• Self-RAG와 Active RAG 또한 RAG과정과 reasoning을 결합하긴 하지만, RaR에는 미치지 못했다. 

• (a)는 모델의 순수 파라미터에 따른 성능, (b)는 inference token 수에 따른 성능, (c)는 inference cost에 따른 성능
• RaR를 적용한 모델은 다른 큰 모델과 비교해도 좋은 성능을 보였다. (a)
• RaR의 경우 inference token의 수가 적을 때는 낮은 성능을 보이는데, 이는 RaR이 반복적으로 생성하면서 더 나은 결과를 생성하는 것에 따른 결과로 보인다. (b)
• RaR을 사용하는 것은 또한 경제적인 선택 (c) 

• 단순히 CoT와 RAG를 섞는 것보다 RaR 과정을 통해 시너지를 내는 것이 더 효과적이었다. 
• RaR는 각 단계별로 검색을 수행하기 때문에 더 구체적인 상황에 대한 검색이 가능해서 이런 결과가 나온 것

 

4. Limitation & Conclusion

• 해당 모델은 학습 없이 이루어지는 프레임워크이기 때문에 베이스 모델의 CoT 성능, in-context 성능에 영향을 받는다. 
• 작은 모델에 대한 해당 프레임워크의 효과는 의문이 든다. 
• 하지만, 해당 프레임워크는 zero-shot CoT를 RAG를 통해 보완하여 더 좋은 thoughts를 만들고, 나아가 더 좋은 답변을 생성할 수 있었다.