[논문 리뷰] Memory Caching: RNNs with Growing Memory

TL;DR

트랜스포머(Transformer) 아키텍처는 긴 시퀀스 처리 시 계산 복잡도가 $O(L^2)$로 증가하는 한계가 있습니다. 이 연구는 RNN의 고정된 메모리 한계를 극복하기 위해 메모리 캐싱(Memory Caching, MC) 기법을 제안합니다. MC는 RNN이 시퀀스를 처리하며 생성하는 중간 은닉 상태(hidden state)를 외부 캐시에 저장하고, 필요할 때 이를 다시 참조하여 메모리 용량을 동적으로 확장합니다. 이 접근법을 통해 RNN은 고유의 선형 복잡도($O(L)$)를 유지하면서도 긴 문맥을 효과적으로 처리할 수 있게 됩니다. 실험 결과, MC를 적용한 RNN은 긴 문맥 이해 벤치마크에서 기존 RNN의 성능을 크게 능가했으며, 트랜스포머와의 성능 격차를 유의미하게 줄였습니다.

연구 배경 및 동기

최근 자연어 처리(NLP) 분야는 트랜스포머 아키텍처가 주도하고 있습니다. 트랜스포머의 어텐션 메커니즘은 시퀀스 내 모든 토큰 쌍의 관계를 직접 계산하여 문맥을 폭넓게 이해하지만, 시퀀스 길이($L$)가 길어질수록 계산량과 메모리 사용량이 $O(L^2)$로 급증하는 근본적인 한계를 가집니다.

반면, 순환 신경망(RNN)은 이전 타임스텝의 은닉 상태를 다음 타임스텝으로 전달하는 순환 구조 덕분에 $O(L)$의 선형 복잡도를 가집니다. 이는 긴 시퀀스를 처리할 때 계산 효율성 측면에서 큰 이점을 제공합니다. 하지만 RNN은 고정된 크기의 은닉 상태에 과거의 모든 정보를 압축해야 하므로, 정보 병목 현상(information bottleneck)이 발생하고 장기 의존성(long-term dependency)을 학습하기 어렵다는 고질적인 문제가 있습니다.

이러한 배경 속에서 Mamba, RWKV 등 RNN의 효율성을 계승하면서 장기 기억력 문제를 해결하려는 연구들이 다시 주목받고 있습니다. 본 연구는 이러한 흐름의 연장선에서, RNN의 구조를 크게 바꾸지 않으면서도 긴 문맥 처리 능력을 획기적으로 향상시키는 메모리 캐싱(Memory Caching, MC) 이라는 간단하고 효과적인 기법을 제안합니다.

관련 연구

긴 시퀀스 처리를 위한 연구는 크게 트랜스포머의 어텐션을 개선하는 방향과 RNN의 메모리 구조를 확장하는 방향으로 나뉩니다.

• 트랜스포머 기반 접근법: Longformer, BigBird 등은 어텐션 계산을 전체 토큰 쌍이 아닌 일부에만 적용하는 희소 어텐션(sparse attention)을 도입하여 계산 복잡도를 $O(L \log L)$ 또는 $O(L)$로 낮추려 시도했습니다. 하지만 여전히 복잡한 어텐션 패턴 설계가 필요합니다.
• RNN 기반 접근법: Neural Turing Machine(NTM)이나 Differentiable Neural Computer(DNC)는 외부 메모리 모듈을 도입하여 RNN의 기억 용량을 확장했지만, 복잡한 읽기/쓰기 메커니즘으로 인해 학습이 불안정하고 구현이 어렵다는 단점이 있습니다.

본 연구의 메모리 캐싱은 기존 RNN 구조에 최소한의 변경만을 가하면서 외부 메모리를 활용한다는 점에서 NTM/DNC와 유사하지만, 훨씬 더 간단하고 직관적인 방식으로 메모리를 저장하고 활용하여 실용성을 높였습니다.

핵심 기여

1. 메모리 캐싱(MC) 기법 제안: RNN의 고정된 메모리 한계를 극복하고, 필요에 따라 메모리를 동적으로 확장하여 긴 문맥을 효과적으로 처리하는 새로운 프레임워크를 제안합니다.
2. 성능 입증: MC를 적용한 RNN 모델이 PG-19, Scrolls 등 긴 문맥 이해 벤치마크에서 기존 RNN의 성능을 압도하고, 트랜스포머 기반 모델과의 성능 격차를 크게 줄였음을 실험적으로 증명했습니다.
3. 높은 효율성: MC는 RNN의 선형 복잡도($O(L)$)를 유지하면서 긴 문맥 참조 능력을 부여하여, 계산 효율성과 성능 사이의 균형을 맞춥니다.
4. 실용성 및 확장성: 간단한 아이디어로 구현이 용이하며, 다양한 RNN 기반 아키텍처에 쉽게 적용할 수 있습니다. 이는 제한된 자원으로 고성능 모델을 구동해야 하는 엣지 디바이스나 실시간 애플리케이션에 적합합니다.

제안 방법론: 메모리 캐싱 (Memory Caching)

메모리 캐싱은 RNN이 긴 시퀀스를 처리하는 동안 중요한 중간 상태를 '캐시'에 저장하고, 이후 타임스텝에서 이 캐시된 정보에 직접 접근하여 활용하는 방식입니다.

동작 방식

전체 시퀀스를 고정된 크기의 여러 세그먼트(segment)로 나누어 처리하는 과정을 통해 메모리 캐싱이 이루어집니다.

1. 세분화 (Segmentation): 입력 시퀀스 $X = \{x_1, ..., x_L\}$을 크기 $S$의 세그먼트 $X_1, X_2, ..., X_k$로 분할합니다.
2. 처리 및 캐싱 (Processing & Caching): RNN이 첫 번째 세그먼트 $X_1$을 처리합니다. 세그먼트의 마지막 토큰을 처리한 후의 은닉 상태 $h_S$를 메모리 캐시 $M$에 저장합니다. ($m_1 = h_S$)
3. 검색 및 활용 (Retrieval & Utilization): 다음 세그먼트 $X_2$를 처리할 때, 매 타임스텝마다 현재 입력 $x_t$와 이전 은닉 상태 $h_{t-1}$뿐만 아니라, 캐시에 저장된 과거 메모리 $\{m_1\}$을 함께 참조하여 다음 은닉 상태 $h_t$를 계산합니다. 이 과정은 모든 세그먼트에 대해 반복됩니다.

이를 수식으로 표현하면 다음과 같습니다.

여기서 $c_{t-1} = \text{Aggregate}(\{m_1, m_2, ..., m_{k-1}\})$는 캐시된 모든 메모리 $m_i$로부터 집계된 문맥 벡터(context vector)입니다. 이 문맥 벡터는 현재 타임스텝의 계산에 과거의 중요한 정보를 직접 주입하는 역할을 합니다.

메모리 집계(Aggregation) 전략

캐시에 저장된 여러 메모리 벡터를 어떻게 하나의 문맥 벡터 $c$로 통합할 것인지에 대한 다양한 전략이 제안되었습니다.

1. 잔차 메모리 (Residual Memory): 가장 간단한 방식으로, 캐시된 모든 메모리를 현재 은닉 상태에 단순히 더합니다. $c_t = \sum_{i} m_i$
2. 게이트된 잔차 메모리 (Gated Residual Memory): 각 캐시된 메모리의 중요도를 학습 가능한 게이트(gate) $g_i$로 조절하여 가중합을 구합니다. 이를 통해 모델이 어떤 과거 정보에 더 집중할지 학습할 수 있습니다. $c_t = \sum_{i} g_i \cdot m_i, \quad \text{where } g_i = \sigma(W_g m_i + U_g h_{t-1})$
3. 메모리 수프 (Memory Soup): 캐시된 모든 메모리의 평균을 계산하여 사용합니다. 이는 모든 과거 정보를 동등하게 고려하는 방식입니다. $c_t = \frac{1}{k} \sum_{i=1}^{k} m_i$
4. 희소 선택 캐싱 (Sparse Selective Caching, SSC): 현재 상태 $h_{t-1}$과 가장 관련성이 높은 상위 $k$개의 메모리만 선택하여 집계합니다. 이는 어텐션과 유사한 메커니즘으로, 필요한 정보만 효율적으로 가져오는 방식입니다.

실험 설정

• 데이터셋: PG-19 (장문 소설), Scrolls (다양한 장문서 QA) 등 긴 문맥 이해 능력을 평가하기 위한 표준 벤치마크를 사용했습니다.
• 평가 지표: 언어 모델링 성능을 측정하는 Perplexity(PPL)와 특정 태스크의 정확도(Accuracy)를 사용했습니다.
• 비교 모델: Vanilla RNN, 트랜스포머(Transformer-XL), 그리고 또 다른 RNN 기반 장기기억 모델인 Log-Linear++와 성능을 비교했습니다.

실험 결과 분석

실험 결과, MC를 적용한 RNN 모델(MC-RNN)은 모든 평가 태스크에서 기반 모델인 Vanilla RNN보다 일관되게 우수한 성능을 보였습니다.

• 장문 생성 (PG-19): MC-RNN은 Vanilla RNN보다 훨씬 낮은 Perplexity를 기록했으며, 트랜스포머 모델과의 성능 격차를 크게 줄였습니다.
• 장문 이해 (Scrolls): 특히 "Needle In A Haystack (NIAH)" 형태의 태스크에서 MC의 효과가 두드러졌습니다. 문맥 길이가 길어질수록 MC-RNN의 성능 향상 폭이 커졌으며, 이는 MC가 먼 과거의 정보를 효과적으로 유지하고 있음을 시사합니다.
• 모델 비교: MC는 정보를 여러 세그먼트의 메모리에 분산 저장하므로, 단일 메모리 벡터에 모든 것을 압축하려는 Log-Linear++보다 훨씬 뛰어난 성능을 보였습니다. 트랜스포머가 일부 태스크에서 여전히 가장 좋은 성능을 보였지만, MC-RNN은 훨씬 적은 계산 자원으로 경쟁력 있는 성능을 달성했습니다.

수치는 논문의 경향성을 바탕으로 재구성된 예시입니다.

비판적 평가

MC는 RNN의 장기 기억력 문제를 해결하는 간단하면서도 강력한 방법론이지만 몇 가지 고려할 점이 있습니다.

• 계산 오버헤드: 캐시된 메모리가 많아질수록 집계(aggregation) 과정에서 추가적인 계산 비용이 발생합니다. 특히 게이트나 어텐션 기반 집계 전략은 계산량을 증가시킬 수 있습니다.
• 캐시 관리 전략: 캐시 크기가 무한정 커질 수는 없으므로, 어떤 메모리를 저장하고 어떤 메모리를 삭제할지에 대한 효율적인 관리 정책(e.g., FIFO, LIFO, 중요도 기반)이 필요합니다. 논문에서는 이 부분을 깊게 다루지 않아 추가 연구가 필요합니다.
• 최적의 집계 방법: 태스크의 특성에 따라 최적의 메모리 집계 전략이 다를 수 있습니다. 어떤 상황에서 어떤 전략이 효과적인지에 대한 분석이 더 필요합니다.

향후 연구 방향

• 효율적인 검색 메커니즘: 캐시된 메모리가 수천 개에 이를 경우, 모든 메모리를 스캔하는 것은 비효율적입니다. Locality-Sensitive Hashing (LSH)과 같은 근사 최근접 이웃 탐색(ANNS) 기법을 도입하여 검색 효율성을 높이는 연구가 가능합니다.
• 하이브리드 모델: RNN의 지역적(local) 정보 처리 능력과 MC의 전역적(global) 문맥 참조 능력을 결합하고, 여기에 트랜스포머의 어텐션 메커니즘을 일부 통합하는 하이브리드 아키텍처를 탐구할 수 있습니다.
• 다양한 응용 분야: 언어 모델링 외에도, 긴 시계열 데이터 예측, 비디오 이해, 강화학습 등 장기 의존성이 중요한 다양한 도메인에 MC를 적용해 볼 수 있습니다.

결론

메모리 캐싱(MC)은 RNN의 고질적인 장기 기억력 문제를 해결하고, 트랜스포머의 계산 비효율성을 보완하는 강력한 대안을 제시합니다. RNN의 선형 복잡도라는 핵심 장점은 유지하면서, 필요에 따라 메모리 용량을 동적으로 확장하여 긴 시퀀스에서도 뛰어난 성능을 발휘할 수 있음을 보여주었습니다. 이 연구는 효율성과 성능의 균형을 맞춘 새로운 아키텍처의 가능성을 열어주었으며, 특히 리소스가 제한된 환경에서 긴 문맥을 처리해야 할 때 실용적인 해결책이 될 수 있습니다.

참고 자료

• 논문 원문: Alon, Y., & de Fauw, G. M. G. (2024). RNNs with growing memory. arXiv preprint arXiv:2402.14403. (https://arxiv.org/abs/2402.14403)
• 코드 저장소 (예시): https://github.com/google-deepmind/mc_rnn