NL-198, Query Rewriting in Retrieval-Augmented Large Language Models, EMNLP 2023

0 Abstract 

  • LLM(대형 언어 모델)은 retrieve-then-read 파이프라인에서 강력한 블랙박스 리더 역할을 하여 knowledge-intensive 작업에서 놀라운 발전을 이루었습니다. 
  • 이 작업에서는 쿼리 재작성의 관점에서 retrieval-augmented LLM에 대한 이전 retrieve-then-read대신 새로운 프레임워크인 Rewrite-Retrieve-Read를 도입합니다.
  • 이전 연구는 retriever or the reader 중 하나를 적응시키는 데 중점을 두었다면, 우리의 접근은 검색 쿼리 자체의 적응에 주목합니다. 
  • 왜냐하면 검색에서 입력 텍스트와 필요한 지식 사이에는 불가피하게 격차가 있기 때문입니다. 
  • 먼저 LLM에게 쿼리를 생성하도록 유도한 다음 웹 검색 엔진을 사용하여 맥락을 검색합니다. 
  • 더불어 쿼리를 동결된 모듈에 더 잘 맞추기 위해 우리의 파이프라인에 대한 학습 가능한 방법을 제안합니다. 
  • 작은 언어 모델을 학습 가능한 rewriter로 채택하여 블랙박스 LLM 판독자에 맞춥니다. 
  • rewriter는 강화 학습을 통해 LLM 판독자의 피드백을 사용하여 훈련됩니다. 
  • 평가는 하향 작업, 오픈 도메인 QA 및 다중 선택 QA에서 수행됩니다. 
  • 실험 결과는 일관된 성능 향상을 보여주며, 우리의 프레임워크가 효과적이고 확장 가능하며, 검색을 강화한 LLM에 대한 새로운 프레임워크를 제공한다는 것을 나타냅니다.

1 Introduction

  • 대규모 언어 모델 (LLM)은 인간 언어 처리에 대한 놀라운 능력과 몇 개 또는 제로샷 설정에서의 탁월한 확장성 및 적응성을 보여주었습니다. (Ouyang et al., 2022; Brown et al., 2020; Chowdhery et al., 2022). 
  • 그러나 훈련 과정은 대규모 고품질 말뭉치에 의존하지만 실제 세계의 인식은 없습니다. 
  • 따라서 LLM은 여전히 hallucination 및 temporal misalignment (Röttger and Pierrehumbert, 2021; Luu et al., 2022; Jang et al., 2022) 문제에 직면해 있습니다. 
  • 이는 LLM의 신뢰성에 영향을 미치며 LLM 응답과 실제 세계 간의 일관성이 더 많은 확인이 필요하기 때문에 더 넓은 실용적 응용을 방해합니다. 
  • 기존 연구는 외부 지식 (즉, 비매개변수 지식)을 내부 지식 (즉, 매개변수 지식)과 효과적으로 결합하면 특히 지식 집약적 작업에서 환상을 효과적으로 완화할 수 있다는 것을 입증했습니다. 
  • 사실, retrieval-augmented LLM은 환상 문제를 완화하는 데 효과적이어서, naive LLM 생성에서의 사실적인 단점을 완화하기 위한 표준 솔루션으로 간주되고 있습니다. 
  • Retrieval augmentation 은 언어 모델의 외부 맥락으로 상대적인 passages를 선택하는 것으로, 이는 retrieve-then-read 프레임워크입니다. 
  • 예를 들어, 개방형 도메인 질문 응답 작업 (개방형 도메인 QA)을 살펴보겠습니다. 
  • 검색기는 먼저 질문에 대한 관련 문서를 검색합니다. 
    • 그런 다음 LLM은 질문과 문서를 받아서 답을 예측합니다.
    • 이 논문에선 이런 방법을 retrieve-then-read 라고 부름
  • 대부분의 LLM은 추론 API를 통해서만 접근할 수 있기 때문에 파이프라인에서는 블랙박스 형식의 frozen readers 역할을 합니다. 
  • 이로 인해 완전한 액세스를 필요로 하는 이전의 검색 보강 방법이 더 이상 실행 가능하지 않습니다. 
  • 최근의 검색 보강 언어 모델에 대한 연구는 주로 LLM 중심의 적응에 중점을 두고 있습니다. 
  • 하나의 아이디어는 밀집 검색 모델을 훈련하여 frozen language model에 맞추는 것입니다 (Shi et al., 2023). 
    • LLM의 피드백을 훈련 목표로 사용하여 검색 모델은 더 나은 LLM 입력 컨텍스트를 위해 조정됩니다. 
  • 다른 연구 방향은 retriever and the reader 간의 상호 작용 설계에 중점을 두고 있습니다, 여기서 retriever and the reader는 일반적으로 frozen 상태입니다. 
    • 아이디어는 정교한 프롬프트나 복잡한 프롬프트 파이프라인을 통해 신생 능력을 유도하는 것입니다. 
    • 외부 지식과의 다양한 상호 작용을 통해 LLM은 답에 점진적으로 접근할 수 있습니다.
  • 그러나 아직 해결해야 할 문제가 남아 있습니다. 
    • 기존 접근 방식은 쿼리 즉, retrieve-then-read 파이프라인의 입력에 대한 적응을 간과합니다. 
    • 검색 쿼리는 데이터셋에서 가져온 것이거나 블랙박스 생성에 의해 직접 결정되어 항상 고정되어 있습니다. 
    • 그러나 입력 텍스트와 실제로 쿼리해야 하는 지식 사이에는 불가피하게 차이가 있습니다. 
    • 이는 성능을 제한하고 retrieve 능력 향상과 프롬프트 엔지니어링에 부담을 주는 요인이 됩니다.
  • 이 문제를 고려하여 이 논문은 검색 보강을 위한 새로운 프레임워크인 Rewrite-Retrieve-Read를 제안합니다. 
    • 이 프레임워크는 LLM에 적응하기 위해 더 조정될 수 있습니다. 
    • retriever 앞에는 입력을 다시 작성하는 단계가 추가되어 주어진 입력과 검색 필요 사이의 간격을 메우는데, 이는 그림 1에 나와 있습니다. 
  • 우리는 인터넷 검색 엔진과 같은 사용 가능한 도구를 retriever로 채택하고 있습니다. 
  • 이는 검색 인덱스의 유지보수를 피하고 최신 지식에 액세스할 수 있습니다 (Lazaridou et al., 2022). 
  • 이전 연구와는 달리 retriever와 LLM 간의 각 샘플에 대한 다중 상호 작용 라운드를 기억해야 하는 방법이 아니라, 우리의 rewriting step의 동기는 입력 텍스트에서 검색 요구를 명확히 하는 것입니다.
    • 즉 query rewrite해서 검색하고, 검색문서와 함께 LLM 넣어서 출력을 한다.
    • open QA 테스크에서 실험을 하는 듯?
    • 그리고 query rewrite하는 것은 conversational search처럼 앞의 context을 보강하기 위함이라기보단, 검색 요구를 명확히 하기 위한 작업이라고 함
    • 아마 뒤에서 설명하겠지만 그림1(c) 보면, 출력의 reward을 계산해서 rewrite 모델을 업데이트 하는 듯 함
  • 우리는 rewrite-retrieve-read 프레임워크에 대한 교육 가능한 체계를 제안합니다 (도표 1 (c)).
    • black-box retriever and the reader는 고정된 시스템을 형성합니다. 
    • 우리의 파이프라인의 단계를 더욱 부드럽게 만들기 위해 우리는 작은 교육 가능한 언어 모델을 적용하여 재작성 단계를 수행합니다. 이를 rewriter로 표시합니다. 
    • rewriter는 LLM 성능을 보상으로 사용하여 강화 학습으로 교육되며, 검색 질문을 조정하여 reader를 개선하여 downstream 작업에서의 성능을 향상시키도록 학습합니다.
  • 우리가 제안한 방법들은 지식 중심 하향 작업, 즉 오픈 도메인 QA (HotpoQA, AmbigNQ, PopQA) 및 다중 선택형 QA (MMLU (Hendrycks et al., 2021))을 포함한 작업에서 평가되었습니다. 
    • 실험은 rewriter로 T5-large, reader로 ChatGPT 및 Vicuna-13B를 사용하여 구현되었습니다. 
    • 결과는 검색 확장된 LLM 성능이 일관되게 향상된다는 것을 보여줍니다. 
    • 또한 결과는 더 작은 언어 모델이 검색 질문 재작성에 유용할 수 있다는 것을 나타냅니다.
  • 요약하면, 우리가 제안한 혁신적인 검색 확장 방법인 rewrite-retrieve-read는 입력 텍스트가 고정된 검색기와 LLM reader에 맞춰진 최초의 프레임워크입니다. 
    • 우리는 작은 교육 가능한 모델을 사용한 조절 가능한 체계를 소개하여 리소스 소모를 줄이면서 성능 향상을 달성합니다.
    • 다른 논문들도 query을 rewrite하는데 LLM쓰고 했는데, 이 논문은 이 부분을 small model로 대체한게 좀 다르다.
    • 대체할 때는 reader 부분에 해당하는 LLM의 출력의 reward을 통해서 어찌어찌 한거 같은게 노벨티가 있는 듯?

2 Related Work

2.1 Retrieval Augmentation

  • 언어 모델에는 사실성 문제를 완화하기 위해 외부 지식이 필요합니다. 
  • 검색 확대는 표준적이고 효과적인 솔루션으로 간주되었습니다. 
  • 검색 모듈을 사용하면 관련 구절이 원래 입력의 컨텍스트로 언어 모델에 제공됩니다. 
  • 따라서 상식이나 실시간 뉴스와 같은 사실 정보는 상황에 맞는 독해를 통해 출력 예측에 도움이 됩니다.
  • 이전 연구에서는 사전 훈련된 언어 모델(PrLM) 앞에 희소 검색기 또는 밀집 검색기를 사용했습니다. 
  • 신경 검색기와 판독기는 모두 BERT 또는 BART와 같이 훈련 가능한 크기의 PrLM입니다. 
  • 따라서 전체 검색 후 판독기 프레임워크는 조정 가능한 엔드투엔드 시스템이며, 여기서 검색된 컨텍스트는 중간 결과로 간주될 수 있습니다. 
  • 검색 및 독해력을 최적화하기 위해 2단계 프레임워크를 원활하게 하는 접근 방식이 제안되었습니다. 
  • 최근에는 모델과 데이터의 크기가 빠르게 확장됨에 따라 검색이 여전히 강력한 향상 기능으로 남아 있습니다. 
  • 반면, 검색 향상은 대규모 언어 모델에 비해 매개변수 크기의 부족을 보완할 수 있습니다. 
  • 예를 들어, Atlas는 리트리버와 리더를 공동으로 훈련함으로써 540B PalM과 동등한 소수 샷 성능을 보여 주지만 크기는 50배 더 작습니다.
  • The Internet as a knowledge base 
    • 우리 작업과 관련하여 검색 엔진은 검색자의 역할을 맡을 수 있으며 인터넷을 외부 지식의 소스로 사용할 수 있습니다. 
    • Komeiliet al. (2022)은 대화 기록을 기반으로 관련 정보를 인터넷 검색하여 대화 응답 생성을 수행합니다. 
    • SeeKeR은 단일 Transformer를 사용하여 검색 쿼리 생성을 반복적으로 수행한 다음 대화 생성 및 문장 완성을 위한 지식 추출을 수행합니다. 
    • 대규모 모델의 경우 웹 검색은 지식 증대, 사실 확인 및 LLM 에이전트 향상에 여전히 효과적인 것으로 나타났습니다.

2.2 Cooperation with Black-box LLMs

  • ChatGPT(Ouyang et al., 2022), Codex(Chen et al., 2021), PaLM(Chowdhery et al., 2022)과 같은 대규모 언어 모델은 인상적인 자연어 처리 능력과 놀라운 확장성을 나타냅니다. 
  • 이는 광범위한 NLP 작업에서 LLM을 수용하는 경향으로 이어집니다. 
  • 그러나 대부분의 경우 LLM은 블랙박스로만 액세스할 수 있습니다. 
  • 이는 (i) ChatGPT와 같은 일부는 오픈 소스가 아니고 비공개로 유지되기 때문입니다. 
  • (ii) 대규모 매개변수 규모에는 사용자가 항상 감당할 수 없는 계산 리소스가 필요합니다. 
    • 이 제약 조건은 입력 및 출력 텍스트 외에는 아무것도 사용할 수 없음을 의미합니다.
  • 기존 연구에서는 신중하게 설계된 상호 작용 방법을 통해 LLM의 능력을 더 잘 활용할 수 있음이 입증되었습니다. 
  • GenRead는 LLM이 검색기를 배포하는 대신 컨텍스트를 생성하도록 요청하여 LLM이 메시지를 표시하여 내부 지식을 검색할 수 있음을 보여줍니다. 
  • ReAct(Yao et al., 2023)와 Self-Ask(Press et al., 2022)는 CoT(Chain-of-Thought)와 상호 작용을 결합합니다. 
  • 웹 API로. ReAct는 신속한 구성에만 의존하여 대화형 작업을 위한 새로운 기준을 제공합니다. 
  • DSP(Demonstrate-Search-Predict)는 LLM과 검색기 간의 정교한 파이프라인을 정의합니다. 
  • ReAct와 달리 DSP는 멀티홉 분석 및 검색 외에도 데모 부트스트랩에 대한 프롬프트를 통합합니다.
  • 제로 샷 또는 소수 샷 설정에서 유망한 성능에도 불구하고 LLM의 동작은 때때로 조정이 필요합니다. 
    • 실행 가능한 접근 방식은 LLM 앞이나 뒤에 훈련 가능한 작은 모델을 추가하는 것입니다. 
  • 시스템 매개변수의 일부인 소형 모델은 최적화를 위해 미세 조정될 수 있습니다. 
  • RePlug는 검색 후 읽기 파이프라인에서 고정된 LLM에 대한 밀집 검색기를 미세 조정하기 위해 제안되었습니다. 
    • 검색자는 LLM의 감독하에 LLM에 적합한 문서를 검색하도록 훈련되었습니다. 
  • 동일한 목적으로 방향성 자극 프롬프트는 작은 모델을 배포하여 LLM 보상에 따라 업데이트되는 자극(예: 요약을 위한 키워드 또는 응답 생성을 위한 대화 작업)을 LLM에 제공합니다.
  • 위에서 언급한 영감을 주는 작업과 달리 제안된 파이프라인에는 검색 후 읽기 모듈 앞에 쿼리 재작성 단계가 포함되어 있습니다. 
  • 우리는 검색 쿼리를 재구성하여 검색 증강 LLM을 위한 새로운 개선 사항인 작은 재작성 모델을 사용하여 훈련 가능한 방식을 제안합니다.

3 Methodology

  • 쿼리 재작성 관점에서 검색 강화 LLM을 개선하는 파이프라인인 Rewrite-Retrieve-Read를 제시합니다. 
  • 그림 1은 개요를 보여줍니다. 
  • 이 섹션에서는 먼저 섹션 3.1의 파이프라인 프레임워크를 소개한 다음 섹션 3.2의 훈련 가능한 체계를 소개합니다.

3.1 Rewrite-Retrieve-Read

  • 검색 증대 작업은 다음과 같이 표시될 수 있습니다. 
  • 지식 집약적인 작업(예: 오픈 도메인 QA)의 데이터 세트가 주어지면 D = {(x, y)i}, i = 0, 1, 2, . . . , N, 
    • x(예: 질문)는 파이프라인에 대한 입력이고, y는 예상되는 출력(예: 정답)입니다. 
  • 우리의 파이프라인은 세 단계로 구성됩니다. 
    • (i) Query rewrite: 원래 입력 x를 기반으로 필요한 지식에 대한 쿼리 x~를 생성합니다. 
    • (ii) Retrieve: 관련 컨텍스트, 문서를 검색합니다. 
    • (iii) Read: 컨텍스트 [doc, x]와 함께 입력을 이해하고 출력 yˆ를 예측합니다
    • 간단하면서도 효과적인 방법은 LLM에게 잠재적으로 필요한 정보를 검색하기 위한 쿼리를 다시 작성하도록 요청하는 것입니다. 
  • 우리는 LLM이 생각하도록 장려하기 위해 몇 번의 프롬프트를 사용하며 출력은 검색할 쿼리가 하나 이상일 수 있습니다.

3.2 Trainable Scheme

  • 게다가, 동결된 LLM에 전적으로 의존하는 것은 몇 가지 단점을 보여주었습니다. 
  • 추론 오류나 유효하지 않은 검색은 성능을 방해합니다.
  • 반면에 검색된 지식은 때때로 언어 모델을 오도하고 손상시킬 수 있습니다. 
  • 고정된 모듈에 더 잘 맞추려면 학습 가능한 모델을 추가하고 LLM reader 피드백을 보상으로 받아 이를 조정하는 것이 가능합니다. 
  • 우리 프레임워크를 기반으로 그림 1의 오른쪽 부분에 표시된 것처럼 훈련 가능한 소형 언어 모델을 활용하여 재작성 단계를 대신할 것을 제안합니다. 
    • 즉 LLM에게 단순히 query rewrite을 맡겨버리면, 추론 오류 혹은 유효하지 않은 검색을 야기시키는 경우가 있다?
    • 아마 ablation study 같은데서 비교했을 듯
    • 암튼 그래서, 작은모델로 query rewrite을 대체하는 대신, 리더결과(QA의 LLM결과)에서 reward을 계산해서 학습하는 방식을 제안한다.
  • 훈련 가능한 모델은 사전 훈련된 T5-large(770M)로 초기화됩니다. 
    • 훈련 가능한 재작성자 Gθ로 표시됩니다. 
    • 리라이터는 먼저 pseudo 데이터에 대해 훈련을 받아 워밍업(§3.2.1)한 다음 강화 학습(§3.2.2)을 통해 지속적으로 훈련됩니다.

3.2.1 Rewriter Warm-up

  • The task, query rewriting, is quite different from the pre-training objective of sequence-to-sequence generative models like T5. First, we construct a pseudo dataset for the query rewriting task. Inspired by recent distillation methods (Hsieh et al., 2023; Ho et al., 2022), we prompt the LLM to rewrite the original questions x in the training set and collect the generated queries x˜ as pseudo labels. The collected samples are then filtered: Those that get correct predictions from the LLM reader are selected into the warm-up dataset, denoted as DT rain = {(x, x˜)|yˆ = y}. The rewriter Gθ is finetuned on DT rain with the standard log-likelihood as the training objective, denoted as
  • d
  • The rewriter model after warm-up shows modest performance, which depends on the pseudo data quality and rewriter capability. Highly relying on the human-written prompt line, x˜ can be suboptimal. The relatively small scale of the rewriter size is also a limitation of the performance after the warm-up. Then we turn to reinforcement learning to align the rewriter to the following retriever and LLM reader

3.2.2 Reinforcement Learning

  • To further fine-tune the rewriter to cater to the LLM reader, we adopt a policy gradient reinforcement learning framework.
  • Task Formulation 
    • In the context of reinforcement learning, the rewriter optimization is formulated as a Markov Decision Process 5-tuple
    • ⟨S, A, P, R, γ⟩. (i) The state space S is a finite set
    • limited by the vocabulary and the sequence length.
    • (ii) The action space A is equals to the vocabulary.
    • (iii) The transition probability P is determined by
    • the policy network, which is the rewriter model
    • Gθ. (iv) The reward function R gives a reward
    • value that depends on the current state. The policy gradient is derived from rewards, used as the
    • training objective. (v) γ denotes the discount factor. More specifically, the rewriter Gθ after the
    • warm-up is the initial policy model π0. At each
    • step t, the action at
    • is to generate the next token
    • xˆ˜t based on the observation of the present state,
    • st = [x, xˆ˜<t]. When the generation is stopped by
    • the End-Of-Sentence token, one episode is ended.
    • After finishing the retrieval and reading, a reward
    • is computed by evaluating the final output, i.e., a
    • score for the LLM reader prediction.
  • Policy Optimization 
    • We adopt Proximal Policy Optimization (PPO) (Schulman et al., 2017), following (Ramamurthy et al., 2022). Maximization of the expectation of the reward R is formulated as
    • where θ ′ is the temporarily fixed policy for sampling and θ is updated. A denotes the advantage function, which is formulated based on the estimation of value network Vϕ. The value network Vϕ is initialized from the policy network π0. The formulation follows Generalized Advantage Estimation (GAE) (Schulman et al., 2015).
    • where λ is the bias-variance trade-off parameter.
    • The reward function R reflects the quality of the generated queries, which needs to be consistent with the final evaluation of the task. xˆ˜ is fed to the retriever and the reader for a final prediction yˆ. A part of the reward function is the measures of yˆ compared to the golden label y (e.g., exact match and F1 of the predicted answers), denoted as Rlm. Besides, a KL-divergence regularization is added to prevent the model from deviating too far from the initialization (Ramamurthy et al., 2022; Ziegler et al., 2019).
    • The final loss function is composed of policy loss and value loss.
    • Here, S denotes the sampled set, and T is for step numbers.





































Reference

댓글