PI-Mem: 병렬 반복 메모리를 이용한 360만 토큰 길이의 장문 맥락 추론 성능 향상

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

장문 맥락 추론은 대규모 언어 모델의 핵심적인 난제로 남아 있으며, 최근의 순환 메모리 접근 방식은 다음과 같은 두 가지 근본적인 문제에 직면합니다. 첫째, 순차적인 청크 단위 업데이트는 초기에 중요한 증거를 나중에 관련 없는 내용으로 덮어쓸 수 있습니다. 둘째, 청크 간의 직렬적 의존성은 병렬성을 제한하고 맥락 길이에 따라 지연 시간을 증가시킵니다. 이러한 문제를 해결하기 위해 우리는 모든 청크를 병렬로 처리하고 공유 메모리를 제한된 횟수의 반복을 통해 정교화하는 방식인 PI-Mem(Parallel-Iterative Memory)을 제안합니다. 각 단계에서 PI-Mem은 현재 메모리에 따라 모든 청크를 병렬로 읽고, 각 청크에서 새로운 또는 보완적인 증거를 선택한 다음, 선택된 증거를 다음 단계의 압축된 공유 메모리로 통합합니다. 불필요한 반복을 방지하기 위해, 우리는 보조적인 효율성 보상을 사용하는 강화 학습을 통해 워크플로우를 최적화하여 모델이 충분한 증거가 축적되면 적응적으로 종료할 수 있도록 합니다. Qwen3.5-35B-A3B 및 Qwen2.5-7B 모델을 사용하여 HotpotQA 벤치마크에서 맥락 길이를 최대 360만 토큰까지 늘려 평가한 결과, PI-Mem은 순환 메모리 기준 성능보다 각각 +6.25점과 +7.81점의 절대적인 성능 향상을 보였으며, 추론 속도는 각각 6.1배와 2.1배 빨라졌습니다. 이러한 결과는 PI-Mem이 장문 맥락 추론에서 정확성과 효율성 간의 균형을 깨뜨리고 매우 긴 문서에 대한 복잡한 다중 단계 질문 응답에 적용 가능한 확장 가능한 접근 방식을 제공함을 보여줍니다.

Long-context reasoning remains a critical bottleneck for large language models, as recent recurrent-memory approaches face two inherent challenges: sequential chunk-wise updates can overwrite early critical evidence with later irrelevant content, and serial inter-chunk dependencies limit parallelism and cause latency to increase with context length. To address these issues, we propose PI-Mem (Parallel-Iterative Memory), a mechanism that processes all chunks in parallel and iteratively refines a shared memory over a bounded number of turns. In each turn, PI-Mem reads all chunks in parallel conditioned on the current memory, selects new or complementary evidence from each chunk, and merges the selected evidence into a compact shared memory for the next turn. To discourage redundant turns, we optimize the workflow through reinforcement learning with an auxiliary turn-efficiency reward, enabling the model to adaptively exit once sufficient evidence has been accumulated. We evaluate PI-Mem with Qwen3.5-35B-A3B and Qwen2.5-7B on the HotpotQA benchmark across context lengths up to 3.6 million tokens and find that it outperforms the recurrent-memory baseline by +6.25 and +7.81 absolute points while achieving 6.1$\times$ and 2.1$\times$ inference speedups, respectively. These results demonstrate that PI-Mem breaks the accuracy--efficiency trade-off in long-context reasoning and provides a scalable approach to complex multi-hop question answering over extremely long documents.