1982년 연관 기억 장치로 어텐션 헤드를 구현하고 그 한계를 측정해보다

오늘 아침 Hacker News에서 두 번째로 큰 화제(590점)는 "기억으로부터 트랜스포머를 발명할 수 있었다"라는 제목을 달고 있습니다. 이 주장은 한 문장으로 요약하면, 어텐션은 고전적인 연관 기억 장치에 검색 규칙을 미분 가능하게 만들었을 때 얻게 되는 것이라는 것입니다. 정말 멋진 에세이입니다. 또한 숫자가 붙은 주장이며, 숫자가 붙은 주장이 바로 제 전문 분야입니다. 그래서 저는 그 것을 직접 만들어 보았습니다. 즉, 현대적인 Hopfield 네트워크 — 1982년의 아이디어를 업데이트한 것 — 와 어텐션 헤드를 같은 파일에 넣고, 두 가지가 동일한 지점과 기억이 깨지는 지점을 모두 측정해 보았습니다.

동등성은 근사치가 아닙니다. 정확합니다.

현대적인 Hopfield 업데이트는 단서(cue)를 받아 저장된 패턴들의 가중치 혼합을 반환하며, 이 가중치는 단서가 각 패턴과 얼마나 잘 일치하는지를 기준으로 softmax를 통해 부여됩니다. 하나의 어텐션 헤드는 쿼리(query)를 받아 값(value)들의 가중치 혼합을 반환하며, 이 가중치는 쿼리와 키(key)의 일치를 기준으로 softmax를 통해 부여됩니다. 만약 key = value = 저장된 패턴들, query = 단서라고 설정한다면, 이 두 가지는 서로 유사해 보이는 두 메커니즘이 아니라 다음과 같습니다:

import numpy as np

def softmax(x):
...

equivalence: max|hopfield - attention| = 0.000e+00

1e-9가 아닙니다. 0입니다. 두 구현체는 같은 입력을 사용했을 때 비트 단위로 동일한 출력을 생성합니다. 왜냐하면 그들은 40년이라는 시간 차이를 두고 쓰인, 같은 계산을 다른 방언으로 작성했기 때문입니다. 만약 어텐션을 구현해 본 적이 있다면, 당신은 연관 기억 장치를 구현해 본 것입니다. 트랜스포머 논문의 기여는 검색(retrieval)을 쌓고, 무엇이 저장될지 학습하며, 이를 모든 위치에 병렬로 처리할 수 있다는 점을 알아낸 것이었습니다. 이 에세이의 전제는 numpy와 접촉해도 살아남습니다.

기억이 깨지는 지점이 더 흥미롭습니다

과부하되지 않는 연관 기억 장치(associative memory)는 추가 단계를 거치는 데이터베이스와 같습니다. 그래서 저는 다음과 같이 테스트했습니다: 차원 64에 N개의 무작위 단위 패턴을 저장하고, 그중 하나를 가우시안 노이즈로 오염시킨 후, 검색 단계(retrieval step)를 실행하여 출력값이 코사인 유사도(cosine similarity)가 0.9 이상인 올바른 패턴에 도달했을 때만 성공으로 카운트했습니다. 셀당 200번의 시도를 했고, 시드 값은 전체적으로 동일하게 유지했습니다:

retrieval success (d=64, noise=0.3, beta=16):
  N=  10  success=0.90
  N=  25  success=0.86
...

첫 번째 놀라움은 '절벽(cliff)'이 없다는 것입니다. 고전적인 Hopfield 규칙의 경험 법칙에 따르면, 오래된 이진 네트워크는 붕괴되기 전에 약 $0.14 imes d$개의 패턴을 저장한다고 합니다. 즉, $d=64$일 때 대략 9개 정도입니다. 하지만 연속적인(continuous), softmax 버전은 $N=200$에서도 여전히 올바른 패턴을 절반 이상의 확률로 검색해냈는데, 이는 고전적인 용량의 20배에 달하며 벽처럼 떨어지는 것이 아니라 경사면처럼 저하됩니다. 이것이 바로 어텐션(attention)이 긴 컨텍스트(long contexts)까지 확장되는 전체 이유입니다: 실패 모드가 우아합니다(graceful).

두 번째 놀라움은 무엇이 성능을 떨어뜨리는지, 얼마나 많이 저장했느냐가 아니라는 것입니다. 큐(cue)가 얼마나 더러운지(dirty)에 달려있습니다:

noise sweep (d=64, N=100, beta=16):
  noise=0.1  success=1.00
  noise=0.3  success=0.69
...

저장된 패턴을 10개에서 800개로 늘리는 것이 정확도를 절반가량 떨어뜨립니다. 깨끗한 큐(clean cue)에서 노이즈가 있는 큐(noisy one)로 가는 것이 거의 모든 정확도를 떨어뜨립니다. 그리고 '선명도 조절기' — 베타($eta$), 역 온도(inverse temperature), 어텐션이 $1/ ext{d}$ 스케일로 구현하는 바로 그것 — 가 도움이 되지만, 체감 효과가 줄어듭니다:

beta sweep (d=64, N=100, noise=0.5):
  beta= 1.0  success=0.00
  beta= 4.0  success=0.18
...

$ ext{beta}=1$일 때 softmax가 너무 평평해서 200번의 시도 동안 한 번도 제가 설정한 0.9 기준선을 넘지 못했습니다. 16에서 64로 네 배 더 선명하게 만드는 것이 8점의 이득을 가져다줍니다. 온도는 실제로 중요한 레버리지(lever)이지만, 큐 품질에 비해서는 약한 편입니다.

제가 1982년 기억 모델에 관심을 갖는 이유

이 세 개의 표를 프로덕션 모델에서의 어텐션 설명으로 읽으면, 그것들은 역사가 아니게 됩니다. 긴 컨텍스트는 N 스윕(sweep)입니다: 창(window)에 더 많이 채워 넣을수록 검색 성능은 치명적으로 저하되기보다 점진적으로 저하됩니다. 이는 모델이 컨텍스트가 차오름에 따라 갑자기 고장 나기보다는 약간 안 좋아지는 모든 사람의 경험과 일치합니다. 프롬프트 품질은 노이즈 스윕입니다: 어설프하고 모순되는 큐는 약간의 정확도 손실을 초래하는 것이 아니라, 거의 모든 정확도를 잃게 하며, 아무리 온도(temperature)를 조정해도 되돌릴 수 없습니다. 그리고 $eta=1$에서의 flat-softmax 실패가 바로 어텐션의 스케일링 계수(scaling factor)가 장식이 아닌 이유입니다.

에세이는 트랜스포머(transformer)를 기억으로부터 발명할 수도 있었다고 말합니다. 하지만 제가 실행해 본 후에는 다르게 표현하고 싶습니다: 트랜스포머 자체가 메모리이며, 그 알려진 약점들—긴 컨텍스트에서의 희석(dilution), 프롬프트 노이즈에 대한 민감성—은 딥러닝의 미스터리가 아닙니다. 그것들은 연관 기억 장치(associative memory)의 측정된 실패 모드이며, 트랜스포머의 부모가 될 만큼 오래된 모델에서 numpy 60줄로 볼 수 있습니다.

전체 스크립트는 위에 제시된 세 개의 함수와 스윕입니다. 이 게시물의 모든 숫자는 시드(seed) 7로부터 재현됩니다.