💬 NLP

PI-Mem: Pushing Long-Context Reasoning to 3.6M Tokens with Parallel-Iterative Memory

이 논문은 모든 컨텍스트 청크를 동시에 처리하고 강화 학습 기반의 적응형 종료를 통해 공유 메모리를 반복적으로 정교화하는 병렬-반복 메모리 메커니즘인 PI-Mem을 소개하며, 이를 통해 대규모 언어 모델이 최대 360만 토큰에 달하는 긴 컨텍스트 추론 작업에서 탁월한 정확도와 상당한 추론 속도 향상을 달성할 수 있게 합니다.

원저자: Dawei Liu, Haixu Song, Shuang Cheng, Shijie Wang, Haozheng Hou, Kaifeng Liu, Ermo Hua, Zhonghang Yuan, Zhijie Zhong, Yuchen Fan, Biqing Qi, Bowen Zhou

게시일 2026-08-05

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Dawei Liu, Haixu Song, Shuang Cheng, Shijie Wang, Haozheng Hou, Kaifeng Liu, Ermo Hua, Zhonghang Yuan, Zhijie Zhong, Yuchen Fan, Biqing Qi, Bowen Zhou

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 360만 페이지에 달하는 거대한 미스터리 소설을 풀려고 노력하고 있다고 상상해 보세요. 당신에게는 매우 똑똑하지만 단기 기억력이 매우 짧은 천재 탐정, 즉 대규모 언어 모델(LLM)이 있습니다. 만약 당신이 이 탐정에게 책 전체를 한꺼번에 건네준다면, 탐정의 두뇌는 과부하가 걸려 1,000,000페이지에 도달할 때쯤에는 10페이지에 있던 단서를 잊어버리기 시작할 것입니다. 이것이 바로 인공지능의 '롱 컨텍스트(long-context)' 문제, 즉 컴퓨터가 마지막 부분을 읽는 동안에도 어떻게 첫 부분에 대한 주의력을 유지하게 할 것인가에 대한 문제입니다.

이를 해결하기 위해 과학자들은 두 가지 주요 기술을 시도했습니다. 첫 번째는 고무줄을 늘리는 것과 같습니다. 모델이 모든 것을 한꺼번에 기억하도록 강요하는 것인데, 이 경우 고무줄이 끊어지거나 너무 느려지는 문제가 발생합니다. 두 번째로 더 대중적인 기술은 '순환 메모리(Recurrent Memory)' 방식입니다. 이는 탐정이 책을 한 장(chapter)씩 읽는 상황을 상상하게 합니다. 탐정은 1장을 읽은 후, 작은 요약 노트를 작성하고, 그 장을 버린 뒤 2장으로 넘어갑니다. 2장을 읽고, 노트를 업데이트하고, 다시 그 장을 버리는 식입니다. 문제는 만약 탐정이 나중에 혼란스러운 장을 읽게 된다면, 1장에서 얻은 중요한 단서를 적어둔 노트를 실수로 덧쓰거나 지워버려 결정적인 단서를 영원히 잃어버릴 수 있다는 점입니다. 이는 탐정이 다음 노트를 완성할 때까지 기다려야 하는 느린 직렬적 과정이며, 이로 인해 조사 전체가 지체됩니다.

여기서 새로운 연구팀이 PI-Mem(Parallel-Iterative Memory)이라는 신선한 아이디어로 등장합니다. PI-Mem은 탐정이 한 장씩 읽고 노트를 업데이트하며 이동하는 대신, 탐정에게 하나의 초능력을 부여합니다. 그것은 바로 모든 장을 동시에 읽되, 약간의 반전이 있는 방식입니다.

실제 논문에서 PI-Mem이 작동하는 방식은 다음과 같습니다. 탐정이에게 '공유 글로벌 메모리(Shared Global Memory)' 게시판이 있다고 상상해 보세요. 탐정은 순차적으로 책을 읽는 대신, 현재 메모리 게시판의 상태를 가이드 삼아 모든 장을 동시에 살펴봅니다. 각 장에 대해 탐정은 다음과 같이 질문합니다. "이 페이지에 내 게시판에 아직 없는 '새로운' 단서가 있는가?" 만약 답이 '예'라면, 탐정은 그 특정 단서만을 뽑아냅니다. 만약 '아니오'라면, 나머지 잡음들은 무시합니다.

모든 장을 병렬로 스캔하고 나면, 탐정은 찾아낸 '새로운' 단서들만을 가져와 메모리 게시판에 병합하여 더 깨끗하고 완전한 그림을 만들어냅니다. 그런 다음, 이 과정을 반복합니다. 탐정은 다시 모든 장을 살펴보지만, 이번에는 업데이트된 게시판을 가이드로 사용합니다. 예를 들어, 200장에서 발견한 조각과 일치하는 것을 방금 찾았기 때문에 5장에서 놓쳤던 단서가 갑자기 이해될 수도 있습니다. 탐정은 새로운 단서가 발견되지 않거나 제한에 도달할 때까지 이러한 '라운드(round)' 형태의 스캐닝을 계속합니다. 마지막으로, 탐정은 360만 페이지의 책 전체가 아니라, 게시판에 통합된 단서들만을 사용하여 미스터리를 해결합니다.

연구진은 이 방법을 두 가지 다른 AI 모델(Qwen3.5 및 Qwen2.5)을 사용하여 테스트했습니다. 이 테스트는 방대한 문서 전체에서 정보를 찾아내야 하는 까다로운 질문들을 다루는 HotpotQA 벤치마크를 활용했습니다. 그들은 컨텍스트 길이를 무려 360만 토큰(대략 거대한 도서관 규모)까지 밀어붙였습니다.

결과는 놀라웠습니다. 기존의 '하나씩 읽는' 방식(순환 메모리)과 비교했을 때, PI-Mem은 단순히 더 나은 답을 내놓는 데 그치지 않고 훨씬 빠르게 답을 찾아냈습니다. 360만 토큰 테스트에서 PI-Mem은 이전의 최고 방식보다 정확도를 6.25에서 7.81 퍼센트 포인트 향상시켰습니다. 더욱 인상적인 점은, 더 큰 모델에서는 6.1배, 더 작은 모델에서는 2.1배 더 빨랐다는 것입니다. 기존 방식은 탐정이 다음 단계로 넘어가기 전에 반드시 매 단계를 마칠 때까지 기다려야 했기에 매우 느렸습니다. 모든 것을 병렬로 읽는 PI-Mem은 느리고 구불구불한 길을 고속도로로 바꾸어 놓았습니다.

연구팀 또한 AI에게 언제 멈춰야 할지를 가르치기 위해 '강화 학습(Reinforcement Learning)'이라는 특별한 훈련 기법을 사용했습니다. 그들은 충분한 증거를 확보했다면 조사를 빠르게 끝낼 경우 AI에게 '보너스'를 주어, 불필요한 추가 읽기 라운드를 수행하지 않도록 유도했습니다. 이를 통해 시스템이 이미 해결한 장을 다시 읽으며 시간을 낭비하지 않도록 보장했습니다.

요약하자면, 이 논문은 AI가 긴 문서를 '생각'하는 방식을 바꾸는 것—즉, 느린 순차적 기록자에서 병렬적이고 반복적인 조사관으로 전환하는 것—을 통해, 우리가 문서의 시작 부분에 있는 단서를 놓치지 않으면서도 방대한 양의 텍스트에서 복잡한 문제를 해결할 수 있음을 시사합니다. 이는 AI에게 모든 것을 한꺼번에 기억하도록 강요할 필요가 없다는 것을 증명합니다. 그저 단서를 찾을 때마다 퍼즐 조각을 정리할 수 있는 더 나은 방법을 제공하면 될 뿐입니다.

기술 요약: PI-Mem

문제 정의

긴 문맥 추론(Long-context reasoning)은 대규모 언어 모델(LLM)의 핵심적인 병목 현상으로 남아 있습니다. 최근의 순환 메모리(recurrent-memory) 접근 방식들은 고유한 컨텍스트 창을 초과하는 입력을 처리하려고 시도하지만, 두 가지 내재적인 문제에 직면합니다:

증거 덮어쓰기(Evidence Overwrite): 순차적인 청크 단위 업데이트는 초기 중요 증거의 관련성이 완전히 확립되기 전에, 이후의 잠재적으로 무관한 내용이 초기 증거를 덮어쓰는 현상을 발생시킵니다.
지연 시간 및 병렬성 제한: 순환 워크플로우의 엄격한 직렬적 청크 간 의존성은 병렬 처리를 방해하여, 추론 지연 시간이 컨텍스트 길이에 따라 선형적으로 증가하게 만듭니다.

포지셔널 인코딩 보간(positional-encoding interpolation)이나 희소 어텐션(sparse attention)과 같은 기존의 대안들은 높은 계산 비용을 초래하거나 구축하기 어려운 합성 코퍼스에 의존하는 경우가 많습니다.

방법론: PI-Mem

저자들은 병렬로 증거를 수집하고 제한된 횟수의 턴(turn) 동안 공유된 메모리 상태를 반복적으로 정제하도록 설계된 메커니즘인 **PI-Mem (Parallel-Iterative Memory)**을 제안합니다.

추론 워크플로우

워크플로우는 Read-Select-Merge 프로세스로 정의되는 사이클(턴)로 작동합니다:

병렬 읽기(Parallel Read): 청크를 순차적으로 처리하는 순환 모델과 달리, PI-Mem은 모든 문서 청크( $C = \{c_i\}$ )를 병렬로 읽습니다. 각 청크는 동일한 글로벌 메모리 상태( $m^{(k-1)}$ )와 질문( $q$ )을 조건으로 처리됩니다. 이는 순차적 덮어쓰기로 인한 초기 증거의 "망각"을 방지합니다.
선택(Select): 각 청크에 대해 모델은 관찰값( $o_i$ )을 생성하고 명시적으로 <check> 신호(yes 또는 no)를 출력합니다. 새로운 또는 보완적인 증거를 나타내는 yes로 표시된 관찰값만 유지됩니다.
병합(Merge): 선택된 관찰값들은 압축되고 업데이트된 컴팩트한 글로벌 메모리( $m^{(k)}$ )로 통합됩니다. 이 단계는 중복을 제거하면서도 필요한 세부 사항을 보존하며 서로 다른 청크로부터의 증거를 압축하고 통합합니다.
반복 정제 및 종료(Iterative Refinement & Exit): 업데이트된 메모리와 함께 사이클이 반복됩니다. 프로세스는 더 이상 새로운 증거를 생성하는 청크가 없거나(모든 <check> 신호가 no인 경우) 최대 턴 제한( $K$ )에 도달할 때 적응적으로 종료됩니다. 최종 답변은 질문과 통합된 메모리만을 사용하여 생성됩니다.

강화 학습을 통한 학습

PI-Mem은 **그룹 상대 정책 최적화(Group Relative Policy Optimization, GRPO)**를 사용하여 엔드투엔드(end-to-end)로 최적화됩니다. 학습 목적 함수는 다음을 결합합니다:

정확도 보상(Accuracy Reward): 최종 답변의 정답 여부에 기반합니다.
턴 효율성 보상(Turn-Efficiency Reward): 중복되는 턴에 대해 페널티를 부여하는 보조 보상( $r_{turn}$ )입니다. 이는 모델이 충분한 증거가 축적되는 즉시 루프를 종료하도록 유도하여, 성능과 추론 지연 시간 사이의 균형을 맞추도록 합니다.

보상은 궤적(trajectory) 내의 모든 모델 호출(읽기, 병합 및 최종 답변)에 브로드캐스트되어 전체 워크플로우를 최적화합니다.

주요 기여

병렬-반복 메커니즘: 공유된 메모리를 조건으로 병렬 증거 추출을 허용함으로써 청크 처리를 순차적 의존성으로부터 분리하는 새로운 아키텍처를 제안합니다.
엔드투엔드 RL 최적화: 반복 정제 루프를 적응적으로 종료하도록 정책을 훈련하기 위해 턴 효율성 보상을 도입하여 중복 계산을 줄였습니다.
실증적 검증: PI-Mem이 전통적인 정확도-효율성 트레이드오프를 깨뜨린다는 것을 보여주는 종합적인 평가를 수행했습니다.

실험 결과

저자들은 RULER 스위트 내의 HotpotQA (HQA) 벤치마크에서 Qwen3.5-35B-A3B와 Qwen2.5-7B를 사용하여 PI-Mem을 평가하였으며, 컨텍스트 길이를 최대 3.6백만 토큰까지 확장했습니다.

정확도:
- Qwen3.5-35B-A3B에서 PI-Mem은 순환 메모리 베이스라인(MemAgent) 대비 평균 점수를 +6.25 절대 포인트 개선하여, 76.56 대비 81.25를 달 기록했습니다. 특정 3.6M 토큰 설정에서 PI-Mem은 MemAgent보다 +7.81 절대 포인트 앞섰습니다.
- Qwen2.5-7B에서 PI-Mem은 평균 점수를 +6.56 포인트 개선하였으며(84.06 vs 77.50), 모든 보고된 길이에서 MemAgent와 대등하거나 능가하는 성능을 보였고, 특히 긴 컨텍스트 영역에서 큰 이득을 얻었습니다.
- PI-Mem은 또한 분포 외(OOD) 작업 및 LongBench v2 벤치마크에서도 우수한 성능을 달성했습니다.
효율성 (지연 시간):
- PI-Mem은 직접 추론(Vanilla) 및 순환 베이스라인 모두에 비해 상당한 추론 속도 향상을 달성했습니다.
- 3.6M 토크의 경우, PI-Mem은 Qwen3.5-35B-A3B 기준으로 MemAgent보다 6.1배, Vanilla보다 16.9배 빨랐습니다.
- Qwen2.5-7B의 경우, 속도는 MemAgent 대비 2.1배, Vanilla 대비 14.4배 빨랐습니다.
증거 보존:
- 다중 값 바늘 찾기(Multi-values Needle-in-a-Haystack, MV-NIAH) 작업에서, PI-Mem은 최종 메모리에 실제 정답 값의 거의 100%를 유지한 반면, 순환 베이스라인은 덮어쓰기로 인해 상당 부분을 손실했습니다.

의의

본 논문은 PI-Mem이 고유 컨텍스트 창을 확장하는 데만 의존하지 않고도 매우 긴 문서에 대한 복잡한 멀티홉 질의응답을 처리할 수 있는 확장 가능한 접근 방식을 제공한다고 주장합니다. 청크를 병렬로 처리하고 공유 메모리를 반복적으로 정제함으로써, 이 방법은 순환 모델이 덮어쓰는 경향이 있는 증거를 효과적으로 보존합니다. 또한, 강화 학습의 통합은 시스템이 추론 깊이와 계산 비용 사이의 트레이드오프를 적응적으로 관리할 수 있게 하여, 높은 정확도와 낮은 지연 시간을 모두 유지하는 초장문 맥락 추론을 위한 유망한 경로를 제시합니다.