메모리포머: 완전연결층을 없애고 메모리 조회로 연산량을 대폭 줄인 트랜스포머

초록

MemoryFormer는 전통적인 완전연결(FC) 레이어를 메모리 내 해시 테이블과 LSH 기반 조회 연산으로 대체하여, 다중 헤드 어텐션을 제외한 대부분의 연산을 제거한다. 입력 임베딩을 비트화·해시한 뒤 관련 벡터들을 가중합함으로써 FC 연산을 근사하고, 역전파를 통해 테이블을 학습한다. 실험 결과, 동일한 모델 규모에서 FLOPs를 80% 이상 절감하면서도 기존 트랜스포머와 비슷한 성능을 달성한다.

상세 분석

MemoryFormer의 핵심 아이디어는 “연산을 메모리 조회로 대체한다”는 점이다. 기존 트랜스포머는 토큰당 d 차원의 임베딩을 h 차원으로 변환하는 FC 레이어를 매번 행렬 곱으로 수행한다. 이는 전체 연산량 O(s·d·h) 를 차지해, 시퀀스 길이 s 가 비교적 짧아도 FLOPs 의 대부분을 차지한다. 논문은 이를 완전히 없애고, 대신 입력 벡터 x 를 d 차원에서 K 개의 서브벡터 z_k (각 차원 τ = d/K) 로 분할한다. 각 서브벡터는 부호(sign) 함수를 통해 2^τ 개의 버킷 중 하나에 매핑되는 LSH 해시 코드를 생성하고, 해당 버킷에 저장된 사전 학습 가능한 벡터