새 MLP 구성법이 트랜스포머의 사실 저장 한계를 낮춘다

모델 차원을 고정하고 저장할 사실 수를 늘릴 때, 제안한 데이터 의존형 구성은 기존 방식보다 최대 63배 나은 저장 용량을 보였다. 이 결과는 트랜스포머 블록 안에서 사실을 저장하고 불러오는 방법을 다룬다.

로베르토 가르시아(Roberto Garcia), 제리 류(Jerry Liu), 로니 정킨스(Ronny Junkins), 크리스 레이(Chris Ré) 등 연구팀은 트랜스포머의 다층 퍼셉트론(MLP)이 연상 기억처럼 작동한다는 설명과, 사실을 저장하는 MLP의 닫힌 형태 구성을 제시했다. 이 구성은 경사하강법으로 학습하지 않고도 사실을 저장하도록 설계됐다. 연구팀은 트랜스포머 블록에서 사실을 재학습 없이 편집하는 데도 이 구성을 활용할 수 있다고 제시한다.

이 구성은 무엇인가?

헤브식 기억은 저장된 키와 값의 쌍을 더해 기억하고, 입력 키가 저장된 키와 얼마나 비슷한지에 따라 값을 꺼내는 방식이다. 이를 열쇠와 열쇠 꾸러미를 대조해 맞는 값을 찾는 일에 빗대면 이해하기 쉽다. 제안한 MLP는 입력과 저장된 키를 특징 공간에서 비교해 유사도에 따라 값을 합산하는 헤브식 기억이다.

연구팀은 사실을 키와 값의 대응 관계로 정의한다. 예를 들어 ‘프랑스’라는 키에 ‘파리’라는 값이 연결되는 식이다. MLP가 사실을 저장했다는 것은 키를 입력했을 때 정답 값의 임베딩(모델이 단어 등을 숫자 벡터로 나타낸 표현)과의 내적 점수가 다른 값 임베딩과의 점수보다 높은 상태를 뜻한다. 언어 모델도 은닉 벡터와 다음 토큰의 출력 임베딩 사이 내적 점수로 후보를 평가한다.

사실을 어떻게 저장하나?

구성은 키 임베딩과 해당 값 임베딩을 이용한다. 먼저 가우스 분포에서 무작위 행렬 A와 G를 뽑고, 다음과 같은 게이트형 MLP를 만든다.

\(\phi(x)=(Ax)\odot(Gx)\)

\(B=\frac{1}{F}\sum_i v_{f(i)}\phi(k_i)^\top\)

\(\mathrm{MLP}(x)=B\phi(x)\)

여기서 F는 저장할 사실 수이고, kᵢ는 키, v_f(i)는 그 키에 연결된 값이다. ϕ는 입력을 특징 공간으로 바꾸며, ⊙는 원소별 곱을 뜻한다. 이 구성에는 경사하강법이 필요하지 않다.

고전적인 헤브식 기억은 키와 값의 외적을 더해 행렬을 만든다. 이 MLP도 같은 방식으로 작동하지만, 키와 질의의 유사도를 원래 공간이 아니라 MLP의 특징 공간에서 잰다. 저장된 키와 같은 질의가 들어오면 일치하는 키의 유사도는 커지고 다른 키의 유사도는 작아져, 가중합이 올바른 값에 가까워진다.

저장 용량은 어느 정도인가?

연구팀은 키와 값 임베딩이 등방성, 즉 방향에 따라 분포가 크게 치우치지 않는 경우 제안한 구성의 매개변수 수 W가 정보 이론상 최적 규모인 \(W=\Theta(md)=\Theta(F\log F)\)에 이른다고 제시한다. 여기서 m과 d는 구성의 차원을 나타내며, F는 저장하는 사실 수다. 비등방성 임베딩에서도 같은 규모를 얻는 유사한 결과를 논문에 제시했지만, 임베딩의 기하학적 특성에 따른 불이익이 붙는다.

연구팀의 비교에 따르면 제안한 헤브식 MLP는 등방성 임베딩에서 기존 신경 접선 커널(NTK) 구성보다 10³~10⁴배 나은 용량을 보였고, 경사하강법으로 학습한 MLP와는 6~10배 이내의 차이를 보였다. 또 공분산으로 특징을 백색화한 커널을 쓰는 변형과, 최소제곱 문제를 풀어 사실을 저장하는 데이터 의존형 변형도 제시했다.

트랜스포머 블록에서는 어텐션이 MLP에 정확한 키 대신 잡음이 섞인 질의를 전달할 수 있다. 연구팀의 분석에 따르면 질의와 정확한 키 임베딩 사이의 L2 오차가 \(\sqrt{d/(F\log F)}\)보다 작으면 블록이 올바른 사실을 회수할 수 있다. 이 조건에서 블록도 \(W=\Theta(md)=\Theta(F\log F)\) 규모의 저장 용량을 갖는다. 데이터 의존형 구성은 모델 차원을 고정하고 사실 수를 늘리는 비교에서 기존 구성보다 최대 63배 나았으며, 경사하강법으로 학습한 MLP와의 차이는 3배 이내였다.

연구팀은 이 구성이 트랜스포머 블록에서 정보 이론상 최적인 비율로 사실을 저장하는 첫 MLP 구성이라고 주장한다. 논문은 사전 학습된 언어 모델의 MLP에서 사실을 직접 추출하거나 편집하는 방법, 어텐션의 KV 캐시 같은 시퀀스 혼합기와 여러 층으로 된 트랜스포머의 기억을 설명하는 문제도 앞으로 다룰 과제로 꼽는다.

논문은 실제 언어 모델에서 얻은 비등방성 임베딩에 대한 유사한 용량 결과를 다루지만, 이 글에는 그 조건별 수치가 제시되지 않았다.