[논문 리뷰] Memory Transformer

Mikhail Burtsev, Yuri Kuratov외 2명|arXiv (Cornell University)|2020. 06. 20.

Analog and Mixed-Signal Circuit Design인용 수 6

한 줄 요약

이 논문은 전역 및 국소 표현을 저장하기 위해 학습 가능한 [mem] 토큰을 도입함으로써 시퀀스 모델링을 향상시키는 메모리 보강형 Transformer 아키텍처인 MemTransformer을 제안한다. 모델은 기계 번역 및 언어 모델링 작업에서 성능 향상을 보이며, 주의 시각화 결과는 효과적인 메모리 읽기/쓰기 및 블록 단위 처리를 보여주며, 메모리 크기가 모델 정확도와 정비례함을 확인한다.

ABSTRACT

Transformer-based models have achieved state-of-the-art results in many natural language processing tasks. The self-attention architecture allows transformer to combine information from all elements of a sequence into context-aware representations. However, information about the context is stored mostly in the same element-wise representations. This might limit the processing of properties related to the sequence as a whole more difficult. Adding trainable memory to selectively store local as well as global representations of a sequence is a promising direction to improve the Transformer model. Memory-augmented neural networks (MANNs) extend traditional neural architectures with general-purpose memory for representations. MANNs have demonstrated the capability to learn simple algorithms like Copy or Reverse and can be successfully trained via backpropagation on diverse tasks from question answering to language modeling outperforming RNNs and LSTMs of comparable complexity. In this work, we propose and study few extensions of the Transformer baseline (1) by adding memory tokens to store non-local representations, (2) creating memory bottleneck for the global information, (3) controlling memory update with dedicated layer. We evaluate these memory augmented Transformers and demonstrate that presence of memory positively correlates with the model performance for machine translation and language modelling tasks. Augmentation of pre-trained masked language model with memory tokens shows mixed results for tasks from GLUE benchmark. Visualization of attention patterns over the memory suggest that it improves the model's ability to process a global context.

연구 동기 및 목표

• 표준 Transformer가 토큰 표현에 분산 저장된 맥락 정보에 접근하는 데에 한계가 있다는 점을 해결하기 위해.
• 학습 가능한 일반 목적 메모리 추가가 기계 번역 및 언어 모델링과 같은 시퀀스 모델링 작업에서 성능 향상에 기여하는지 조사하기 위해.
• 메모리 보강 아키텍처가 복사, 요약, 블록 단위 처리와 같은 복잡한 메모리 연산을 어떻게 지원할 수 있는지 탐색하기 위해.
• 추론 중 메모리 크기가 변할 때 메모리 컨트롤러의 강건성과 일반화 능력을 평가하기 위해.
• 메모리 보강 기법이 BERT 및 Transformer-XL과 같은 사전 훈련된 모델에 하류 GLUE 작업에서 전이 가능한지 평가하기 위해.

제안 방법

• 입력 시퀀스의 시작 부분에 [mem] 토큰을 도입하여 전역 및 국소 표현을 위한 통합 메모리 저장소로 사용한다.
• 세 가지 변형을 제안한다: MemTransformer(기본 메모리 보강), MemCtrl([mem] 토큰 갱신을 전용 Transformer 레이어로 제어), MemBottleneck([시퀀스 간 주의 없이] 메모리 간 주의만 허용).
• 표준 목표를 사용해 역전파를 통해 훈련함으로써 메모리 쓰기, 읽기 및 메모리 내 연산의 엔드 투 엔드 학습을 가능하게 한다.
• 자기 주의 메커니즘을 활용해 입력 토큰과 [mem] 토큰 간의 주의(읽기/쓰기) 및 [mem] 토큰 간 주의(메모리 내 처리)를 가능하게 한다.
• 주의 패턴의 시각화를 통해 메모리 행동을 분석하고, 대각선 주의(증폭), 블록 복사, 수직 주의(읽기)와 같은 연산을 식별한다.
• 사전 훈련된 모델(예: BERT, Transformer-XL)에 메모리 보강을 적용하고, GLUE 벤치마크 작업에서 미세 조정한다.

실험 결과

연구 질문

• RQ1 학습 가능한 [mem] 토큰의 추가가 기계 번역 및 언어 모델링 작업에서 표준 Transformer의 성능 향상에 기여하는가?
• RQ2 MemTransformer의 주의 패턴에서 쓰기, 읽기, 복사, 요약과 같은 학습된 메모리 연산의 증거가 드러나는가?
• RQ3 메모리 크기가 추론 중 모델 성능과 강건성에 미치는 영향은 어떠한가, 특히 메모리 크기가 변경될 경우?
• RQ4 메모리 보강이 하류 GLUE 작업에서 BERT 및 Transformer-XL과 같은 사전 훈련된 모델의 성능 향상에 기여하는가?
• RQ5 메모리 컨트롤러는 다양한 메모리 크기 간에 일반화 가능한가, 그리고 추론 중 메모리가 감소할 때 점진적으로 성능 저하가 발생하는가?

주요 결과

• [mem] 토큰의 추가로 인해 Transformer 아키텍처는 기계 번역 및 언어 모델링 작업에서 일관된 성능 향상을 보였다.
• 모델 성능은 메모리 크기와 정비례함을 확인하여, 더 큰 메모리 용량이 표현 학습에 기여함을 시사한다.
• 주의 패턴의 시각화 결과는 메모리 읽기/쓰기 연산 및 블록 복사, 대각선 증폭과 같은 메모리 내 처리가 존재함을 확인한다.
• 메모리 손실 테스트 결과, 사전 훈련된 MemTransformer 모델의 성능은 메모리의 존재에 크게 의존함을 확인하여 기능적 중요성을 입증한다.
• 메모리 컨트롤러는 강건성을 보이며, 추론 중 메모리 크기를 줄일 경우 성능 저하가 점진적으로 발생함을 확인하여 일반화 능력을 시사한다.
• 확장 후 메모리 보강 모델의 미세 조정이 성능 향상에 기여함을 확인하여, 지속적 훈련과 결합했을 때 메모리 보강 기법이 효과적임을 입증한다.