지식과 추론을 분리한 DRIFT: 효율적 장문 컨텍스트 처리

초록

DRIFT는 경량 지식 모델이 질의에 맞춰 문서 조각을 압축해 암묵적 사실 토큰으로 변환하고, 이를 대형 추론 모델의 임베딩 공간에 투사해 장문 입력을 효율적으로 처리한다. 32배 압축에서도 LongBench v2 점수를 20.87%→29.22%로 끌어올리며, 256k 토큰 문서에 대해 평균 7배 빠른 추론 속도를 달성한다.

상세 분석

DRIFT는 “지식 추출 ↔ 추론”을 명시적으로 분리하는 이중‑모델 프레임워크다. 작은 지식 모델 ψₖₙₒ는 질의 Q와 문서 청크 Cⱼ를 입력받아, <|CPS|> 토큰을 추가한 뒤 마지막 레이어의 은닉 상태를 “암묵적 사실 토큰”(implicit fact tokens, Tⱼ)으로 출력한다. 이 토큰은 고차원(d) 벡터이며, MLP 투사기 π를 통해 추론 모델 θᵣₑₐ의 임베딩 공간에 정렬된 “암묵적 사실 임베딩”(E)으로 변환된다. 추론 모델은 기존 텍스트 대신 E와 질의 임베딩 E(Q)를 받아 복잡한 논리 연산을 수행한다.

핵심 기술은 (1) 동적 버킷 압축이다. 고정 비율이 아닌 토큰 길이 구간(예: 64–128, 128–256)별 상한을 기준으로 출력 토큰 수를 결정해, 정보가 희소한 구간에서도 충분히 압축된 표현을 유지한다. 이는 “핵심 문장 ↔ 불필요 문장” 비대칭성을 반영해 압축 효율을 크게 높인다.

(2) 세 단계 학습 파이프라인.

• LFRP (Latent Fact Reconstruction Pretraining): 질의와 무관하게 정적 압축 비율 cₛₜₐ=8로 문서를 압축하고, 고정된 추론 디코더를 이용해 원문 재구성을 목표로 손실을 계산한다. 이때 그래디언트는 오직 ψₖₙₒ와 π에만 흐르며, 토큰이 원문 복원에 충분히 정보를 담도록 학습한다.
• QAFT‑DC (Query‑Aware Fine‑Tuning – Dynamic Compression): 질의를 포함한 압축 지시문 I_dyn을 사용해, ψₖₙₒ가 질의‑관련 정보를 선택적으로 인코딩하도록 유도한다. 압축 비율을 가변적으로 조정하면서도 재구성 손실과 질의‑응답 손실을 동시에 최소화한다.
• QAFT‑QA (Query‑Aware Fine‑Tuning – Answer Generation): 최종 단계에서 θᵣₑₐ를 미세조정해, 암묵적 사실 임베딩을 기반으로 실제 답변을 생성한다.

실험에서는 Mistral‑7B 기반 추론 모델에 DRIFT를 적용했으며, 256k 토큰 문서를 32배 압축(≈8k 토큰)했을 때 평균 7× 속도 향상과 메모리 사용량 감소를 기록했다. LongBench v2에서 20.87%→29.22%로 정확도가 크게 상승했으며, 64×·128× 압축에서도 경쟁력 있는 성능을 유지했다. 또한 300K 규모의 Document‑QA‑Evidence 데이터셋을 자체 구축해, 청크‑단위 정밀 라벨링과 증거 추출을 제공함으로써 학습 안정성을 확보했다.

비교 대상인 COCOM, xRAG, E2LLM 등은 여전히 검색 단계에 의존하거나 정적 압축에 머물러, 질의‑특정 정보 손실이 빈번했다. DRIFT는 질의‑조건부 압축과 투사 기반 임베딩 정렬을 통해 이러한 한계를 극복한다. 한편, 현재 구현은 사전 정의된 청크 분할과 고정된 투사 MLP에 의존하므로, 매우 비구조적 텍스트(예: 코드, 표)에서는 추가적인 전처리와 토큰화 전략이 필요할 수 있다. 또한, 두 모델 간 인터페이스가 고정돼 있어, 다른 규모의 추론 모델이나 멀티모달 입력에 대한 확장성은 향후 연구 과제로 남는다.

요약하면, DRIFT는 (1) 지식‑추론 분리를 통한 컨텍스트 효율화, (2) 동적 버킷 압축으로 정보 손실 최소화, (3) 단계별 사전학습·미세조정으로 압축‑추론 파이프라인을 최적화한다는 세 가지 혁신을 제시한다. 이는 장문·지식‑집약형 작업에서 LLM의 실용성을 크게 높이는 방향성을 제공한다.