기술 요약: MapFormers

문제 제기

현재 트랜스포머 기반 AI 시스템은 분포 내 (in-distribution) 작업에서 놀라운 성능을 보임에도 불구하고, 분포 외 (out-of-distribution, OOD) 시나리오에 직면할 때 상당한 취약성을 드러냅니다. 이러한 시스템은 종종 인간과 동물이 쉽게 처리하는 새로운 상황에 일반화하지 못합니다. 저자들은 이러한 격차를 생물학적 뇌가 **인지 지도 (cognitive maps)**를 구축할 수 있게 해주는 구조적 편향의 부재로 돌립니다. 인지 지도란 환경의 추상적이고 불변인 구조를 그 구체적이고 가변적인 내용으로부터 분리해 내는 내부 모델을 의미합니다.

RNN 에서 행동 의존적 행렬을 사용하는 Tolman-Eichenbaum Machine (TEM) 과 같은 기존 인지 지도 모델링 시도들은 두 가지 치명적인 한계를 겪고 있습니다:

1. 확장성 부족: 위치를 업데이트하기 위해 순차적인 행렬 곱셈이 필요하여 병렬 처리를 방해합니다.
2. 사전에 지정된 행동: 행동 집합 (예: 기본 방향) 이 데이터에서 학습되는 것이 아니라 하드코딩되어야 합니다.

또한, 표준 위치 인코딩 (RoPE 등) 은 정적이거나, 알려지지 않은 환경에서의 경로 통합에 필요한 동적 행동 조합을 지원하지 않는 휴리스틱에 의존합니다.

방법론

이 논문은 지도 학습 없이 다음 토큰 예측을 통해 인지 지도를 학습하고 경로 통합을 수행하도록 설계된 트랜스포머 아키텍처 클래스인 MapFormers를 소개합니다. 핵심 혁신은 리 군 (Lie-group) 이론에서 파생된 입력 의존적 행렬을 사용하여 위치 인코딩을 업데이트하는 데 있습니다.

이론적 프레임워크

• 리 군 형식주의: 저자들은 행동을 리 군의 요소 (특히 $SO(2)$와 같은 회전과 같은 콤팩트 군) 로 모델링합니다. 가능한 모든 행동에 대해 별도의 행렬을 학습하는 대신, 모델은 **리 대수 생성자 (Lie-algebra generators)**를 학습합니다.
• 지수 사상을 통한 경로 통합: 행동은 리 대수에서 무한소 생성자로 표현됩니다. 경로 통합 (행동 시퀀스 조합) 은 대수에서 이러한 생성자를 합산 (선형 누적) 한 후, 리 군에서 유한 변환을 얻기 위해 단일 지수 사상을 적용하여 수행됩니다.
  • 가환 지도 (예: 2D 그리드 탐색) 의 경우, 단일 행렬 지수 함수를 따르는 누적 합산을 통해 경로 통합을 병렬로 계산할 수 있습니다.
  • 이 접근법은 유도 편향을 행동의 수에서 대수의 구조적 복잡성으로 줄입니다.

아키텍처 변형

이 논문은 절대 및 상대 위치 인코딩을 고유한 메모리 모델과 통합하는 MapFormers 의 두 가지 변형을 제안합니다:

1. MapEM (에피소드적 메모리):

   • 인지 지도를 표현하기 위해 절대 위치 임베딩을 사용합니다.
   • 내용 (관측치) 과 구조 (위치) 를 별도의 신경 집단으로 분리합니다.
   • 내용과 위치의 결합 ($Q_G = \text{vec}(Q^T \cdot P)$) 에 어텐션을 계산하여, 위치 인코딩이 내용과 독립적으로 업데이트되는 어텐션 마스크 역할을 효과적으로 수행합니다.
   • 특정 상태의 견고한 회상이 필요한 작업에 적합합니다.

2. MapWM (작업 기억):

   • RoPE 와 유사하지만 동적인 상대 위치 인코딩을 사용합니다.
   • 입력 의존적 회전 행렬 ($R_{\theta PI}$) 을 통해 키와 쿼리를 서로 다른 부분 공간으로 회전시킴으로써 구조를 암시적으로 표현합니다.
   • 회전 각도 $\theta$는 입력 토큰의 저차원 투영에서 유도되므로, 모델은 어떤 토큰이 위치를 업데이트하는 행동으로 작용하고 어떤 토큰이 관측치로 위치를 변경하지 않는지를 학습할 수 있습니다.

입력 의존적 메커니즘

두 아키텍처 모두 각 토큰에 대한 "통합 지속 시간"($\Delta_t$) 을 결정하기 위해 2 단계 투영을 사용합니다:

1. 내부 투영: 입력 토큰을 저차원 행동 표현으로 매핑합니다.
2. 외부 투영: 이 표현을 블록별 통합 지속 시간으로 투영합니다.
   최종 회전 각도는 $\theta_t = \omega \odot \Delta_t$이며, 여기서 $\omega$는 학습 가능한 각속도 벡터입니다. 경로 통합은 이러한 각도의 누적 합 ($\theta_{0 \to t} = \text{cumsum}(\theta_t)$) 으로 계산됩니다.

주요 기여

1. MapFormers 아키텍처: 리 대수 생성자에 기반한 입력 의존적 위치 인코딩을 사용하여 구조를 내용으로부터 분리함으로써 인지 지도를 학습하는 새로운 트랜스포머 변형.
2. 병렬 경로 통합: 가환 리 군 (예: $SO(2)^k$) 을 활용함으로써 이전 RNN 기반 인지 지도 모델의 순차적 병목 현상을 극복하고 병렬 경로 통합을 달성합니다.
3. 메모리 모델의 통합: 절대/상대 위치 인코딩과 에피소드/작업 기억 사이의 이론적 연결을 제공하여, 둘 다 동일한 리 이론 원칙에서 유도될 수 있음을 보여줍니다.
4. 비가환 지도로의 확장성: 여러 생성자를 사용하여 비가환 문제 (예: 가족 관계도) 로 프레임워크가 확장되었으나, 이는 순차적 경로 통합이 필요합니다.

실험 결과

저자들은 MapFormers 를 다양한 인지 능력을 대상으로 하는 세 가지 공식 작업에서 평가했습니다:

1. 선택적 복사 (게이팅):

   • 작업: 방해 토큰을 무시하면서 시퀀스를 복사합니다.
   • 결과: MapFormers(EM 및 WM 모두) 는 완벽한 OOD 일반화를 달성한 반면, 표준 베이스라인 (RoPE, TAPE) 은 실패했습니다. 이는 동적 게이팅 메커니즘을 학습할 수 있음을 보여줍니다.

2. 그리드 탐색 (1D 및 2D):

   • 작업: 교차하는 행동과 무작위 관측치 시퀀스에서 관측치를 예측하며, 모델은 잠재적 위상 구조를 추론하고 위치를 경로 통합해야 합니다.
   • 결과: MapFormers 는 길이와 밀도가 다른 OOD 시퀀스에서 거의 완벽한 정확도를 달성한 반면, 베이스라인 (CoPE 및 PathAtt 포함) 은 일반화에 실패했습니다.
   • 신경 분석: 모델은 관측치 토큰을 영 회전 (항등) 으로, 행동 토큰을 특정 회전 방향으로 투영하도록 학습하여 내용과 구조를 효과적으로 분리했습니다.

3. Dyck-2 (중첩 계층 구조):

   • 작업: 스택 추적이 필요한 괄호 시퀀스의 유효한 연장을 예측합니다.
   • 결과: MapFormers 는 단일 레이어만 사용하여 훈련에서 본 것보다 더 길고 깊은 시퀀스에 완벽하게 일반화했습니다. 베이스라인은 여러 레이어가 필요했음에도 불구하고 OOD 데이터에서 실패했습니다.
   • 메커니즘: 모델은 서로 다른 괄호 유형을 위한 직교 차원과 여는/닫는 괄호를 위한 반대 방향을 사용하여 위치 인코딩에 스택을 직접 임베딩했습니다.

4. 자연어 확장:

   • 작업: OpenWebText 에 대한 사전 학습.
   • 결과: 12 레이어 MapWM 모델은 RoPE 베이스라인 ($18.79$대$19.14$) 보다 일관된 퍼플렉시티 개선을 보였으며 더 나은 길이 외삽 능력을 입증했습니다. 그러나 BLIMP 구문 벤치마크에서는 MapWM 이 RoPE 보다 유의미한 향상을 보이지 않아, 가환 경로 통합만으로는 모든 구문 모델링에 충분하지 않을 수 있음을 시사합니다.

중요성과 주장

이 논문은 MapFormers 가 휴리스틱에 의존하는 대신 경로 통합을 통해 잠재적 구조를 포착함으로써 체계적인 OOD 일반화를 가능하게 하는 중요한 구조적 편향을 제공한다고 주장합니다.

• 견고성: 모델들은 추상적 관계를 내용으로부터 분리함으로써 관찰되지 않은 시퀀스 길이, 깊이 및 환경 밀도에 대한 견고한 일반화를 가능하게 함을 입증했습니다.
• 확장성: 가환 지도에 대한 병렬 계산을 사용함으로써 이러한 원칙이 대규모 실세계 도메인으로 확장될 수 있음을 자연어 결과를 통해 입증했습니다.
• 생물학적 타당성: 이 접근법은 "무엇"과 "어디"를 분리하는 생물학적 인지 지도와 현대 딥러닝 사이의 격차를 해소하여, 환경의 기하학적 구조를 이해함으로써 새로운 문제를 추론할 수 있는 AI 시스템으로 가는 길을 제시합니다.

저자들은 비가환 경로 통합은 병렬화할 수 없으며, 현재 가환 제약이 언어의 특정 복잡한 구문 구조 모델링을 제한할 수 있음을 지적하며 제한 사항에 대해 겸손한 태도를 유지했습니다. 그들은 향후 연구가 더 복잡한 관계 추론을 위해 하이브리드 아키텍처나 비가환 확장을 탐구할 수 있다고 제안합니다.