Context Structure Reshapes the Representational Geometry of Language Models

기술 요약: 문맥 구조가 언어 모델의 표현 기하학을 재형성한다

문제 제기
대규모 언어 모델(LLM)은 새로운 과업에 적응하고 문맥 창(In-Context Learning, 이하 ICL) 내의 구조를 학습하는 능력을 보여준다. 선행 연구들은 LLM이 선형 외삽을 통한 다음 토큰 예측을 용이하게 하기 위해, 깊은 층에서 입력 시퀀스를 더 "곧은(straighter)" 신경 궤적으로 조직한다는 점을 밝혀냈다. 그러나 이러한 표현의 직선화(representational straightening)가 모든 형태의 ICL에 적용되는 보편적인 메커니즘인지에 대해서는 여전히 불분명하다. 특히, 문맥의 구조와 LLM이 행동을 적응시키기 위해 내부 표현을 재형성하는 방식 사이의 관계는 아직 완전히 이해되지 않았다. 본 논문은 표현의 직선화가 다양한 인컨텍스트 과업 전반에 걸쳐 일관되게 발생하는지, 아니면 특정 과업 구조에 따라 달라지는지를 조사한다.

방법론
저자들은 서로 다른 인지적 요구와 데이터 구조를 탐색하도록 설계된 세 가지 별개의 과업 클래스에 대해 Gemma-2-27B 모델(베이스/사전 학습 버전)을 분석하였다:

1. 자연어 장거리 의존성(Natural Language Long-Range Dependencies): LAMBADA 데이터셋을 사용하여, 모델이 문장의 마지막 단어를 예측하기 위해 어떻게 광범위한 문맥을 통합하는지 조사하였으며, 이를 의미적 구조가 파괴된 "무작위 텍스트" 대조군과 비교하였다.
2. 그리드 월드 과업(Grid World Tasks): 선형 토큰 시퀀스로부터 잠재적 그래프 구조를 추론해야 하는 합성 과업들이다. 여기에는 다음이 포함된다:
   • 직접 매핑(Direct Mapping): 고유한 단어들로 이루어진 6x6 격자에서의 무작위 보행.
   • 계층적 매핑(Hierarchical Mapping): 잠재 노드와 연관된 "자식" 토큰들의 노이즈 섞인 관찰로부터 숨겨진 4x4 잠재 그래프를 추론.
   • 0-샷 전이(0-shot Transitions): 문맥 내에서 명시적으로 관찰되지 않은 잠재 노드 간의 전이를 추론하는 모델의 능력 테스트.
3. 구조적 예측 과업(Structured Prediction Tasks):
   • 퓨샷 학습(Few-Shot Learning): 표준 벤치마크(예: 국가-수도, 알고리즘 추론)를 "Q: [입력]"과 "A: [출력]"이 교차하는 형식으로 구성.
   • 수수께끼(Riddles): BIG-bench 데이터셋의 의미적 모호성 해소 과업으로, 0-샷과 8-샷 성능을 대조.

기하학적 및 행동적 지표
표현의 기하학을 정량화하기 위해 저자들은 트랜스포머 층 전반에 걸쳐 네 가지 상호 보완적인 지표를 사용하였다:

• 직선화(Straightening): 첫 번째 층에서 후속 층으로 진행됨에 따라 국소 곡률(연속적인 전이 벡터 사이의 각도)이 감소하는 정도로 측정.
• 멩거 직선화(Menger Straightening): 세 개의 연속적인 점을 포함하는 외접원의 반지름을 기반으로 한 보완적 곡률 측정치.
• 유효 차원(Effective Dimensionality): 표현 공간의 부피를 측정하기 위해 PCA 고윳값의 참여 비율(Participation Ratio)을 통해 추정.
• 신장도(Elongation): 궤적의 이방성(첫 번째 대 두 번째 주성분 사이의 분산)을 통해 측정.

행동적 성능은 로짓 차이(그리드 월드용) 및 정확한 일치도(퓨샷 과업용)를 통해 평가되었다.

주요 결과
본 연구는 문맥 구조에 따라 LLM의 표현이 변화하는 방식에 있어 근본적인 이분법이 존재함을 밝혀냈다:

1. 연속적 예측 과업 (자연어 및 그리드 월드):

   • 장거리 의존성 통합이나 잠재 그래프 구조 추론을 요구하는 과업의 경우, 문맥 길이가 길어질수록 특히 중간 층(레이어 15–25)에서 표현의 직선화가 일관되게 증가한다.
   • 이러한 직선화는 행동적 개선(예: 유효한 이웃과 유효하지 않은 이웃 사이의 높은 로짓 분리)과 강한 상관관계를 보인다.
   • 계층적 그리드 월드에서는 0-샷 전이 조건에서도 직선화가 지속되며, 이는 모델이 단순히 토큰 쌍을 암기하는 것이 아니라 기저의 잠재 그래프를 학습함을 시사한다.
   • 기하학적으로, 이러한 과업들은 표현이 더 낮은 차원의, 더 이방성(신장된) 매니폴드로 수축되는 양상을 보인다.

2. 구조적 매핑 과업 (퓨샷 및 수수께끼):

   • 이산적인 예시를 출력에 매핑하는 과업(예: 국가-수도, 수수께끼)의 경우, 표현의 직선화가 일관되지 않으며 과업의 계산 핵심부와 분리된다.
   • 단계적 분리(Phase Dissociation): 퓨샷 예시를 질문(Question), 전이(Transition, 포맷팅), 답변(Answer) 단계로 분해했을 때:
     • 전이 단계(예: "\n A:")는 예시가 추가됨에 따라 점진적인 직선화를 보이며, 예측 가능한 연속체처럼 행동한다.
     • 질문 단계는 첫 번째 예시 이후 직선화가 감소하는 경향을 보인다.
     • 결정적으로, 답변 단계(실제 솔루션이 생성되는 구간)는 점진적인 직선화를 보이지 않으며, 수수께끼와 같은 경우 문맥이 증가함에 따라 직선화가 크게 감소한다.
   • 이는 과업의 핵심 로직(입력을 출력으로 매핑)을 해결하는 메커니즘이 연속적 예측에서 관찰되는 신경 궤적의 선형화에 의존하지 않음을 나타낸다.

의의 및 주장
본 논문은 ICL이 단일한 기하학적 메커니즘에 의해 구동되는 단일한 프로세스가 아니라고 주장한다. 대신, 저자들은 LLM이 특정 과업 구조에 따라 계산 전략을 동적으로 선택하는 "스위스 아미 나이프(Swiss Army knife)"처럼 기능한다고 제안한다.

• 연속적 예측 과업의 경우, 모델은 매니폴드를 선형화하고 외삽을 용이하게 하기 위해 표현의 직선화 전략을 사용한다.
• 구조적 매핑 과업의 경우, 모델은 반드시 궤적 직선화로 나타나지 않는 별개의 메커니즘을 채택한다.

결론적으로, 저자들은 표현의 직선화가 연속적 예측 환경에서의 학습에 대한 강력한 징후이기는 하지만, 모든 ICL 능력의 보편적인 지표는 아니라고 결론짓는다. 이 연구 결과는 기계론적 해석 가능성(mechanistic interpretability)에 대해 주의 깊은 관점을 제공한다. 단일한 기하학적 징후(예: 직선화)에 의존하는 것은 다양한 문맥에 적응하는 모델의 전체 스펙트럼을 포착하는 데 실패할 수 있다. 본 연구는 해석 가능성 가설을 검증할 때 일반화된 표현 구조를 가정하기보다, 다양한 과업 구조에 걸쳐 테스트할 필요가 있음을 강조한다.