Hebbian 메모리 기반 MLP가 학습 없이 Transformer에 사실을 저장함

TL;DR

• Hebbian 메모리 관점에서 만든 닫힌 형태의 다층 퍼셉트론(MLP)은 경사 하강법(gradient descent) 없이 사실 집합을 저장하며, 등방성 임베딩에서 정보 이론적 최적 용량인 Θ(F log F)를 달성함.
• 구성은 무작위 가우스 행렬로 게이트형 MLP의 특징을 만들고, 저장할 키와 값의 외적을 합산하는 방식임.
• MLP는 특징 공간에서 질의와 저장된 키의 유사도를 계산하고, 그 유사도에 따라 저장된 값의 가중합을 반환하는 Hebbian 메모리로 해석됨.
• Transformer 블록은 어텐션이 전달하는 키의 잡음이 √(d/(F log F)) 이하일 때 올바른 사실을 회수하며, 데이터 의존형 구성은 기존 방식보다 최대 63배 높은 용량을 보임.
• 이 접근은 사전 학습된 대형 언어 모델(LLM)의 사실 저장 분석, 저장 지식의 추출·편집, 시퀀스 믹서와 다층 Transformer의 메모리 연구로 확장될 여지가 있음.

사실 저장

• 선행 연구는 LLM이 사실 F개를 저장하는 데 정보 이론적 최적 수준에 가까운 Θ(F log F)개 매개변수를 사용한다고 실험적으로 관찰함. 목표는 이 용량을 MLP에서 구성적으로 달성하고 Transformer가 해당 MLP를 활용할 수 있음을 보이는 것임.
• 사실 집합은 키 임베딩 K={k₁,…,kF}, 값 임베딩 V={v₁,…,vF}, 그리고 키 인덱스를 값 인덱스에 연결하는 대응 f:[F]→[F]로 정의됨. 예를 들어 수도 정보 집합에서는 “France”의 키 임베딩이 “Paris”의 값 임베딩에 대응함.
• MLP는 키 kᵢ를 입력했을 때 출력이 올바른 값 v_f(i)와 가장 큰 내적을 갖는 경우 해당 사실 집합을 저장한 것으로 정의됨. 이는 언어 모델이 은닉 벡터와 출력 임베딩의 내적으로 다음 토큰 점수를 계산하는 방식과 일치함.

간단한 MLP 구성

• 키 kᵢ와 해당 값 v_f(i)가 주어지면, 가우스 분포에서 무작위 행렬 A,G∈Rᵐˣᵈ를 표본 추출하고 게이트형 MLP를 구성함. 특징은 ϕ(x)=(Ax)⊙(Gx)이며, 출력은 MLP(x)=Bϕ(x)임.
• 출력 행렬 B는 각 값 벡터와 해당 키의 특징 벡터 외적을 합산한 뒤 사실 수 F로 나눈 값임: B=(1/F)∑ᵢ v_f(i)ϕ(kᵢ)ᵀ. 이 구성에는 경사 하강법이 필요하지 않음.
• 고전적인 Hebbian 메모리는 키와 값의 외적 합 W=∑ᵢvᵢkᵢᵀ를 저장하고, 질의와 키의 유사도를 가중치로 삼아 값을 합산함. 제안된 MLP도 같은 원리지만, 유사도를 원래 입력 공간이 아닌 MLP 특징 공간 ϕ에서 측정함.
• 질의 q에 대한 출력은 저장된 값 v_f(i)와 특징 공간 유사도 ⟨ϕ(kᵢ),ϕ(q)⟩의 가중합이며, 이 유사도는 커널 K(kᵢ,q)로 표현됨. 저장된 키 q=kⱼ를 질의하면 일치하는 항의 유사도는 크고 불일치 항의 유사도는 작아져 올바른 값에 대응하는 결과가 나옴.

MLP 저장 용량

• 키와 값 임베딩이 균일 구면 분포와 같은 등방성 분포일 때, 매개변수 W개를 사용하는 구성은 사실 F개를 W=Θ(md)=Θ(F log F)의 정보 이론적 최적 용량으로 저장함.
• 비등방성 임베딩, 즉 LLM에서 얻는 임베딩에 대해서도 논문은 유사한 결과를 제시함. 저장 용량은 임베딩 기하에 따른 페널티 요인을 제외하면 같은 비율로 증가함.
• 제안된 Hebbian MLP는 등방성 임베딩에서 최적 용량을 달성함. 기존 신경 탄젠트 커널(NTK) 구성보다 10³~10⁴배 높은 용량을 보이며, 경사 하강법으로 학습한 MLP와의 차이는 6~10배 이내임.
• 추가 구성으로 공분산을 백색화한 커널을 사용하는 백색화 변형과, 사실 집합 저장을 위해 최소제곱 문제를 푸는 데이터 의존형 변형을 제시함.

Transformer 블록의 저장 용량

• 독립된 MLP에 사실을 저장하는 것과 Transformer 블록 안에서 사실을 활용하는 것은 다름. 블록 안에서는 어텐션이 정확한 키 kᵢ가 아니라 잡음이 섞인 벡터를 MLP에 전달함.
• 제안된 구성은 잡음이 있는 질의를 처리할 수 있음. 예를 들어 어텐션이 정확한 키 대신 q=k₁+ε를 전달해도 구성된 MLP와 전체 Transformer 블록이 올바른 값 v₁=Emb(Paris)를 복원함.
• 질의 벡터와 조회 대상 사실의 정확한 키 임베딩 사이의 L2 오차로 정의되는 어텐션 잡음이 √(d/(F log F))보다 작으면 Transformer 블록은 올바른 사실을 회수함. 이 조건에서 블록은 MLP와 동일한 정보 이론적 최적 용량 W=Θ(md)=Θ(F log F)를 유지함.
• 어텐션 잡음이 제한된 조건에서 제안된 MLP를 포함한 Transformer 블록은 최적 용량을 달성함. 데이터 의존형 구성은 기존 구성보다 최대 63배 높은 용량을 보임.
• Transformer 블록이 사실을 정보 이론적 최적 비율로 저장할 수 있음을 입증한 첫 MLP 구성임. 제안된 구성은 기존 연구보다 단순하고 Transformer 안에서 더 잘 확장되며, 데이터 의존형 구성은 경사 하강법 학습 MLP와의 차이가 3배 이내이고 사실 수가 늘어날수록 기존 구성보다 최대 63배 높은 용량을 보임.

다음 연구 과제

• 간단한 수학적 대상인 Hebbian 커널 메모리를 통해 용량 보장이 있는 사실 저장 MLP와 Transformer 블록을 구성할 수 있음. 다만 추가 연구 과제가 남아 있음.
• 사전 학습된 LLM의 MLP 내부에 사실이 어떻게 저장되는지 이론으로 설명할 수 있는지, 해당 MLP에 저장된 지식을 직접 추출하거나 편집할 수 있는지가 연구 질문임.
• 어텐션의 키·값(KV) 캐시와 같은 시퀀스 믹서의 매개변수에 사실 저장 이론을 확장할 수 있는지, 다층 모델에서 어텐션과 MLP 계층이 지식 저장에 어떻게 협력하는지가 연구 질문임.