[논문 리뷰] Memorizing Transformers

Yuhuai Wu, Markus N. Rabe외 2명|arXiv (Cornell University)|2022. 03. 16.

Topic Modeling인용 수 39

한 줄 요약

이 논문은 디코더-전용 트랜스포머에 approximate kNN을 통해 비미분 가능 외부 메모리를 대규모로 추가하여 과거의 입력을 기억하고 검색함으로써 다양한 긴 맥락 데이터셋에서 언어 모델링을 향상시킨다.

ABSTRACT

Language models typically need to be trained or finetuned in order to acquire new knowledge, which involves updating their weights. We instead envision language models that can simply read and memorize new data at inference time, thus acquiring new knowledge immediately. In this work, we extend language models with the ability to memorize the internal representations of past inputs. We demonstrate that an approximate kNN lookup into a non-differentiable memory of recent (key, value) pairs improves language modeling across various benchmarks and tasks, including generic webtext (C4), math papers (arXiv), books (PG-19), code (Github), as well as formal theorems (Isabelle). We show that the performance steadily improves when we increase the size of memory up to 262K tokens. On benchmarks including code and mathematics, we find that the model is capable of making use of newly defined functions and theorems during test time.

연구 동기 및 목표

• 가중치를 업데이트하기보다 추론 시 과거 입력을 기억하도록 하여 언어 모델이 새로운 지식을 습득하도록 확장하는 것을 동기화한다.
• 외부 저장소에서 정확한 메모리를 검색하는 간단하고 확장 가능한 kNN-증강 어텐션 메커니즘을 제안한다.
• 메모리 크기가 여러 긴 맥락 도메인과 모델 규모에 걸쳐 perplexity를 향상시킨다는 것을 입증한다.

제안 방법

• 이전 단계에서 모델이 생성한 키/값으로 업데이트되는 (key, value) 페어의 비미분 가능 외부 메모리를 추가한다.
• 하나의 어텐션 layer에서 각 쿼리에 대해 메모리에 대해 근사 kNN 검색을 수행하고, 학습된 게이트를 통해 메모리 어텐션과 로컬 밀집(attention)을 결합한다.
• 계속되는 맥락을 위한 Transformer-XL 스타일 캐싱과 로컬 맥락을 유지하기 위한 슬라이딩 인과 마스크를 사용한다.
• 훈련 단계와 메모리 내용 사이의 분포 이동으로 인한 구식화를 완화하기 위해 키와 쿼리를 정규화한다.
• 검색 품질과 계산 효율성의 균형을 맞추기 위해 근사 kNN 구현을 사용하고, 수만 개의 토큰에 달하는 메모리 크기를 가능하게 한다.

실험 결과

연구 질문

• RQ1 kNN을 통해 액세스되는 대형 외부 메모리가 긴 형식의 텍스트와 코드 데이터셋에서 perplexity를 향상시키는가?
• RQ2 메모리 크기가 성능에 어떤 영향을 미치며 수익 체감점이 있는가?
• RQ3 사전 학습된 모델이 미세 조정을 통해 외부 메모리를 효과적으로 활용할 수 있는가, 아니면 처음부터 학습이 필요한가?
• RQ4 모델이 메모리에서 검색할 때 어떤 패턴을 보이는가(예: 함수 이름, 정의, 보조정리 등)?

주요 결과

• 외부 메모리는 긴 형식의 데이터셋과 구조에서 일관되게 perplexity를 향상시킨다(예: 긴 웹 텍스트, 책, 코드, 형식적 증명).
• 메모리 크기를 늘리면 매우 큰 메모리 규모까지 지속적으로 성능 향상이 나타나며, 코드 및 수학 작업에서 특히 이점이 크다.
• 게이팅 메커니즘은 헤드를 외부 메모리 선호로 편향시켜, 미분 가능 메모리 기울기가 없어도 효과적으로 장거리 검색을 가능하게 한다.
• 모델은 메모리를 활용해 정의, 보조정리, 함수 이름을 기억하고, 형식적 증명과 코드에 대한 검색과 유사한 동작을 입증한다.
• 사전 학습 모델을 메모리를 활용하도록 미세 조정하면 기억된 학습 체제와의 격차를 신속히 줄일 수 있어 실용성을 시사한다.