[논문 리뷰] Recurrent Memory Transformer

Aydar Bulatov, Yuri Kuratov외 1명|arXiv (Cornell University)|2022. 07. 14.

Topic Modeling인용 수 25

한 줄 요약

이 논문은 장기 시퀀스 모델링을 향상시키기 위해 전용 메모리 토큰을 입력에 통합하고 세그먼트 간 복귀를 가능하게 하는 메모리 증강형 세그먼트 수준의 순환 Transformer인 Recurrent Memory Transformer(RMT)을 제안한다. RMT는 복사, 뒤집기, 연관 검색과 같은 장기 컨텍스트 작업에서 Transformer-XL을 능가하지만, 최대 10배의 메모리 크기 감소로도 동등한 언어 모델링 성능를 달성하여, 변형 없이도 메모리 효율성과 장기 의존성 학습 능력이 뛰어나다는 것을 입증한다.

ABSTRACT

Transformer-based models show their effectiveness across multiple domains and tasks. The self-attention allows to combine information from all sequence elements into context-aware representations. However, global and local information has to be stored mostly in the same element-wise representations. Moreover, the length of an input sequence is limited by quadratic computational complexity of self-attention. In this work, we propose and study a memory-augmented segment-level recurrent Transformer (RMT). Memory allows to store and process local and global information as well as to pass information between segments of the long sequence with the help of recurrence. We implement a memory mechanism with no changes to Transformer model by adding special memory tokens to the input or output sequence. Then the model is trained to control both memory operations and sequence representations processing. Results of experiments show that RMT performs on par with the Transformer-XL on language modeling for smaller memory sizes and outperforms it for tasks that require longer sequence processing. We show that adding memory tokens to Tr-XL is able to improve its performance. This makes Recurrent Memory Transformer a promising architecture for applications that require learning of long-term dependencies and general purpose in memory processing, such as algorithmic tasks and reasoning.

연구 동기 및 목표

• 표준 Transformer의 이차적 자기주의 복잡도로 인한 장기 시퀀스 처리 한계를 해결하기 위해.
• Transformer 아키텍처를 수정하지 않고도 장기 의존성과 전반적 컨텍스트를 개선하기 위해.
• 전용 메모리 토큰과 세그먼트 수준의 복귀가 장기 컨텍스트 작업에서의 메모리 효율성과 성능 향상에 얼마나 효과적인지 탐색하기 위해.
• 메모리 증강형 Transformer가 Transformer-XL과 같은 최첨단 모델보다 우월하거나 동등한 성능을 내는지 평가하기 위해.
• RMT의 어텐션 메커니즘이 메모리 사용 방식, 특히 해석 가능한 읽기/쓰기 패턴을 통해 어떻게 작동하는지 분석하기 위해.

제안 방법

• 모델은 입력 시퀀스에 전용 메모리 토큰을 추가하여 전역 및 국소 정보를 저장하는 데 전용 저장소로 활용한다. 이를 통해 입력 토큰 표현과 분리된다.
• 시퀀스는 세그먼트 단위로 처리되며, 이전 세그먼트의 메모리 히든 상태가 현재 세그먼트로 전달되어 복귀와 장기 컨텍스트 모델링이 가능해진다.
• 메모리 연산(읽기, 쓰기, 재작성)은 어텐션 메커니즘을 통해 엔드 투 엔드로 학습되며, 훈련 중 메모리 상태를 통해 기울기가 역전파된다.
• 모델의 아키텍처 변경 없이도 구현 가능하며, 수정은 입력 및 출력 시퀀스 수준에서만 적용된다.
• RMT는 메모리 연산과 시퀀스 표현 처리를 동시에 제어하여 압축되고 컨텍스트에 민감한 표현을 가능하게 한다.
• 이 방법은 기존 모델과 호환 가능하다. RMT는 Transformer-XL의 캐시와 조합하여 성능 향상을 더욱 높일 수 있다.

실험 결과

연구 질문

• RQ1 메모리 증강형 순환 Transformer 아키텍처가 표준 Transformer와 Transformer-XL을 초월해 장기 시퀀스 모델링을 향상시킬 수 있는가?
• RQ2 전용 메모리 토큰의 사용이 입력 토큰과 전역 컨텍스트 간의 정보 혼합을 줄여 알고리즘 작업 성능 향상에 기여하는가?
• RQ3 RMT가 메모리 크기를 크게 줄였음에도 불구하고, Transformer-XL과 동등한 언어 모델링 성능를 달성할 수 있는가?
• RQ4 RMT의 어텐션 패턴은 Transformer-XL과 어떻게 다를까? 특히 메모리 접근 및 정보 흐름 측면에서.
• RQ5 RMT는 Transformer-XL의 캐싱 메커니즘과 효과적으로 조합될 수 있는가? 이를 통해 성능 향상이 가능할까?

주요 결과

• 시퀀스 길이가 임계값을 초과할 경우, 특히 2개 이상의 세그먼트가 존재할 때, RMT는 복사, 뒤집기, 연관 검색과 같은 장기 컨텍스트 알고리즘 작업에서 Transformer-XL을 능가한다.
• 언어 모델링 작업에서는 RMT가 Transformer-XL과 동등한 성능를 달성하지만, 최대 10배 적은 메모리 크기를 사용하여 뛰어난 메모리 효율성을 입증한다.
• 어텐션 맵의 시각화 결과 RMT가 전체 세그먼트를 메모리 토큰에 압축하는 방식으로 학습하는 것으로 나타났으며, Transformer-XL에서 관찰되는 표현 혼합 현상은 피한다.
• 4개의 세그먼트가 있는 작업에서, 메모리 크기가 6인 Transformer-XL은 뒤집기 작업에서 평균 정확도 0.8에 머물렀지만, RMT는 메모리 저장소를 격리함으로써 정확도 1.0을 달성했다.
• RMT의 읽기 메모리에서 쓰기 메모리로의 재작성 기능 덕분에 최근 세그먼트 정보의 장기 유지가 가능해져, 다중 세그먼트 시퀀스에서 성능 향상이 이루어졌다.
• RMT를 Transformer-XL의 캐시와 조합하면 언어 모델링 성능이 향상되며, 이는 두 접근 방식 간 상호보완적인 강점이 있음을 시사한다.