[논문 리뷰] RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval

ICLR 2024 [paper]
Sarthi et al.
31 Jan 2024

 

Introduction

Background

예시(NarrativeQA dataset)
"신데렐라는 어떻게 행복한 결말을 맞이했는가?"

1. 기존의 retrieval-augmented 방법들은 연속적인 짧은 청크만 검색하기 때문에 전체 문서 컨텍스트에 대한 이해가 제한된다. 
   즉, 텍스트의 여러 부분에서 지식을 통합해야 하는 주제와 관련된 질문(위 예시)은 top-k개의 짧고 연속적인 텍스트를 이용하여 질문에 답하기에는 충분한 컨텍스트가 포함되지 않는다.
2. 기술의 발전에 따라 모델이 처리할 수 있는 컨텍스트 길이가 확장되면서, "모델에 많은 양의 컨텍스트를 입력해주면 되지 않을까"라는 의문이 제기되었다.
   그러나 Liu et al. (2023) 과 Sun et al. (2021) 에 따르면, 컨텍스트의 길이가 증가함에 따라서 모델이 긴 컨텍스트에 포함된 관련된 정보를 충분히 잘 활용하지 못하면서 성능이 저하되는 경향이 있다.
3. 실용적인 측면에서, 긴 컨텍스트를 사용하는 것은 비용과 시간이 많이 소모된다.

 

RAPTOR overview

• 재귀적으로 텍스트 청크를 임베딩, 클러스터링, 요약하는 방식으로, 하단에서부터 상단으로 요약 레벨이 다른 트리 구축
• Inference 시에는 구축된 트리에서 검색함으로써, 긴 문서에 대해 다양한 요약 레벨에서 정보를 얻어 통합
• RAPTOR 검색을 GPT-4와 함께 사용했을 때 QuALITY 벤치마크(중간정도 길이의 passages를 기반으로 한 객관식 QA)에서 절대 정확도를 20%까지 향상

트리 구축과정

1. 텍스트 청킹 및 임베딩
2. 클러스터링
3. 요약
4. 반복하며 하단에서 상단으로 트리 생성 (추가 클러스터링이 불가능해질 때까지)

 

RAPTOR

텍스트 청킹 및 임베딩

1. retrieval corpus를 100-토큰 길이로서 짧고 연속적인 텍스트로 분할
2. 문장이 100-토큰을 초과하면 중간 문장을 자르는 대신, 전체 문장을 다음 청크로 이동
3. 이 텍스트를 SBERT를 사용하여 임베딩 (첫 SBERT 임베딩은 트리의 Leaf layer)

 

문서 클러스터링

각 텍스트 세그먼트는 다양한 정보를 담고있을 수 있기 때문에, 여러 클러스터에 속할 수 있도록 클러스터의 수를 고정하지 않는 soft clustering을 사용한다. 클러스터링 알고리즘은 Gaussian Mixture Models(GMMs)을 사용하며, GMMs은 여러 개의 가우시안 분포가 있을 때 각 텍스트 세그먼트가 어떤 분포에 속하는 지에 대한 확률을 계산한 다음, 이를 종합하여 전체 분포를 나타낸다. 이를 통해 각 텍스트 벡터가 어떻게 분포되어 있는지 쉽게 이해할 수 있다.

 

GMMs 통계 모델로 텍스트 벡터의 확률분포를 나타내는 방법은 아래와 같다.

1. 텍스트 벡터의 확률 $P(\mathbf{x} \mid k)$:
   • 특정 텍스트 벡터 x가 k번째 가우시안 분포에 속할 확률
   • $P(\mathbf{x} \mid k) = \mathcal{N}(\mathbf{x}; \mu_k, \Sigma_k)$
     
     • $\mu_k$: k번째 가우시안 평균 벡터
     • $\Sigma_k$: k번째 가우시안 분포의 공분산 행렬
2. 전체 확률분포 $P(\mathbf{x})$:
   • 전체 확률분포 $P(\mathbf{x})$는 여러 가우시안 분포의 가중치합으로 나타낸다.
   • $\sum_{k=1}^{K} \pi_k \mathcal{N}(\mathbf{x}; \mu_k, \Sigma_k)$
     • $\pi_k$: k번째 가우시안 분포의 가중치 (합은 1)
     • $K$: 가우시안 분포의 개수

텍스트 벡터의 확률분포를 정의했다면, 그 다음으로 유사한 벡터끼리 클러스터링해야 한다. 그러나 고차원의 벡터 공간에서 유사성을 측정하는 것은 거리 계산 메트릭이 제대로 작동하지 않을 수 있다. 따라서 차원을 축소하기 위해 Uniform Manifold Approximation and Projection(UMAP)을 사용한다. 기존의 UMAP의 n_neighbors는 로컬 구조와 글로벌 구조의 보존에 대한 밸런스를 조절한다. RAPTOR에서는 계층적인 클러스터링 구조를 만들기 위해서 n_neighbors를 다소 변형하였다. 1) 먼저 글로벌 클러스터를 식별한 다음, 2) 글로벌 클러스터 내에서 로컬 클러스터링을 수행한다. 이 2단계 프로세스는 광범위한 주제에서 세부적인 디테일까지 텍스트 간의 관계를 포착할 수 있다. 또한, 로컬 클러스터의 결합된 컨텍스트가 요약 모델의 토큰 제한을 초과한 경우, 클러스터 내에서 다시 클러스터링을 적용한다.

 

최적의 클러스터 수를 결정하기 위해 Bayesian Information Criterion (BIC)를 사용한다. BIC는 모델 복잡도에 대해서 페널티나 보상을 부여한다.

1. GMM에서의 BIC
   • $\text{BIC} = \ln(N)k - 2\ln(\hat{L})$
     • $N$: 텍스트 세그먼트의 수
     • $k$: 모델 파라미터 수 (입력 벡터의 차원과 클러스터 수에 대한 함수)
     • $\hat{L}$: likelihood function의 최댓값

BIC로 구한 최적의 클러스터 수를 가지고, Expectation-Maximization 알고리즘은 GMM 파라미터, 평균, 공분산, mixture weights를 추정한다.

 

요약

각 클러스터의 노드들은 gpt-3.5-turbo에 입력되어 요약된다. 요약을 위해 사용된 프롬프트는 아래와 같다.

system: You are a Summarizing Text Portal
user: Write a summary of the following, including as many key details as possible: {context}:

 

 

Querying

RAPTOR 트리에서 쿼리와 관련된 정보를 검색하기 위해서 Tree traversal, Collapsed tree 두 가지 방법을 제안한다.

Tree traversal (트리 순회)

1. RAPTOR 트리의 root layer에서 시작한다. query embedding과 초기 레이어의 노드들의 embedding 사이의 코사인 유사도를 계산한다.
2. 코사인 유사도 점수를 기반으로 top-k개를 고른다. ($S_1$)
3. S1 set 요소들의 자식 노드들로 넘어간다. query embedding과 자식 노드들의 embedding 사이의 코사인 유사도를 계산한다.
4. 코사인 유사도 점수를 기반으로 top-k개를 고른다. ($S_2$)
5. 위 스텝을 leaf nodes에 도달할 때까지 반복한다. ($S_1, S_2, \ldots, S_d$)
6. 선택된 모든 노드들의 텍스트는 검색 컨텍스트의 형태로 합친다.

 

Collapsed tree (축소 트리)

1. RAPTOR 트리를 하나의 레이어로 flatten한다. flatten된 노드들 집합은 C로 나타내며, 원래 트리의 모든 레이어 노드들을 포함한다.
2. query embedding과 축소된 C의 모드들의 embedding 사이의 코사인 유사도를 계산한다.
3. 코사인 유사도 점수를 기반으로 top-k개를 고른다. 모델의 입력 제한을 초과하지 않도록 미리 정의한 최대 토큰 수에 도달할 때까지 결과 노드를 계속 추가한다.

 

비교

두 방법을 QASPER 데이터셋에서 실험한 결과, Collapsed tree 방법이 일관되게 더 나은 성능을 보였다. 즉, 모든 노드들을 동시에 검색하며 주어진 질문에 대해 세분화되고 정확한 정보들을 검색한다.

(※ 개인적으로도 Tree traversal 방법은 초기에 선택된 노드들의 자식 노드들에 대해서만 검색을 하다보니 다양한 검색 기회를 잃어버릴 것 같은 생각..)

Collapsed tree 방법의 단점은 모든 노드들에 대해서 코사인 유사도 계산을 수행해야 된다는 점인데, 이마저도 FAISS 같은 효율적이고 빠른 KNN 라이브러리를 사용해서 해결할 수 있다.

따라서 Collapsed tree 접근 방법을 채택하였고, 최대 토큰은 2000개로 설정했는데 이는 대략적으로 top-20 노드들을 검색하는 것과 같다.

 

Experiments

Datasets

NarrativeQA

• 책, 영화 녹취록 기반 QA 데이터셋, 총 1,572개
• 전체 이야기를 포괄적으로 이해해야 하며, 긴 텍스트를 이해하는 능력 테스트
• 표준 BLEU (B-1, B-4), ROUGE (R-L), and METEOR (M) metrics로 성능 측정

QASPER

• 1,585개의 NLP 논문 기반 QA 데이터셋, 총 5,049개
• 답변 유형: Answerable/Unanswerable, Yes/No, Abstractive/Extractive
• 정확도는 표준 F1로 성능 측정

QuALITY

• 질문이 평균 5,000개의 토큰 길이로 구성된 객관식 QA 데이터셋
• 전체 문서에 걸친 추론이 뒷받침되어야 하며, 중간 길이의 문서에서 검색하는 능력 테스트
• 대다수의 human annotators가 잘못 주석처리한 질문들을 포함한, 어려운 데이터셋인 QuALITY-HARD 포함

 

Controlled Baseline Comparisons

위의 3가지 데이터셋에서 SBERT, BM25, DPR을 임베딩 모델로 사용하면서, RAPTOR 트리 구조가 있을 때와 없을 때를 비교한다. 언어모델은 UnifiedQA 3B를 사용한다. 아래 두 표에서 확인할 수 있듯이, RAPTOR가 모든 retriever와 결합할 때 모든 데이터셋에서 기본 retriever보다 성능이 향상되었다.

NarrativeQA Performance With + Without RAPTOR

QuALITY and QASPER Performance With + Without RAPTOR

 

SBERT를 사용한 RAPTOR의 성능이 가장 좋기 때문에 모든 후속 실험에서 사용되었다. 아래 표는 GPT-3, GPT-4, UnifiedQA 3개의 LLM을 사용하여 RAPTOR(SBERT), BM25, DPR을 비교한다. 전반적으로 RAPTOR가 적용된 방법들이 성능이 우수한 것을 확인할 수 있다.

QASPER 데이터셋에서의 F-1 점수 비교 (GPT-3, GPT-4, UnifiedQA 3B)

QUALITY 데이터셋에서의 accuracy 비교 (GPT-3, UnifiedQA 3B)

NarrativeQA 데이터셋에서의 여러 모델 간 성능 비교

 

Contribution of the tree structure

각 레이어가 RAPTOR의 검색 능력에 얼마나 기여하는 지 살펴본다.

다른 트리 레이어에 대해 querying할 때의 RAPTOR 성능

 

위 표에서 열은 시작 포인트(highest layer)를 나타내고, 행은 querying할 레이어 수를 나타낸다. 표에서 확인할 수 있듯이, 모든 레이어를 활용하는 전체 트리 검색이 특정 레이어에만 초점을 맞춘 검색보다 성능이 뛰어나다. 이러한 결과는 RAPTOR에서 전체 트리 구조의 중요성을 강조해준다. 검색을 위해 원본 텍스트와 상위 수준의 요약을 모두 제공함으로써 RAPTOR는 고차원 주제에 대한 내용부터 세부적인 내용까지 넓은 범위의 질문을 효과적으로 처리할 수 있다.

 

Conclusion

• 다양한 요약 레벨의 정보로 LLM의 지식을 증강하는, 트리 기반 검색 시스템 RAPTOR
• 클러스터링 및 요약의 반복을 통해 계층적인 트리 구조 생성
• 구축된 트리로부터 querying