💬 NLP

PReM: Learning What to Preserve and When to Refresh for Context Compression

PReM은 전용 메모리 레이어와 특수 토큰을 통해 내부 레이어별 KV 메모리를 동적으로 보존하고 갱신하는 법을 학습함으로써, 외부 압축기에 의존하거나 성급한 유지 결정을 내리지 않고도 효율적인 롱 컨텍스트 추론을 가능하게 하는 새로운 컨텍스트 압축 프레임워크입니다.

원저자: Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

게시일 2026-07-17

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Bohan Yu, Lei Shen, Chenxi Zhou, Chen Han, Junlin Liu, Wenbo Su, Yu Cheng, Bo Zheng

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

거대한, 복잡한 미스터리를 풀려고 노력하고 있다고 상상해 보세요. 하지만 단 하나의 단서 대신, 당신은 32,000페이지의 문서가 담긴 도서관을 건네받았습니다. 당신의 뇌(이 경우에는 거대 언 언어 모델이라는 컴퓨터 프로그램)는 질문에 답하는 데 필요한 특정 사실들을 찾아내기 위해 이 모든 것을 읽어야 합니다. 문제는 그 전체 도서관을 한 번에 "작업 기억(working memory)"에 담아두는 것이 매우 무겁고 느리다는 점입니다. 이는 마치 배낭에 도서관 전체를 넣고 경주를 하는 것과 같습니다. 결국 당신은 지치게 될 것이고, 정작 필요한 책을 떨어뜨릴 수도 있습니다.

이를 해결하기 위해 과학자들은 지금까지 두 가지 주요 기술을 시도해 왔습니다. 첫 번째는 "압축(compression)"으로, 당신이 읽기를 시작하기도 전에 도서관을 아주 작은 요약본 하나로 줄이는 것입니다. 두 دوم 번째는 "캐싱(caching)"으로, 도서관 전체를 보관하되 가장 인기 있는 페이지들만 골라 보는 방식입니다. 하지만 두 방법 모두 결함이 있습니다. 바로 문제를 풀기 전에 무엇을 남길지 결정한다는 점입니다. 만약 미스터리가 진행되는 도중에 상황이 바뀌어 책 뒷부분에 있는 페이지가 갑자기 필요해진다면, 당신은 곤경에 처하게 됩니다. 이미 버려진 페이지를 다시 가져올 수 없기 때문입니다. 왜냐하면 결정이 너무 일찍 내려졌기 때문입니다. 이 논문은 이 문제를 다루는 새로운 방법을 소개하며, 정적인 요약 대신 이야기가 전개됨에 따라 스스로 업데이트할 수 있는 메모리 시스템이 필요하다고 제안합니다.

여기 PReM(Preserve and Refresh Memory)이 있습니다. 이는 모델의 메모리를 정적인 파일 캐비닛이 아닌, 동적이고 살아있는 노트처럼 취급하는 영리한 새로운 프레임워크입니다. 핵심 아이디어는 간단하지만 강력합니다. 한 번만 무엇을 남길지 결정하는 것이 아니라, 지금 당장 무엇을 남길지 결정하고, 이야기가 요구하는 즉시 그것들을 교체할 준비를 하는 것입니다.

PReM이 어떻게 작동하는지 재미있는 비유를 통해 설명하겠습니다. 당신이 탐정 팀(AI의 서로 다른 레이어들)과 함께 사건을 해결하고 있다고 상상해 보세요. 기존 방식에서는 팀원들이 사건 파일을 통째로 읽고, 가장 마음에 드는 10페이지만 골라낸 뒤 나머지는 금고에 잠가 버립니다. 나중에 금고에 있는 페이지가 필요해지더라도 방법이 없습니다.

PReM은 규칙을 바꿉니다. 전체 사건 파일을 유지하되, 이를 작은 "덩어리(chunks)"(마치 챕터처럼)로 조직합니다. 탐정이 해결책을 쓰기 시작하면, 특별한 "메모리 토큰"(이를 마법의 책갈피라고 부르며, **⟨m⟩**으로 표기함)이 신호 역할을 합니다. 탐정이 이 책갈피를 쓴다는 것은, "잠깐! 내 책상 위에 무엇이 있는지 다시 생각해 봐야겠어"라고 말하는 것과 같습니다.

그 순간, 전용 "메모리 매니저"(AI의 특정한 레이어)가 깨어납니다. 이 매니저는 현재의 질문과 마법의 책갈피를 살펴본 뒤, 모든 챕터를 빠르게 스캔합니다. 그리고 결정합니다: "좋아, 다음 이야기 단계에서는 반드시 챕터 3과 챕터 12가 필요해. 나머지는 일단 잊어도 좋아." 그런 다음 기존의 챕터들을 책상에서 치우고 새로운 챕터들로 교체하며, 불필요한 것들은 아주 작은 요약본으로 압축하면서 중요한 세부 사항들은 선명하게 유지합니다. 이 과정은 글을 쓰는 도중에 발생하며, 사전에 이루어지는 것이 아닙니다.

논문은 이 "새로고침(refresh)" 메커니즘이 게임 체인저임을 보여줍니다. 32,000 토큰 컨텍스트(엄청난 양의 텍스트입니다!)를 사용하는 테스트에서, PReM은 메모리를 16배 또는 32배까지 압축할 수 있었습니다. 한 번에 책상 위에 올려두는 정보량이 훨씬 적었음에도 불구하고, PReM은 도서관 전체를 펼쳐두려는 모델들보다 실제로 질문에 더 잘 답변했습니다. 예를 들어, 70억 개의 파라미터를 가진 모델에서 PReM은 기존의 가장 좋은 방법들보다 정확도를 최대 12.55포인트 향상시켰습니다.

연구진은 또한 이 "메모리 매니저"가 AI의 뇌 어디에 위치해야 하는지에 대해 흥미로운 사실을 발견했습니다. 그들은 메모리 매니저를 AI의 초기, 중간, 또는 후기 레이어에 배치하여 테스트했습니다. 그 결과, 중간과 후기 레이어가 무엇을 남길지 결정하는 데 훨씬 더 뛰어났던 반면, 초기 레이어는 이 기술을 배우기에 너무 느리다는 것을 발견했습니다. 이는 마치 신참 형사(초기 레이어)가 아니라 선임 형사(중간 레이어)가 어떤 증거가 관련 있는지 판단하게 하는 것과 같습니다.

PReM은 더 작은 모델(30억 파라미터)이 PReM을 사용할 때, 다른 압축 기술을 사용하는 더 큰 모델(70억 파라미터)보다 성능이 뛰어날 수 있다는 놀라운 사실을 보여주었습니다. 이는 메모리를 언제 새로고침할지 아는 것이 얼마나 큰 뇌를 가졌는지 못지않게 중요하다는 것을 시사합니다.

또한 이 논문은 몇 가지 아이디어를 배제합니다. 단순히 AI의 자연스러운 어텐션(현재 무엇에 집중하고 있는지 보는 것)을 사용하는 것만으로는 무엇을 남길지 결정하기에 충분하지 않으며, 특화된 학습 시스템이 필요하다는 것을 보여줍니다. 또한 텍text를 처음에 한 번 압축하고 나서 잘 되기를 바라는 방식은 효과가 없음을 보여줍니다. 즉, "정적"인 접근 방식(전체 답변 동안 동일한 압축 요약본을 유지하는 것)은 특히 복잡한 다단계 질문에서 더 나쁜 결과를 초래합니다.

요약하자면, PReM은 효율적인 AI의 미래가 단순히 데이터를 줄이는 것이 아니라, 숨 쉬는 메모리를 구축하는 것에 있다고 제안합니다. PReM은 대화의 현재 단계에 필요한 증거를 보존하는 법을 배우고, 다음 단계를 위해 새로운 증거를 가져오도록 "새로고침" 버튼을 누를 정확한 타이밍을 압니다. 이렇게 함으로써, 미스터리를 풀기 위해 필요한 단서를 잊지 않으면서도 AI를 빠르고 효율적으로 유지합니다.

기술 요약: PReM (Preserve and Refresh Memory)

1. 문제 정의

대규모 언어 모델(LLM)의 효율적인 롱 컨텍스트(long-context) 추론은 계산 비용을 줄이는 것과 필요한 문맥적 증거에 대한 접근성을 유지하는 것 사이의 근본적인 긴장 관계에 직면해 있습니다. 컨텍스트 길이가 증가함에 따라 Key-Value(KV) 캐시가 성장하며, 이는 디코딩 지연 시간과 메모리 사용량을 크게 증가시킵니다.

이를 완화하기 위한 기존 방식들, 예를 들어 KV 캐시 압축 및 컨텍스트 압축은 두 가지 주요 한계점을 가집니다:

정적 증거 유지 (Static Evidence Retention): 대부분의 방법은 어떤 컨텍스트를 유지할지에 대해 초기에 단 한 번의 결정을 내립니다. 이러한 정적 압축은 서로 다른 추론 단계가 서로 다른 증거 조각들을 필요로 할 수 있는 다단계 추론의 진화하는 요구에 적응하지 못합니다.
결합되지 않은 증거 선택 (Decoupled Evidence Selection): 외부 압축기에 의존하는 방식은 압축 과정과 생성을 분리합니다. 이는 추가적인 계산을 유발하며, 모델이 특정 추론 단계에 맞춰 동적으로 증거 선택을 조정하는 것을 방해합니다.

본 논문은 효과적인 압축이 단순한 크기 축소를 넘어, 유용한 증거를 동적으로 보존하고 이후의 추론 단계에서 다른 정보가 필요할 때 이를 갱신(refresh)할 수 있어야 한다고 주장합니다.

2. 방법론: PReM

PReM (Preserve and Refresh Memory)은 롱 컨텍스트를 LLM 내부의 레이어별 KV 메모리로 유지하는 컨텍스트 압축 프레임워크입니다. 외부 압축기에 의존하는 대신, PReM은 모델의 내부 상태를 기반으로 무엇을 보존할지 그리고 언제 갱신할지를 결정하는 법을 학습합니다.

핵심 메커니즘

레이어별 KV 메모리 표현 (Layer-Wise KV Memory Representation): 롱 컨텍스트는 모델 내부에 저장된 레이어별 KV 상태( $M_\ell$ )로 인코딩됩니다. 이 상태들은 연속적인 청크(chunk)로 분할됩니다.
전용 메모리 레이어 (Dedicated Memory Layer): 특정 레이어( $\ell_m$ $ℓ_{m}$ )가 메모리 선택 결정을 내리도록 지정됩니다. 이 레이어는 컨텍스트 청크의 점수를 매기기 위해 전용 선택 쿼리 투영 행렬( $W^c_\ell$ $W_{ℓ}^{c}$ )을 사용합니다.
- 청크 선택 (Chunk Selection): 모델은 현재 쿼리(질문 또는 갱신 토큰으로부터의 쿼리)를 기반으로 청크의 점수를 매깁니다.
- 보존 및 풀링 (Preserve-and-Pool): 선택된 청크들은 토큰 수준의 정밀도로 보존됩니다. 선택되지 않은 청크들은 하나의 평균화된 KV 대표값으로 압축됩니다.
메모리 토큰 $\langle m \rangle$ : 특수 토큰은 메모리를 갱신하기 위한 명시적인 내부 신호 역할을 합니다. 모델이 $\langle m \rangle$ 을 생성하면, 메모리 레이어는 컨텍스트 청크를 재평가하여 새로운 관련 증거를 선택하고, 후속 레이어를 위해 압축된 KV 메모리를 업데이트합니다. 이를 통해 모델은 현재 추론 단계에 따라 언제 갱신할지를 제어할 수 있습니다.

학습: 단계 분리 갱신 학습 (Phase-Separated Refresh Training)

메모리 선택을 생성과 정렬하도록 학습시키기 위해, PReM은 단계 분리 갱신 학습 전략을 도입합니다. 각 추론 단계는 두 개의 뚜렷한 순전파(forward pass)로 나뉩니다:

메모리 선택 단계 (Memory-Selection Phase): 모델은 입력(질문 또는 $\langle m \rangle$ )을 처리하여 청크 점수를 계산하고 상위- $k$ 개의 증거 청크를 선택합니다. 이 단계는 모델이 증거가 아닌 청크보다 증거 청크의 순위를 높게 매기도록 유도하는 메모리 선택 손실( $L_{ms}$ )로 감독됩니다.
메모리 조건부 생성 단계 (Memory-Conditioned Generation Phase): 모델은 선택 단계의 결과로 생성된 새로 갱신된 메모리 상태를 조건으로 하여 답변 세그먼트( $y_s$ )를 생성합니다.

학습과 추론 사이의 일관성을 보장하기 위해, 학습 목적 함수에는 경계 교차 엔트로피 항( $L_{bd}$ )이 포함됩니다. 이 항들은 선택 단계를 첫 번째 답변 토큰과 연결하고, 답변 세그먼트의 끝을 다음 $\langle m \rangle$ 토큰과 연결하여, 모델이 메모리 업데이트와 생성 사이의 자기회귀적(autoregressive) 연결을 학습하도록 보장합니다.

3. 주요 기여

정식화 (Formulation): 본 논문은 컨텍스트 압축을 학습된 메모리 보존 및 갱신 문제로 재정의하여 외부 압축기의 필요성을 제거했습니다.
아키텍처 (Architecture): PReM은 전용 메모리 레이어와 특수 $\langle m \rangle$ 토큰을 도입하여, LLM의 내부 KV 상태 내에서 동적이고 단계별인 메모리 업데이트를 가능하게 합니다.
학습 전략 (Training Strategy): 메모리 선택을 생성과 분리하여 학습함으로써, 모델이 해당 단계에 필요한 특정 메모리 상태를 조건으로 생성을 수행하도록 보장하는 단계 분리 갱신 학습을 제안했습니다.
실증적 검증 (Empirical Validation): 학습된 동적 메모리 업데이트가 정적 압축 및 전체 컨텍스트 프롬프팅보다 우수함을 입증하는 종합적인 실험을 수행했습니다.

4. 실험 결과

실험은 32K 토큰 컨텍스트를 가진 6개의 벤치마크(TriviaQA, SQuAD, NaturalQuestions, 2WikiMQA, HotpotQA, MuSiQue)에 대해 Qwen2.5-Instruct의 3B 및 7B 변형 모델을 사용하여 진행되었습니다.

베이스라인 대비 성능 (Performance vs. Baselines): 16배 및 32배 압축 비율 모두에서, PReM은 KV 캐시 방식(SnapKV, StreamingLLM, CAKE)과 컨텍스트 압축 방식(LongLLMLingua, EXIT, ICAE)을 포함한 8개의 강력한 베이스라인보다 우수한 성능을 보였습니다.
- 3B 백본의 16배 압축 시, PReM은 가장 강력한 베이스라인 대비 평균 Exact Match(EM)와 F1에서 각각 5.30 및 5.15 포인트 향상되었습니다.
- 7B 백본의 32배 압축 시, 향상 폭은 EM 8.20, F1 9.89 포인트에 달했습니다.
전체 컨텍스트와의 비교 (Comparison to Full Context): PReM은 3B 및 7B 모델 모두에서 직접적인 32K 토큰 프롬프팅(압축 없음)보다 우수한 성능을 보였으며, 이는 학습된 메모리 선택이 모델에 필터링되지 않은 전체 컨텍스트를 노출하는 것보다 더 효과적임을 시사합니다.
소형 모델의 우월성 (Small Model Superiority): 3B PReM 모델은 출시된 7B 소프트 컨텍스트 압축 베이스라인(예: Activation Beacon)보다 높은 성능을 보였으며, 이는 학습된 보존 및 갱신이 더 작은 백본 규모를 보완할 수 있음을 나타냅니다.
효율성 (Efficiency): PReM은 경쟁력 있는 추론 지연 시간을 유지하면서 피크 GPU 메모리를 크게 줄였습니다 (예: 전체 메모리 사용 시 30.0GB 대비 약 19.57GB). $\langle m \rangle$ 토큰의 사용은 디코딩 계산(갱신 빈도)과 메모리 절감 사이의 조절 가능한 트레이드오프를 허용합니다.

5. 의의 및 주장

본 논문은 PReM이 효과적인 롱 컨텍스트 추론을 위해서는 정적인 컨텍스트 단축이 아닌 생성과 결합된 메모리 업데이트가 필요함을 입증한다고 주장합니다. 컨텍스트를 내부 KV 메모리로 유지하고 이를 동적으로 갱신하도록 학습함으로써, 모델은 각 추론 단계에 필요한 특정 증거에 적응적으로 접근할 수 있습니다.

저자들은 자신들의 접근 방식이 정적 유지와 분리된 선택의 한계를 해결하며, 모델이 "무엇을 보존하고 언제 갱신할지"를 학습할 수 있음을 보여준다고 강조합니다. 이는 멀티홉(multi-hop) 추론 작업에서 더 우수한 성능을 보이며, 기존 압축 기술에 비해 더 나은 효율성-품질 트레이드오프를 제공합니다. 이 연구는 생성과 공동으로 학습된 내부 메모리 관리가 과도한 계산 비용 없이 LLM을 더 긴 컨텍스트로 확장하는 데 있어 실행 가능한 경로임을 시사합니다.