긴문맥 LLM 해석을 위한 효율·신뢰 멀티토큰 기여 추적

초록

FlashTrace는 긴 추론 체인을 갖는 대형 언어 모델에서 다중 토큰 출력에 대한 기여도를 한 번에 계산하고, 중간 추론 토큰을 재귀적으로 추적해 원본 입력까지 중요도를 전파한다. span‑wise 집계와 재귀 기법을 결합해 O(N) 시간에 속하며, RULER, MATH, MorehopQA 등에서 130배 이상의 속도 향상과 기존 방법보다 높은 복원율을 달성한다.

상세 분석

본 논문은 현대 LLM이 수천 토큰에 달하는 체인‑오브‑생각(Chain‑of‑Thought)이나 도구 호출을 통해 복잡한 추론을 수행하는 상황에서, 기존 토큰‑단위 기여도 기법이 직면한 두 가지 근본적인 한계를 정확히 짚어낸다. 첫 번째는 효율성 병목이다. 전통적인 Integrated Gradients, AttnLRP, IFR 등은 목표 토큰 하나당 전체 컨텍스트(N)와의 상호작용을 O(N) 비용으로 계산한다. 목표가 M개의 연속 토큰이라면 전체 비용은 O(M·N)으로 급격히 증가해 5천~1만 토큰 길이의 문서에서는 수시간이 소요된다. 두 번째는 신뢰도 저하이다. 추론 토큰(T)이 중간 단계에서 생성되면, 자동 회귀 구조상 다음 토큰에 대한 조건부 확률이 T에 크게 의존하게 되며, 기여도는 자연스럽게 T에 집중된다. 결과적으로 원본 입력(I)의 기여가 희석돼 복원율(recovery rate)이 급격히 떨어진다.

FlashTrace는 이 두 문제를 동시에 해결한다. Span‑wise Aggregation 단계에서는 각 트랜스포머 레이어의 어텐션 헤드, FFN, residual 흐름을 토큰‑대‑스팬 형태로 미리 집계한다. 구체적으로, 목표 스팬 S의 모든 토큰에 대한 변환 함수 f_{j→i}를 합산해 하나의 스팬 표현 y_S를 만든 뒤, L1‑proximity metric을 이용해 입력 토큰 j와 y_S 사이의 기여도를 한 번의 전방패스만으로 계산한다. 이렇게 하면 O(N) 시간에 전체 스팬에 대한 기여도가 얻어지며, M에 대한 선형 의존성이 사라진다.

두 번째 혁신인 Recursive Attribution은 중간 추론 토큰이 흡수한 기여를 다시 입력 단계로 되돌린다. 첫 번째 홉에서는 출력 스팬 O에 대한 기여를 계산해 높은 점수를 받은 추론 토큰 집합 T̂를 선택한다. 이후 T̂를 새로운 목표 스팬으로 설정하고 동일한 span‑wise 절차를 재실행한다. 이 과정을 K번(보통 1~2 hop) 반복하면, 기여가 단계별로 역전파되어 최종적으로 원본 입력 I에 대한 점수가 회복된다. 논문은 “한 번의 재귀 홉만으로도 복원율이 2배 이상 상승한다”는 실험적 증거를 제시한다.

실험에서는 RULER(긴 문서 검색), MATH(수학 증명), MorehopQA(다중 홉 질문) 세 가지 벤치마크에 대해 Qwen‑3‑8B와 GPT‑4o 등 다양한 모델에 적용했다. 속도 측면에서 FlashTrace는 기존 AttnLRP와 IFR 대비 평균 130× 빠른 처리 시간을 기록했으며, 특히 10k 토큰 컨텍스트에서 2초 내외로 기여도를 산출한다. 신뢰도 측면에서는 복원율이 CoT 상황에서 26%→9%로 급락하던 기존 방법에 비해 FlashTrace는 26%→23% 수준을 유지, 즉 원본 입력 토큰을 거의 그대로 찾아낸다. 또한, 재귀 홉 수를 늘릴수록 기여도가 점진적으로 입력 쪽으로 이동하는 “importance diffusion” 현상이 관찰돼, 모델 내부의 인과 흐름을 정량적으로 파악할 수 있다.

이론적 분석에서는 기존 proximity metric을 스팬 수준으로 확장하면서, 어텐션 가중치와 FFN 변환을 선형 결합해 근사함으로써 메모리 사용량을 O(L·H) (L: 레이어 수, H: 헤드 수) 수준으로 제한한다. 또한, 재귀 단계마다 가중치 스팬을 정규화해 과도한 확산을 방지하고, 중요한 토큰에 집중하도록 설계했다. 한계점으로는 (1) 스팬이 매우 비연속적일 경우 현재 구현이 비효율적일 수 있고, (2) 재귀 깊이가 깊어질수록 누적 근사 오차가 발생할 가능성이 있다. 향후 연구에서는 비연속 스팬 지원과 오차 보정 메커니즘을 도입할 여지가 있다.

총괄하면, FlashTrace는 **효율성(O(N) 시간) + 신뢰성(중간 토큰을 통한 인과 역추적)**이라는 두 축을 동시에 만족시키는 최초의 멀티토큰 기여도 프레임워크이며, 긴 문맥을 다루는 실제 LLM 기반 에이전트에 실시간 해석 가능성을 제공한다.