메모리 캐싱: 성장하는 메모리를 갖춘 RNN

기존 RNN은 메모리 크기가 고정되어 긴 정보를 기억하지 못하고, Transformer는 긴 정보를 처리할 때 연산 비용이 너무 많이 든다. 이 논문은 RNN의 효율성을 유지하면서도 필요한 정보를 선별적으로 저장하고 불러오는 '메모리 캐싱'을 통해 두 모델의 장점을 결합한 새로운 대안을 제시한다.

RNN은 과거의 모든 정보를 하나의 작은 가방(Hidden State)에 억지로 구겨 넣는 것과 같다. 시퀀스가 길어지면 가방이 넘쳐서 옛날 정보는 빠져나가게 되는데, 이것이 RNN이 긴 문서를 처리하지 못하는 근본적인 이유이다. 반면 Transformer는 모든 정보를 바닥에 펼쳐놓고 보지만, 정보가 많아질수록 바닥 면적(연산량)이 기하급수적으로 넓어져야 한다.

이 논문은 RNN이 정보를 처리할 때마다 중요한 순간의 가방 상태를 사진으로 찍어 '앨범(캐시)'에 보관하는 방식을 제안한다. 현재 토큰을 처리할 때, 지금 들고 있는 가방뿐만 아니라 앨범에 있는 과거의 사진들을 훑어보며 관련 있는 정보를 다시 꺼내 쓴다. 이는 마치 책을 읽으면서 중요한 페이지에 포스트잇을 붙여두었다가 나중에 다시 들여다보는 것과 같은 원리이다.

결과적으로 RNN은 여전히 가벼운 가방을 들고 다니면서도(효율성), 필요할 때는 앨범을 뒤져 아주 예전의 정보까지 정확하게 기억해낼 수 있게 된다(확장성). 이를 통해 연산량은 시퀀스 길이에 비례해서만 늘어나면서도, 기억력은 Transformer에 근접하는 효과를 거둔다.

전체 시퀀스를 N개의 세그먼트로 나누고, 각 세그먼트가 끝날 때의 메모리 상태를 캐시에 저장한다. 현재 시점 t의 출력 yt를 계산할 때, 현재의 온라인 메모리 Mt와 과거의 캐시된 메모리들을 집계 함수 Agg를 통해 결합한다. 이 과정에서 각 캐시된 메모리는 현재 쿼리 qt에 대한 응답을 생성하는 독립적인 압축기로 작동한다.

Gated Residual Memory (GRM) 전략은 입력 종속적 게이트를 사용한다. [현재 입력 ut와 각 세그먼트의 평균 표현값 간의 내적 계산] → [Softmax를 통한 가중치 할당] → [가중치와 각 세그먼트 메모리의 출력값 곱셈] → [현재 온라인 메모리 출력과 합산] 순으로 연산하여 쿼리와 가장 유사한 맥락의 과거 정보를 복원한다.

Sparse Selective Caching (SSC)은 효율성을 위해 라우터 메커니즘을 도입한다. [현재 토큰과 과거 세그먼트들 간의 관련성 점수 산출] → [상위 k개의 가장 관련 있는 세그먼트만 선택] → [선택된 세그먼트의 메모리만 로드하여 연산] 과정을 거친다. 이는 Mixture of Experts(MoE)와 유사하게 동작하여, 초장기 시퀀스에서도 모든 캐시를 읽지 않고 필요한 정보만 빠르게 추출할 수 있게 한다.

760M 및 1.3B 파라미터 규모의 실험에서 MC를 적용한 SWLA, DLA, Titans 모델은 기본 모델 대비 WikiText-103 및 상식 추론 태스크에서 일관된 성능 향상을 보였다. 특히 Titans 모델에 MC를 결합했을 때 평균 정확도가 약 0.8% 상승하며 가장 우수한 효율을 기록했다.

긴 문맥 회상 능력을 측정하는 Needle-In-A-Haystack(NIAH) 테스트에서 MC의 효과가 극명하게 나타났다. 기존 DLA 모델이 16K 길이에서 44%의 정확도에 그친 반면, MC(GRM)를 적용한 모델은 82.4%를 기록하며 두 배 가까운 성능 향상을 보였다. 이는 캐싱된 메모리가 장기 기억 보존에 결정적인 역할을 함을 입증한다.

LongBench와 같은 복합 긴 문맥 이해 벤치마크에서도 MC 강화 모델들은 Transformer와 대등한 수준의 정확도를 달성했다. 동시에 훈련 및 추론 처리량(Throughput) 분석 결과, 시퀀스 길이가 길어질수록 Transformer 대비 월등히 높은 속도를 유지하여 실무적 효율성을 증명했다.

MC는 RNN의 상태 업데이트 규칙을 변경하지 않고 출력 단계에서 캐시를 참조하는 비침습적(Non-invasive) 구조를 가진다. 이는 기존의 Linear Attention이나 Titans 등 다양한 순환 아키텍처에 즉시 통합될 수 있는 범용성을 제공한다.

수학적으로 MC는 yt = γt Mt(qt) + Σ γi Mi(qt) 형태로 정의된다. 선형 메모리 모듈을 사용할 경우, 세그먼트 크기를 1로 설정하고 적절한 게이팅을 적용하면 이 구조가 Gated Global Attention과 수학적으로 동등해짐을 증명하여 Transformer와의 이론적 연결 고리를 제시했다.

Memory Soup 변체는 가중치 평균(Weight Souping) 개념을 메모리 파라미터에 직접 적용한다. 여러 세그먼트의 파라미터를 입력 데이터에 기반해 보간하여 새로운 입력 종속적 메모리 모듈을 생성하며, 이는 비선형 메모리 모듈(DLA, Titans)에서 단순 출력 합산보다 더 정교한 정보 복원을 가능하게 한다.

세그먼테이션 전략에 따른 Trade-off 분석 결과, 고정 크기 세그먼트는 성능 면에서 유리하고 로그 스케일 세그먼트는 연산 효율 면에서 유리함을 확인했다. 로그 스케일은 O(L log L)의 복잡도를 가지지만, 아주 오래된 과거 정보에 대한 해상도가 낮아지는 특성이 있어 태스크 성격에 따른 선택이 필요하다.

로그 스케일 세그먼테이션을 사용할 경우 오래된 과거 토큰에 대한 해상도가 낮아져 정밀한 회상이 필요한 작업에서 성능이 저하될 수 있다. 또한 캐시된 메모리를 읽어오는 과정에서 추가적인 메모리 대역폭 사용이 발생하며, 이는 하드웨어 구현 시 고려해야 할 요소이다.