기술 요약: 트랜스포머의 기하학적 사실적 회상

문제 정의

대형 언어 모델 (LLM) 은 방대한 사실적 연관 집합을 저장하고 검색할 수 있는 능력을 갖추고 있습니다. 이 능력을 이해하는 데 있어 핵심적인 질문은 트랜스포머 아키텍처가 이러한 연관성을 어떻게 기억하는지입니다. 기존 이론적 프레임워크는 이를 연상 기억 (associative memory) 문제로 모델링하는 경우가 많으며, 여기서 내부 가중치 행렬 (특히 MLP 또는 어텐션 메커니즘) 은 입력 - 출력 임베딩 쌍을 무작위이며 거의 직교하는 벡터로 저장합니다. 이러한 관점에서 $N$개의 사실을 기억하는 데 필요한 매개변수 수는 $N$에 비례하여 선형적으로 증가합니다.

그러나 경험적 관찰은 다른 메커니즘을 시사합니다. 즉, 트랜스포머는 **기하학적 기억 (geometric memorization)**을 활용하여 학습된 임베딩이 데이터의 관계 구조를 직접 인코딩할 수 있다는 것입니다. 본 논문은 트랜스포머가 무차별적인 연상 저장에 의존하기보다 임베딩 공간을 기하학적으로 구조화함으로써 사실적 연관성을 더 효율적으로 기억할 수 있는지 조사합니다. 저자들은 $N$개의 주체, $R$개의 관계, 그리고 공유된 속성 집합이 포함된 통제된 설정에 초점을 맞추며, 모델은 주체 - 관계 쌍과 연관된 속성을 예측하거나 다중 홉 쿼리에서 이러한 관계를 구성해야 합니다 (예: "$x$의 아내의 어머니는 누구인가?").

방법론

본 논문은 합성 데이터 및 실제 데이터를 통한 경험적 검증을 이론적 구성과 결합합니다.

이론적 프레임워크

저자들은 임베딩을 학습 가능한 매개변수로 간주할 때 트랜스포머가 사실적 회상 과제를 해결할 수 있는 능력을 분석합니다. 그들은 두 가지 영역을 구분합니다:

1. 단일 홉 회상 (Single-Hop Recall): 주체와 관계가 주어졌을 때 속성을 예측합니다.
2. 다중 홉 회상 (Multi-Hop Recall): 중간 출력 없이 $k$개의 관계를 순차적으로 구성하여 직접 예측하는 경우와 사고의 사슬 (Chain-of-Thought, CoT) 을 사용하는 경우.

이론적 분석은 학습 가능한 임베딩을 가진 트랜스포머에 대한 명시적 해를 구성하는 데 의존합니다. 이러한 구성의 핵심은 주체 임베딩 내에서의 선형 중첩 (linear superpositions) 활용입니다. 사실을 별도의 직교 벡터로 저장하는 대신, 주체 임베딩은 해당 주체의 모든 $R$개 가능한 속성을 나타내는 벡터들을 포함하는 블록 벡터로 구성됩니다. 그런 다음 MLP 는 조회 테이블이 아니라 ReLU 게이팅을 통해 관련 속성 블록을 추출하는 **관계 조건부 선택기 (relation-conditioned selector)**로 작용합니다.

저자들은 또한 $R$-정규 엣지 컬러 그래프에 대한 카운팅 논증을 활용하여 $k$-홉 회상에 필요한 매개변수 수와 임베딩 차원에 대한 정보 이론적 하한을 유도합니다.

경험적 설정

경사 하강법이 이러한 이론적 구조를 발견하는지 확인하기 위해 저자들은 합성 데이터에 대한 광범위한 실험을 수행합니다:

• 데이터셋: 공유된 속성 풀로 매핑되는 $N=4096$개의 주체와 $R \in \{2, \dots, 16\}$개의 관계.
• 아키텍처: 균일하거나 학습된 어텐션과 2 층 MLP 를 가진 단일 층 트랜스포머.
• 통제: 학습 가능한 임베딩을 가진 모델과 이전 연구의 고정 임베딩 영역을 모방하기 위해 무작위 초기화 시 **동결 (frozen)**된 입력 엔티티 임베딩을 가진 모델을 비교합니다.
• 분석: 훈련 후 분석으로는 중첩 구조를 확인하기 위한 선형 판독, MLP 선택성을 테스트하기 위한 인과적 개입, 그리고 학습된 메커니즘의 일반성을 테스트하기 위한 "동결 및 교체 (freeze-and-swap)" 전이 실험이 포함됩니다.
• 실제 세계 탐지: 저자들은 사전 훈련된 LLM(예: Qwen, Llama, Phi) 을 탐지하여 주체 은닉 상태가 이론적 예측과 유사한 선형 관계 인코딩을 보이는지 확인합니다.

주요 기여

1. 공유 속성을 통한 로그arithmic 기억

본 논문은 MLP 가 작고 관계 조건부일 경우, 단일 층 트랜스포머가 임베딩 차원 $d = O(R \log N)$을 사용하여 $N \times R$개의 사실을 기억할 수 있음을 증명합니다.

• 메커니즘: 주체 임베딩은 해당 $R$개의 연관 속성 벡터들의 선형 중첩을 인코딩합니다.
• MLP 의 역할: MLP 는 매핑을 가중치에 저장하는 것이 아니라 ReLU 게이팅을 통해 쿼리된 관계에 해당하는 특정 속성 벡터를 추출하는 선택기로 작용합니다.
• 의의: 이는 연상 기억 모델에서와 같이 $N$에 선형적인 임베딩 차원을 $N$에 로그arithmic인 수준으로 줄여주며, MLP 크기를 증가시키지 않습니다.

2. 다중 홉 회상을 위한 용량 - 깊이 트레이드오프

저자들은 사고의 사슬 (CoT) 없이 $k$-홉 사실적 회상을 위한 근본적인 트레이드오프를 확립합니다:

• CoT 없이: $k$-홉 쿼리를 해결하려면 $k$에 지수적인 임베딩 차원 ($d = \tilde{O}(R^k)$) 이거나 사실 수에 선형적인 MLP 폭 ($O(N \cdot R)$) 이 필요합니다. 이는 중간 출력 없이 모델이 단일 은닉 상태 내에 전체 평가 트리를 인코딩해야 함을 보여주는 정보 이론적 하한 (정리 4.2) 을 통해 증명됩니다.
• CoT 사용: 모델이 중간 관계 단계를 생성하도록 허용하면 (Chain-of-Thought) 이 병목 현상을 깨뜨립니다. $d = \tilde{O}(R + k)$의 임베딩 차원과 $\tilde{O}(R)$의 MLP 폭을 가진 단일 층 트랜스포머는 임의의 $k$에 대해 충분합니다. CoT 는 다중 홉 문제를 일련의 단일 홉 문제로 효과적으로 축소하여 임베딩 공간에 전체 $k$-홉 트리를 저장할 필요성을 우회합니다.

3. 기하학적 구조의 경험적 검증

실험은 경사 하강법이 예측된 기하학적 구조를 자연스럽게 발견함을 확인합니다:

• 용량 임계값: 학습 가능한 임베딩의 경우 임계값 $d = \Theta(R \log N)$이 경험적으로 관찰되는 반면, 동결된 임베딩은 $d = \Theta(N^R)$ (또는 유사한 다항식 스케일링) 를 필요로 하여 구조화된 임베딩의 이론적 이점을 확인합니다.
• 중첩: 선형 판독은 주체 임베딩이 속성 벡터들의 선형 중첩으로 잘 근사됨을 보여줍니다.
• 일반적 선택: 인과적 개입은 MLP 가 관계별 선택기로 작용함을 확인합니다. 특히, 주체 임베딩을 새로운 무작위 전단사 (bijections) 를 인코딩하도록 재초기화할 때, 동결된 MLP는 높은 제로샷 정확도를 달성합니다. 이는 MLP 가 특정 사실의 기억이 아닌 일반적인 선택 메커니즘을 학습했음을 보여줍니다.

의의 및 주장

본 논문은 트랜스포머의 기하학적 사실적 기억에 대한 이론적 및 경험적 설명을 제공하며, 지배적인 연상 기억 관점에 대한 대안을 제시합니다.

• 효율성: 임베딩을 관계 위상을 직접 인코딩하도록 구조화함으로써, 트랜스포머는 무작위 임베딩 연상 기억 모델에 비해 훨씬 적은 매개변수 ($N$에 로그arithmic) 로 사실적 회상을 달성할 수 있습니다.
• 회상 메커니즘: 이 발견들은 사실적 회상이 구조화된 기하학적 표현을 통해 구현되며, MLP 가 일반적인 관계 조건부 선택기로 작용함을 시사합니다. 이는 선택 메커니즘이 특정 사실 값과 분리되어 있으므로 임베딩 공간이 적절히 재초기화될 경우 모델이 새로운 사실로 일반화할 수 있는 이유를 설명합니다.
• 사고의 사슬 (Chain-of-Thought) 의 역할: 이 작업은 다중 홉 추론에서 사고의 사슬의 효능에 대한 이론적 설명을 제공합니다. CoT 가 자기 autoregressive 단계를 통해 계산을 분산시킴으로써 직접적인 $k$-홉 예측에 내재된 용량 병목 현상을 우회하여 얕은 모델이 깊은 추론 과제를 해결할 수 있음을 보여줍니다.
• 실제 세계 관련성: 사전 훈련된 LLM 에서 선형 관계 인코딩이 관찰된다는 사실은 이러한 기하학적 메커니즘이 단순한 이론적 산물이 아니라 자연어로 훈련된 모델에 존재함을 시사하며, 이전 경험적 작업에서 관찰된 관계 해독의 선형성과 같은 현상을 설명할 수 있습니다.

저자들은 겸손한 입장을 견지하며, 그들의 설정은 관계를 불투명한 토큰으로, 속성을 임의의 실체로 취급하는 반면, 실제 세계의 지식은 풍부한 의미 구조를 포함한다고 지적합니다. 그들은 이러한 규칙성을 가진 데이터로 이러한 기하학적 설명을 확장하는 것이 향후 연구 방향임을 제안합니다.