SproutRAG: Attention-Guided Tree Search with Progressive Embeddings for Long-Document RAG

긴 문서의 증거를 효과적으로 모으려면 문맥의 다층 구성을 보존해야 한다. SproutRAG는 sentence-level chunk를 계층적으로 묶는 attention-가이드 트리를 학습해 다중 granul 합성 임베딩으로 검색을 수행하고, LLM 호출 없이도 다층 증거를 확보한다. 이를 통해 정보 효율 IE를 강화하고, 대규모 문서에서도 정확한 증거를 신속하게 찾아낸다.

단락 1: 문제 정의와 한계 — Transformer의 Self-Attention은 시퀀스 길이에 따라 계산량이 증가하고, 문서를 작은 청크로 자르는 방식은 맥락과 관계를 잃게 한다. 기존 접근은 고정 청크나 요약에 의존하거나 LLM 호출이 필요해 비용이 증가한다. 단락 2: 해결 원리 — SLLM으로 문장을 임베딩하고, 여러 layer와 head의 attention을 α 파라미터로 가중 합산해 문장 간 관계를 나타내는 A를 만든다. A를 바탕으로 대칭화된 M, 그리고 leaves의 임베딩 e(s)에서 자식을 합성해 progressive embedding e(p)를 얻고, bottom-up으로 이진 트리를 구성한다. 단락 3: 달라지는 점 — 트리는 단순히 근접도에 의존하는 평균 대신 학습된 가중치를 사용해 문장 간 의미적 근접도에 맞춘 연결을 형성한다. 내부 노드는 서브트리의 progressive 임베딩을 보유하고, 다층 증거를 검색하기 위한 계층적 빔 탐색으로 여러 해상도에서 후보를 모은다. 단락 4: 요약 효과 — end-to-end 학습으로 증거 품질과 트리 구조를 동시에 개선하고, retrieval 시 external LLM 호출 없이 멀티-그레나리티 검색이 가능해진다. 실험에서 IE가 평균 6.1% 향상되며, 특정 벤치마크에서 최대 8.06 포인트의 개선을 보인다.

단락 1: 전체 접근 방식 — D를 sentence-level로 split하고, SLLM으로 e(s_i)와 Attn^(l,h)를 얻은 뒤 w_{l,h}를 학습 가능 파라미터로 정의하여 A_{ij}를 얻는다. 패턴: 입력 Attn^{(l,h)} -> 가중 합산 -> A_{ij}; 의미: 문장 간 관계의 누적 기여를 나타내는 관계 매트릭스. 단락 2: 핵심 메커니즘 — A_{ij} = sum_{l,h} w_{l,h} Attn^{(l,h)}{ij}, w{l,h} = exp(alpha_{l,h}) / sum_{l',h'} exp(alpha_{l',h'}); M_{ij} = (A_{ij} + A_{ji}) / 2; p 노드 임베딩 e(p) = (e(u) + e(v)) / 2; M_{pr} = max(M_{ur}, M_{vr}). 패턴: 입력 문장들 u,v를 합쳐 p를 만들고, M를 업데이트한다. 단락 3: 트리 구축 및 검색 — Leaves를 초기 노드로 삼아 가장 큰 M_{uv}를 가진 쌍을 합치고, 새 노드의 임베딩과 남은 노드와의 관계를 갱신한다. 결과적으로 𝒯는 Leaves-기반의 트리이며, Online Retrieval은 쿼리 e(q)로 각 노드 v의 sim(q,v) = e(q)ᵀ e(v) / (||e(q)|| ||e(v)||)를 계산하고, 깊이별 beam search로 다층에서 후보를 수집한다. 단락 4: 학습 및 구현 상세 — 손실 함수: L = L_ret + λ L_attn. L_ret은 쿼-패시지 대조 학습으로, q와 양성 p^+ 사이의 sim을 최대화하고 hard negatives와 구분하도록 한다. L_attn은 내부 노드가 형성하는 트리 구조를 정렬하도록, 문장 쌍 (s_i, s_j) 간의 A_{ij} + A_{ji}를 log로 스케일링해 학습한다. 트리 구축에 필요한 파라미터 α, w_{l,h}, 하이퍼 파라미터 λ를 학습한다. 단락 5: 구현/학습 세부 — SLLM은 1.3B 파라미터이고, 문서는 문장 단위로 나눠 청크를 최대 두 문장으로 설정한다. Off-line indexing에서 30K 예제의 query–passage를 샘플링해 학습하고, 8×NVIDIA A100 GPU에서 학습한다. 인퍼런스 시에는 query 인코딩, 트리 순회, reranking만 수행한다. 주요 하이퍼 파라미터: b=5, λ=0.1, τ=0.05, warmup 5%, 3 에폭, learning rate 2e-5(SLLM), 1e-3(aggregation).

단락 1: Retrieval quality — SproutRAG는 4개 벤치마크에서 IE가 최고치를 기록, Dragonball에서 가장 큰 이점(+8.06), SCI-DOCS(+4.65), LegalBench-RAG(+4.90), MS MARCO(+6.83)로 나타난다. 또한 재현 가능한 모든 벤치마크에서 Recall과 Precision의 균형이 좋고, noise가 적은 증거를 다층적으로 확보한다. 단락 2: End-to-End 성능 및 효율 — HotpotQA에서 F1 76.47, WebQuestions에서 F1 68.12, Dragonball의 ROUGE-L 0.372, METEOR 0.389, BERTScore 0.671, 온라인 토큰 4,382, 지연 193 ms를 달성한다. 다른 시스템 대비 더 나은 성능-효율 비율을 보이며, 온라인 토큰 수를 크게 줄이고 latency를 낮춘다. 단락 3: Ablation 및 효율성 분석 — Not trained, L_ret 제거, L_attn 제거 등의 변형에서 IE가 하락한다. Uniform attention이나 Leaf-only Retrieval 등의 대체 설계는 성능 저하를 초래한다. Beam width(b)나 λ의 변화에 따른 민감도는 있고, 5의 기본 설정에서 가장 좋은 Recall-Precision 트레이드를 제공한다. 전반적으로 학습된 attention aggregation, 내부 노드 활용, 빔 탐색의 조합이 최적의 성능을 제공한다.

단락 1: 전체 아키텍처 — 문서를 sentence-level chunks S = {s1,...,sn}로 나눈 뒤, SLLM으로 임베딩 E와 attention 행렬 Attn^(l,h)을 얻고, w_{l,h}를 통해 A_{ij}를 구성한다. [어떤 값을 입력으로] Attn^(l,h) → 가중 합산 → A_{ij}이며, [그 값의 의미] 서로 다른 heads와 layers의 상호작용 정보를 반영하는 문장 간 관계 매핑이다. 단락 2: 수학적/알고리즘적 기반 — w_{l,h} = exp(α_{l,h}) / Σ exp(α_{l',h'}), A_{ij} = Σ_{l,h} w_{l,h} Attn_{ij}^{(l,h)}, M_{ij} = (A_{ij} + A_{ji})/2; e(p) = (e(u) + e(v))/2; M_{pr} = max(M_{ur}, M_{vr}). 이진 트리(bottom-up)에서 가장 큰 M_{uv}를 가진 쌍을 반복 합성하여 노드 p를 생성한다. 단락 3: Prior work 대비 차별점 — uniform averaging의 proximity bias를 줄이기 위해 학습된 가중 합산을 사용하고, 멀티-레이어 간의 관계를 반영하는 트리를 구성한다. Leaf는 문장 임베딩, 내부 노드는 서브트리의 progressive 임베딩을 저장한다. 단락 4: 구현 및 학습 세부 — 인덱싱 시 offline에서 SLLM으로 임베딩, Attn 정보를 이용해 트리 구성; 온라인 검색은 query 인코딩, 트리의 노드를 다층에서 검색, 후보를 재랭킹하여 Top-K를 생성한다. 학습은 L_ret(대조 학습)과 L_attn(attn 정규화)을 합친 손실로 수행한다. 4) 학습/평가 구성 — 30K 예제의 CLaRa에서 학습, 8 GPU, 1.3B SLLM 사용, 슬라이스는 최대 두 문장, beam width 5, λ 0.1, τ 0.05, warmup 5%.

제시된 한계점으로는 1) 현재 트리가 이진 트리(binary tree)로 구성되어 다-다대다(degree > 2) 관계를 포착하는 데 한계가 있을 수 있음; 다분기 구조로 확장 가능성. 2) Offline 인덱싱 및 학습 비용이 필요해, 초기 비용이 발생하며 재학습 없이 벤치마크 간 일반화되더라도 adaptation 비용은 있음. 3) 트리 구성이 고정되며 쿼리 의존적 재배치에 한정되어 있어, 쿼리 특성에 따른 증거 재구성의 유연성이 제한될 수 있음.