알렉스 왕은 트랜스포머를 기억의 관점에서 설명한다

알렉스 왕(Alex Wang)은 선형층, 다층 퍼셉트론(MLP), 셀프 어텐션을 서로 다른 메모리 저장 방식으로 해석하고, 저장된 정보를 얼마나 잘 되찾는지에 따라 용량을 비교한다. 무작위 키와 값으로 계산한 추정치에서 키 차원이 64일 때 선형 메모리는 약 64개, 제곱 ReLU MLP는 약 2,700개, 셀프 어텐션은 약 120만 개의 메모리를 저장할 수 있다.

글에는 관련 코드 저장소와 유사도 간섭을 살펴보는 대화형 실험이 포함돼 있다.

메모리 용량은 회상 간섭으로 추정한다

이 글에서 메모리는 키와 값의 쌍이다. 키는 유사도를 기준으로 정보를 찾는 데 쓰고, 값은 되찾으려는 정보다. 저장한 키로 질의하면 해당 값뿐 아니라 다른 키와의 유사도에 따른 값도 함께 회상될 수 있다. 왕은 이때 목표 값에 비해 다른 메모리가 얼마나 방해하는지를 ‘간섭 대 신호 비율’로 삼는다. 간섭이 신호만큼 커지는 지점을 메모리 용량으로 정의한다.

가장 단순한 방식은 각 키·값 쌍의 외적을 가중치에 더하는 선형 저장소다. 무작위 단위 길이 키와 비슷한 크기의 무작위 값을 가정하면, 이 방식의 용량은 키 차원 d_k 정도다. 저장된 메모리를 한꺼번에 맞추도록 선형 회귀로 가중치를 정하면 간섭을 줄여 용량을 약 2d_k까지 높일 수 있다.

MLP는 키와 질의의 유사도에 비선형 함수를 적용해 유사도가 낮은 키의 영향을 억제한다. ReLU는 음의 유사도를 0으로 만들며, 제곱 ReLU는 작은 유사도의 영향도 더 강하게 줄인다. 왕의 추정에서 d_k = 64일 때 ReLU MLP의 용량은 약 128개, 제곱 ReLU MLP는 약 2,700개다. 세제곱 ReLU는 약 3만 5,000개로 추정한다.

셀프 어텐션은 유사도에 지수 함수를 적용하고 소프트맥스로 정규화한다. 이 방식은 유사도가 높은 키에 가중치를 집중시킨다. d_k = 64이고 온도 계수 β = √d_k = 8인 조건에서 용량은 약 120만 개로 계산된다. 이 수치는 무작위 키와 값에 대한 이론적 추정이며, 학습된 모델이 실제로 보유한 사실의 수를 직접 측정한 결과는 아니다.

셀프 어텐션에서는 키와 값이 학습된 가중치가 아니라 문맥에 있는 토큰에서 계산된다. 따라서 메모리 수는 문맥 길이에 따라 늘어난다. 문맥이 길어질수록 소프트맥스의 온도 계수 β를 키우는 방법도 소개한다. 글은 Qwen의 LogN-Scaling, Su가 제안한 방식, Scalable-Softmax와 YaRN을 비교하며, 실험 조건에 따라 간섭이 더 낮아질 수 있다고 설명한다. 다만 β를 지나치게 높이면 소프트맥스가 포화돼 다른 키로 전달되는 기울기가 사라지고 학습이 나빠질 수 있어 무한정 키울 수는 없다.

왕은 이 분석이 현대 신경망의 용량을 바라보는 한 가지 관점일 뿐이라고 강조한다. 추정치는 한 층이 무작위 벡터를 얼마나 회상하는지에 초점을 맞추며, 학습 데이터가 실제 키와 값에 미치는 영향은 반영하지 않는다. 따라서 표의 수치를 모델이 실제로 기억하는 사실의 개수로 해석해서는 안 된다.