Memory Layers at Scale

Memory+는 키-값 메모리를 LLM을 위한 실용적인 희소 용량 구성 요소로 만드는데, 특히 단순히 더 많은 FLOPs를 쓰는 것보다 사실 기억이 더 중요할 때 의미가 있다.

이 논문은 Transformer를 위한 학습 가능한 키-값 메모리 계층을 확장하여, 모델이 FLOPs를 비례적으로 늘리지 않고도 대량의 희소 용량을 추가할 수 있게 한다. 제안된 Memory+ 설계는 곱-키 조회, 입력 의존적 게이팅, 여러 계층에 걸친 메모리 공유, 메모리-값 샤딩, 맞춤형 CUDA 커널을 사용하여 수십억 파라미터 메모리를 학습 가능하고 병렬화 가능하게 만든다. 134M에서 8B 밀집 파라미터에 이르는 모델과 최대 1T 토큰의 사전학습 전반에서, 메모리 증강 모델은 연산량을 맞춘 밀집 기준선과 MoE/PEER 계열 희소 기준선을 능가하며, 특히 NaturalQuestions와 TriviaQA 같은 사실 QA에서 큰 향상을 보인다. 결과는 한계도 보여준다. 메모리 계층은 밀집 순방향 연산을 대체하기보다는 보완하며, 밀집 용량을 너무 많이 대체하면 성능이 떨어진다.

핵심 요약

• 방법 측면에서 각 메모리 계층은 입력 질의에 대해 상위 k개의 학습 가능한 키를 검색하고 해당 값들을 결합하여, 시퀀스 활성값에 대한 어텐션과는 구별되는 저렴한 학습형 조회 테이블 역할을 한다.
• 엔지니어링 측면에서 곱-키는 검색 비용을 줄이고, 값들은 GPU 전반에 샤딩되며, 맞춤형 CUDA 커널은 H100에 근접한 메모리 대역폭에 도달하여, 보고된 스케일링 연구에서 최대 128B 파라미터 메모리를 가능하게 한다.
• 결과적으로 Memory+는 두 배 이상의 연산 예산을 가진 밀집 모델을 능가하고, 연산량과 파라미터 수를 맞췄을 때 MoE 및 PEER를 능가하며, 사실 기억 벤치마크에서 가장 뚜렷한 향상을 보인다.
• 시사점은 최적의 설계가 밀집 추론 용량과 희소 메모리 용량 사이의 균형을 유지하며, 이는 미래 LLM이 연산 집약적 변환과 효율적인 지식 저장을 분리할 수 있음을 시사한다는 것이다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.