Attention 기반 모델의 내부 Retrieval: INTRA(InTrinsic Retrieval via Attention)

전통적 RAG 파이프라인은 retriever와 generator 간 표현 공간의 차이로 비효율이 존재한다. INTRA는 decoder cross-attention의 쿼리-기반 매칭을 활용해 외부 retriever 없이도 증거를 검색하고 생성에 직접 활용한다. 이로써 멀티홑 QA에서 증거 재구성 및 엔드-투-엔드 품질이 향상된다.

Transformer의 cross-attention은 쿼리-문맥에 대한 정보 선택을 수행한다. INTRA는 이를 활용해 decoder가 필요로 하는 쿼리 상태(qℓ)와 encoder의 조각 ki 간의 매칭 점수를 MaxSim 방식으로 산정하고, 상위 chunk를 SINTRA로 선택해 y를 생성한다. Reverse-QWK를 사용해 다층에서의 키-프로젝션을 쿼리 쪽으로 이동시켜 K¯(S) 하나로 모든 레이어를 커버하고, mean-pooling된(chunk) 벡터를 이용해 계산 복잡도를 감소시킨다. 초기 컨텍스트 S0는 x와의 유사도로 구성되되 최종 SINTRA의 점수로 전체 코퍼스를 스캔한다. 이 프레임워크는 외부 retriever 없이도 retrieval와 generation을 하나의 모델에서 처리한다.

전체 접근은 x를 입력으로 받아 Enc로 증거 ki를 얻고 ki의 모음 K(S)을 구성한다. S0는 s(0)i = MaxSim(kx, ki)로 초기화되며, Dec g를 통해 qℓ를 산출하고 MaxSim(qℓ, ki)로 si를 계산한다. αℓ로 가중된 MaxSim의 합으로 si를 얻고, SINTRA = {i | si가 상위인 인덱스}로 선택한다. y = Dec(x, K(SINTRA))로 생성한다. Reverse-QWK는 qeℓ = (qℓ W⊤K,ℓ) ⊙ γK,ℓ로 정의하고 zℓ = Attention(qeℓ, K¯(S), K¯(S))로 cross-attention을 수행한다. Mean-pooled 임베딩 bki를 사용해 MaxSim 연산의 비용을 줄인다. 학습은 retrieval tokens ρ와 layer-aggregation α를 조정하는 soft cross-entropy 손실로 수행한다.

완전한 증거 기억 측정에서 INTRA가 다중 홉 벤치마크에서 우수하다. HotPotQA: R@5=59.9%, R@10=70.9%, R@20=76.1%; 2Wiki: R@5=40.7%, R@10=50.3%, R@20=55.2%; MuSiQue: R@5=12.8%, R@10=18.9%, R@20=23.7%; NQ: R@5=29.1%, R@10=38.3%, R@20=45.9%. 엔드-투-엔드 EM/F1은 HotPotQA 46.4/58.0, 2Wiki 49.2/53.2, MuSiQue 14.0/23.0, NQ 51.2/60.3으로 보고된다. ablation은 S0의 초기 컨텍스트, 다중 retrieval 토큰의 중요성을 시사한다. TTFT 등 효율성 측면에서 INTRA는 재사용 가능한 컨텍스트를 통해 RAG 대비 생성 시간과 초기 토큰 도달 시간을 줄이고, 1B 토큰 풀이 2.56 TB의 8-bit 저장으로 가능하다.

프레임워크는 K(S)로 구성된 pre-encoded chunk를 재사용하고, MaxSim 기반의 late-interaction으로 si를 계산한다. Reverse-QWK로 다층 cross-attention의 keys를 재계산하지 않고도 점수를 유지한다. 평균 풀링된 벡터(bKi)를 사용해 효율성을 확보하고, S0 초기 컨텍스트를 통해 검색의 시작점을 제공한다. 학습은 retrieval 토큰과 α를 조정하는 soft cross-entropy 손실을 사용한다.

고정된 컨텍스트 풀에 의존하며, 웹 규모의 오픈 도메인 검색에 일반화됐는지 불확실하다. 단일 구현군(T5Gemma2 + Reverse-QWK) 기반이며 다른 아키텍처/모달리티에 대한 일반화는 확인되지 않았다. 또한 decoder-만 모델이므로 decoder-only 아키텍처에는 적용되지 않는다.