입력 의존 위치 임베딩으로 학습하는 인지 지도 Transformer

초록

MapFormer는 입력에 따라 변하는 위치 임베딩을 이용해 Transformer 내부에 인지 지도를 형성한다. 절대·상대 위치 인코딩을 각각 에피소드 기억(EM)과 작업 기억(WM) 모델에 적용해 행동 토큰은 위치를, 관찰 토큰은 내용을 업데이트하도록 학습한다. 실험에서는 2D 네비게이션과 선택적 복사 과제에서 기존 모델을 크게 앞서며, 시퀀스 길이가 늘어나도 거의 완벽한 OOD 일반화를 보였다.

상세 분석

본 논문은 인간·동물의 인지 지도 능력을 인공 신경망에 도입하기 위해, Transformer의 위치 인코딩을 입력‑의존적인 행렬로 교체하는 새로운 아키텍처인 MapFormer를 제안한다. 핵심 아이디어는 구조‑내용 분리를 통해 “어디(where)”와 “무엇(what)”을 별도로 학습하도록 하는데, 이는 기존 정적 위치 임베딩이 제공하지 못하는 구조적 편향을 제공한다.
두 변형, Map EM과 Map WM은 각각 절대 위치 임베딩과 상대적 회전형 Positional Encoding(RoPE)을 활용한다. Map EM에서는 행동 토큰이 생성하는 회전 각도 θ를 저차원 투사 ω·Δ로부터 얻어, 누적 합(cumsum) 연산을 통해 경로 통합(path integration)을 병렬적으로 수행한다. 이렇게 얻은 입력‑의존 회전 행렬 Rθ는 토큰의 위치를 인지 지도상의 좌표로 매핑한다. 반면 Map WM은 RoPE와 동일한 구조를 유지하되, θ를 학습 가능한 파라미터로 두어 실제 회전 행렬을 동적으로 생성한다. 두 모델 모두 키·쿼리(Q,K)를 Rθ로 회전시켜 구조 정보를 별도 어텐션 흐름에서 처리하고, 내용 정보는 일반 어텐션(A X)에서 다룬다.
수학적으로는 SO(2)와 같은 1‑차원 콤팩트 리 군을 이용해 행동‑의존 행렬을 Lie algebra(so(2)) 수준에서 합산하고, 최종적으로 exponential map을 통해 그룹 원소로 복원한다. 이는 기존 RNN 기반 모델이 순차적으로 행렬을 곱해야 하는 비효율성을 제거하고, 병렬 연산이 가능한 Transformer와 자연스럽게 결합한다. 또한, 블록‑대각 회전 행렬을 여러 스케일에 걸쳐 학습함으로써 큰 그리드에서도 래핑 현상을 방지한다.
실험에서는 (1) 선택적 복사(Selective Copy) 과제와 (2) 2‑D 네비게이션(ced‑navigation) 과제를 사용해 OOD 일반화 능력을 평가한다. 선택적 복사에서는 행동 토큰이 아닌 관찰 토큰을 무시하고 정확히 복사해야 하는데, 기존 RoPE 기반 Transformer는 방해 토큰을 구분하지 못해 성능이 급락한다. MapFormer는 입력‑의존 위치 행렬 덕분에 행동·관찰을 명확히 구분해 99 % 이상의 정확도를 유지한다. 2‑D 네비게이션에서는 훈련 시보다 긴 시퀀스와 더 큰 그리드로 확장했을 때도 거의 완벽한 경로 예측을 수행한다. 이는 모델이 구조를 추상화해 내용과 독립적으로 학습했기 때문이다.
비교 대상인 TEM‑t, 기존 RoPE Transformer, 그리고 최신 선택적 State‑Space Model(SSM)과도 비교했으며, MapFormer가 길이 일반화와 구조 학습 측면에서 일관되게 우수함을 보였다. 이론적 분석에서는 대각 행렬만을 사용하는 현재 SSM이 비가환 그룹을 표현하지 못해 인지 지도 학습에 한계가 있음을 지적한다. 반면 MapFormer는 입력‑의존 회전 행렬을 통해 보다 풍부한 군 구조를 모델링한다.
한계점으로는 현재 실험이 2‑D 격자와 제한된 관계(상하좌우)로만 검증됐으며, 복잡한 비가환 구조(예: 3‑D 회전)나 대규모 실제 환경에 대한 확장성 검증이 부족하다. 또한, 입력‑의존 행렬을 학습하는 과정이 추가 파라미터와 연산 비용을 유발하지만, 병렬화 덕분에 전체 학습 시간은 기존 Transformer와 비슷한 수준이다. 향후 연구에서는 비가환 리 군(SO(3) 등) 적용, 멀티‑모달 인지 지도 학습, 그리고 뇌 신경생리학적 데이터와의 정량적 비교가 기대된다.