[논문 리뷰] Large Memory Layers with Product Keys

Guillaume Lample, Alexandre Sablayrolles외 3명|HAL (Le Centre pour la Communication Scientifique Directe)|2019. 07. 10.

Topic Modeling인용 수 50

한 줄 요약

제품 키(product keys)를 기반으로 한 확장 가능한 메모리 강화 계층을 도입하여, 정확한 최근접 이웃 탐색과 함께 매우 큰 메모리 용량을 가능하게 하고, 대규모 언어 모델링을 위한 트랜스포머에 통합되어 더 깊은 baselines보다 더 빠른 추론 속도와 더 나은 정확도를 달성합니다.

ABSTRACT

This paper introduces a structured memory which can be easily integrated into a neural network. The memory is very large by design and significantly increases the capacity of the architecture, by up to a billion parameters with a negligible computational overhead. Its design and access pattern is based on product keys, which enable fast and exact nearest neighbor search. The ability to increase the number of parameters while keeping the same computational budget lets the overall system strike a better trade-off between prediction accuracy and computation efficiency both at training and test time. This memory layer allows us to tackle very large scale language modeling tasks. In our experiments we consider a dataset with up to 30 billion words, and we plug our memory layer in a state-of-the-art transformer-based architecture. In particular, we found that a memory augmented model with only 12 layers outperforms a baseline transformer model with 24 layers, while being twice faster at inference time. We release our code for reproducibility purposes.

연구 동기 및 목표

• 모델 용량을 증대시키되 계산량은 비례적으로 증가하지 않도록 크고 학습 가능한 메모리 계층을 도입한다.
• product keys를 이용한 확장 가능한 키-값 메모리에서 정확한 최근접 이웃 탐색을 제공한다.
• 대규모 언어 모델링 작업에서 트랜스포머 내에서 메모리 계층의 효과를 입증한다.
• 메모리-강화 모델이 더 깊은 baselines보다 더 낮거나 비교 가능한 계산량으로 더 나은 성능을 발휘할 수 있음을 보여준다.
• 메모리 크기, 네트워크 내 위치, 정규화가 성능 및 사용에 어떤 영향을 주는지 분석한다.

제안 방법

• 쿼리 네트워크, 두 서브-키 프로덕트 키 공간, 값 조회 테이블로 메모리를 정의한다.
• 두 코드북의 외적에 의해 형성된 프로덕트 키를 사용하여 큰 키 집합을 암묵적으로 정의한다.
• 쿼리와의 내적 유사도로 상위-k 키를 선택한 다음 해당 값의 가중 합계를 계산한다(희소 읽기).
• 각 헤드가 자체 쿼리 네트워크와 서브-키를 가지되 값을 공유하는 다중-헤드 메모리 어텐션을 사용한다.
• 키 커버리지와 사용을 개선하기 위해 쿼리 네트워크의 희소 메모리 업데이트와 배치 정규화를 사용한다.
• 선택된 트랜스포머 레이어에서 FFN을 메모리 계층으로 바꿔 잔여 연결(x → x + PKM(x))를 더하여 사용한다.

실험 결과

연구 질문

• RQ1 프로덕트-키 메모리(PKM)가 정확한 탐색을 유지하면서도 과소비 오버헤드 없이 확장 가능한 용량을 제공하는가?
• RQ2 メ모리 크기, 헤드 수, k-NN, 메모리 위치가 대형 트랜스포머의 perplexity, 메모리 사용량, 속도에 어떤 영향을 미치는가?
• RQ3 프로덕트 키 메모리와 플랫 키 메모리의 성능, 사용성, 효율성에 차이가 있는가?
• RQ4 대규모 언어 모델링 데이터셋에서 메모리 사용과 모델 성능을 최적화하는 학습 구성(BatchNorm, 희소 업데이트)은 무엇인가?

주요 결과

• 12-layer 메모리와 프로덕트 키가 같은 차원의 24-layer 트랜스포머보다 CC-News에서 추론 시간이 대략 절반 수준으로 더 나은 성능을 보인다.
• 메모리 크기를 늘리면 perplexity가 일관되게 낮아지며, 배치 정규화가 키 사용을 개선할 때 메모리 크기 1M 슬롯에서 주목할 만한 이득이 발생한다.
• 프로덕트 키는 같은 크기에서 플랫 키보다 더 높은 메모리 사용량, 더 낮은 perplexity, 더 빠른 추론 속도를 제공한다.
• 메모리 위치는 중요하며 중간 트랜스포머 레이어(약 4–5)에 메모리를 삽입하면 가장 좋은 성능을 얻을 수 있다.
• 512^2 총 슬롯과 각 헤드에 32개의 키를 가진 4개의 메모리 헤드가 속도와 perplexity 간의 유리한 절충점을 제공한다; 프로덕트 키는 플랫 키 대비 파라미터 수를 줄이고 추론 시 더 큰 규모로 효율적으로 확장한다.
• 쿼리 네트워크의 배치 정규화가 대형 메모리의 키 사용과 perplexity를 크게 향상시킨다.