[논문 리뷰] Language Models Represent Space and Time

Wes Gurnee, Max Tegmark|arXiv (Cornell University)|2023. 10. 03.

Language and cultural evolution인용 수 41

한 줄 요약

본 논문은 Llama-2와 Pythia 언어 모델이 공간과 시간에 대한 선형적이고 다중 스케일의 표현을 학습하며, 식별 가능한 공간/시간 뉴런을 포함하고 있다는 것을 보여주고, 이는 다음 토큰 예측으로부터 학습된 초기 세계 모델을 시사한다.

ABSTRACT

The capabilities of large language models (LLMs) have sparked debate over whether such systems just learn an enormous collection of superficial statistics or a set of more coherent and grounded representations that reflect the real world. We find evidence for the latter by analyzing the learned representations of three spatial datasets (world, US, NYC places) and three temporal datasets (historical figures, artworks, news headlines) in the Llama-2 family of models. We discover that LLMs learn linear representations of space and time across multiple scales. These representations are robust to prompting variations and unified across different entity types (e.g. cities and landmarks). In addition, we identify individual "space neurons" and "time neurons" that reliably encode spatial and temporal coordinates. While further investigation is needed, our results suggest modern LLMs learn rich spatiotemporal representations of the real world and possess basic ingredients of a world model.

연구 동기 및 목표

• LLMs가 피상적 통계 이상으로 일관된 시공간 표현을 학습하는지 조사한다.
• 내부 활성화를 추출하고 분석하여 공간 좌표(위도/경도)와 시간 좌표(타임스탬프)를 매핑한다.
• 이 표현들의 선형성, 강건성, 그리고 규모와 엔터티 유형에 걸친 교차 엔터티 일관성을 평가한다.
• 공간과 시간에 관련된 개별 뉴런을 식별하여 모델이 이 특성을 활용하고 있음을 확립한다.

제안 방법

• 세계, 미국, NYC; 역사적 인물; 예술; 헤드라인 등 공간 및 시간 관련 엔티티의 다중 규모 데이터셋 여섯 개를 구성한다.
• Llama-2 및 Pythia 모델의 각 층에서 마지막 토큰 활성에 선형 리지 프로브를 실행하여 실제 좌표(위도/경도) 또는 타임스탬프를 예측한다.
• R^2 및 스피어먼 순상관 계수로 프로브를 평가하고, 국지적 지리 정밀도에 대한 오차를 보정하기 위해 근접 오차를 활용한다.
• 공간/시간 특징의 디코딩 가능성을 확인하기 위해 비선형 프로브(MLP)와 비교하여 선형성을 검증한다.
• 프롬프트를 다양화하여 프롬프트 민감도를 평가하고 데이터세트 전반에서 프로브 성능을 관찰한다.
• 활성화를 프로브 방향에 투사하고 뉴런 가중치를 분석하여 “space neurons”와 “time neurons”를 식별한다.

Figure 1: Spatial and temporal world models of Llama-2-70b. Each point corresponds to the layer 50 activations of the last token of a place (top) or event (bottom) projected on to a learned linear probe direction. All points depicted are from the test set.

실험 결과

연구 질문

• RQ1 LLM이 내부 표현에 공간적 및 시간 정보를 인코딩하는가?
• RQ2 공간 및 시간 표현이 중간에서 말단 층의 활성으로 선형적이고 해독 가능한가?
• RQ3 이 시공간 표현은 프롬프트 변 variation에 대해 강건하며 엔터티 유형 간에 통일되어 있는가?
• RQ4 LLM 내부의 개별 뉴런이 공간 또는 시간 좌표를 인코딩하여 이 특징들이 분산적으로 사용되고 있음을 나타내는가?

주요 결과

• 공간 및 시간 특징은 계층과 스케일에 걸쳐 선형 프로브로 복구될 수 있다.
• 표현은 초기에서 중간 층에서 개선되며 모델의 중간 지점에서 정체되고, 더 큰 모델일수록 성능이 더 좋다.
• 비선형 프로브는 선형 프로브에 비해 큰 개선을 제공하지 않으며, 공간/시간의 선형 해독 가능성을 뒷받침한다.
• 프롬프트가 프로브 성능에 미치는 영향은 제한적이며, 무작위 토큰은 성능을 저하시킬 수 있다; 헤드라인 뒤의 마침표 토큰은 이를 개선할 수 있다.
• 공간 및 시간 표현은 다른 엔티티 유형(도시 vs 명소) 간에 일반화되며, 개별 뉴런이 프로브 방향과 일치하여 이 특징들의 실제 사용을 시사한다.
• 데이터를 PCA로 축소해도 프로브가 여전히 정보성을 유지하여 강건한 기반 시공간 구조를 시사한다.

Figure 2: Out-of-sample $R^{2}$ for linear probes trained on every model, dataset, and layer.