계층적 메모리를 활용한 사전 학습: 장기 꼬리 지식과 일반 지식의 분리

초록

현대 언어 모델의 인상적인 성능 향상은 현재 파라미터 스케일링에 의존하고 있다: 더 큰 모델은 더 많은 세계 지식을 저장하고 더 나은 추론 능력을 보인다. 그러나 모든 세계 지식을 파라미터로 압축하는 것은 불필요하며, 각 프롬프트마다 사용되는 지식은 일부에 불과하고, 제한된 추론 시간 메모리와 컴퓨팅 자원을 가진 엣지 디바이스에서는 실용적이지 않다. 우리는 이러한 단점을 메모리 증강 아키텍처와 기존 하드웨어 패러다임에 맞춘 사전 학습 전략으로 해결한다. 우리는 세계 지식을 인코딩한 대규모 계층적 파라미터 메모리 뱅크에 접근하는 소형 언어 모델을 소개한다. 사전 학습과 추론 과정에서 우리는 작은, 문맥에 의존적인 메모리 블록을 가져와 모델에 추가한다. 우리의 사전 학습은 긴 꼬리 형태의 세계 지식을 메모리 파라미터에 저장하는 방법을 배우는 반면, 소형 언어 모델은 일반적인 지식과 일반적인 추론 능력을 포착하는 앵커 역할을 한다. 조 단위 토큰 규모의 실험을 통해 우리는 상당한 성능 향상을 보여준다: 160M 파라미터 모델이 4.6B 메모리 뱅크에서 가져온 18M 파라미터 메모리로 증강되었을 때, 2배 이상의 파라미터를 가진 일반 모델과 비슷한 성능을 얻는다. 광범위한 실험을 통해 우리는 트랜스포머에서 파라미터 메모리의 최적 유형과 크기를 연구하고, 이를 21B 파라미터 이상으로 스케일링한다. 우리는 제안된 계층적 피드포워드 메모리가 사전 학습 중에 추가되든 사후에 추가되든 트랜스포머 아키텍처 전반에서 견고하게 작동함을 발견한다.

English

The impressive performance gains of modern language models currently rely on scaling parameters: larger models store more world knowledge and reason better. Yet compressing all world knowledge into parameters is unnecessary, as only a fraction is used per prompt, and impractical for edge devices with limited inference-time memory and compute. We address this shortcoming by a memory-augmented architecture and a pretraining strategy aligned with existing hardware paradigms. We introduce small language models that access large hierarchical parametric memory banks encoding world knowledge. During pretraining and inference, we fetch a small, context-dependent memory block and add it to the model. Our pretraining learns to store long-tail world knowledge in the memory parameters, while the small language model acts as an anchor capturing common knowledge and general reasoning abilities. Through trillion-token-scale experiments, we show significant gains: a 160M-parameters model augmented with an 18M-parameters memory fetched from a 4.6B memory bank obtains comparable performance to a regular model with more than 2x the parameters. Through extensive experiments, we study the optimal type and size of parametric memories in transformers, scaling them to over 21B parameters. We find that our proposed hierarchical feed-forward memories work robustly across transformer architectures, whether added during pretraining or post-hoc.

계층적 메모리를 활용한 사전 학습: 장기 꼬리 지식과 일반 지식의 분리

Pretraining with hierarchical memories: separating long-tail and common knowledge

초록

Support