기술적 요약: 대형 언어 모델에서 다음 토큰 예측에 대한 기하학적 관점

문제 제기

대형 언어 모델 (LLM) 이 모델의 깊이에 걸쳐 예측 정보를 어떻게 처리하고 정제하는지 이해하는 것은 기계적 해석 가능성 분야에서 여전히 핵심적인 과제입니다. 트랜스포머 아키텍처는 공유된 잔여 스트림 (residual stream) 에 덧셈 업데이트를 적용하지만, 다음 토큰 예측에 필요한 정보의 기하학적 진화는 명확하지 않습니다. LogitLens 와 같은 표준 프로빙 (probing) 방법은 "표현 드리프트 (representational drift)"로 고통받는데, 이는 중간 상태가 최종 언임베딩 (unembedding) 기저와 정렬되지 않아 초기 레이어에 존재하는 예측 신호에 대해 "맹목적"이게 만듭니다. 또한, 중첩 (superposition) 의 원리에 따르면 예측 특징은 고차원 잔여 스트림 내에 숨겨진 겹치는 방향에 존재할 수 있습니다. 본 논문은 다음과 같은 질문에 답합니다: 예측 정보는 잔여 스트림의 어디에 존재하며, 네트워크 깊이를 통해 전파됨에 따라 그 기하학적 구조는 어떻게 진화하는가?

방법론

저자들은 표현 렌즈 (representation lenses) (중간 잔여 스트림에서 다음 토큰을 예측하도록 학습된 학습된 아핀 맵) 를 단순한 디코딩 메커니즘이 아닌 기하학적 진단 도구로 재창조하는 프레임워크를 제안합니다.

1. 표현 렌즈: TunedLens 접근법을 따르며, 저자들은 KL 증류 (distillation) 를 통해 중간 표현을 최종 출력 분포에 정렬시키기 위해 레이어별 아핀 보정 $(A_\ell, b_\ell)$을 학습합니다.
2. 예측 판독 부분공간: 전체 보정 행렬을 분석하는 대신, 저자들은 보정 행렬 $A_\ell$에 대해 특이값 분해 (SVD) 를 수행합니다. 그들은 예측 판독 부분공간 $S^{(k)}_\ell$을 상위 $k$개의 우특이벡터 (right singular vectors) 가 spanning 하는 $k$차원 부분공간으로 정의합니다. 이 부분공간은 예측 정보가 집중된 잔여 스트림의 지배적인 방향을 나타냅니다.
3. 그라스만 다양체 궤적: 부분공간 시퀀스 $\{S^{(k)}_\ell\}_{\ell=1}^L$은 그라스만 다양체 $Gr(k, d)$상의 궤적으로 간주됩니다. 이 궤적의 진화는 연속적인 부분공간 사이의 주각 (principal angles) 의 평균 코사인으로 정의된 **판독 부분공간 유사성 (RSS)**을 사용하여 정량화됩니다.
4. 해상도 분석: 저자들은 가시성 (층 간 기하학적 판별력) 과 에너지 보존 (스펙트럼 충실도) 사이의 파레토 프론티어를 추적하기 위해 해상도 매개변수 $k$를 숨겨진 차원 $d$의 1% 에서 50% 까지 변화시킵니다. 그들은 이러한 지표를 균형시키는 중간 영역 ($k \in \{5\%, 10\%, 15\%\}$) 을 식별합니다.
5. 위상 감지: RSS 프로파일은 단봉형 (rise, plateau, descent) 모양을 보입니다. 평활화된 프로파일에서 2 단계 Kneedle 기준을 사용하여 저자는 네트워크 깊이를 세 가지 distinct 한 기하학적 위상으로 분할하는 두 개의 절단점 ($\hat{b}_1, \hat{b}_2$) 을 감지합니다.
6. 실험 설정: 이 프레임워크는 SlimPajama 데이터셋으로 학습된 1B 에서 32B 파라미터까지의 QWEN2.5 와 OLMO2 두 계열의 여덟 가지 모델에 적용됩니다. 지표로는 숨겨진 상태의 유효 랭크, 업데이트 직교성 ($|\cos(\Delta_\ell, h_{\ell-1})|$), 그리고 FFN 대 MHA 업데이트 노름의 비율이 포함됩니다.

주요 기여

본 논문은 두 가지 주요 기여를 합니다:

1. 깊이별 분석을 위한 기하학적 프레임워크: 그라스만 다양체 상에서 예측 정보의 궤적을 추적하는 방법으로, 예측 정보가 가장 초기 레이어부터 존재하지만 중첩 상태로 숨겨져 있어 디코딩을 위해서는 저차원 투사가 필요함을 밝힙니다.
2. 세 가지 등장 기하학적 위상의 식별: 분석은 자기회귀 모델의 순전파 (forward pass) 에서 일관된 삼분 구조를 드러내며, 이는 랭크, 업데이트 직교성, 그리고 부분공간 안정성에 관한 distinct 한 기하학적 행동으로 특징지어집니다.

결과 및 위상 특성화

모든 모델과 규모에서 유사성 프로파일은 단봉 분포를 따르며, 세 가지 위상을 정의합니다:

위상 1: 시딩 멀티플렉싱 (Seeding Multiplexing)

• 기하학적 행동: 예측 부분공간이 재배향되고 안정화됩니다 (RSS 상승). 숨겨진 상태의 유효 랭크가 증가하여 이 위상 끝부분 근처에서 정점에 달합니다.
• 업데이트: 업데이트는 잔여 스트림에 대략 직교합니다. QWEN2.5 에서는 FFN 기여가 지배적이며, OLMO2 에서는 기여가 더 균형을 이룹니다.
• 예측 상태: 최종 토큰은 상위 후보들 사이에 존재하지만 (Hit@1 $\approx$ 20-33%), 지배적이지는 않습니다. LogitLens 는 비효율적으로 유지되며, 이는 신호가 중첩되어 있음을 나타냅니다.
• 스케일링: 이 위상의 지속 시간은 모델 깊이에 대해 약하게 스케일링됩니다 (기울기 $\le$ 0.26), 이는 후보 시딩에 대한 고정 비용을 시사합니다.

위상 2: 호스팅 오버라이딩 (Hoisting Overriding)

• 기하학적 행동: 예측 부분공간은 매우 안정적으로 유지됩니다 (RSS 거의 평탄). 숨겨진 상태의 유효 랭크가 안정화됩니다.
• 업데이트: 업데이트는 잔여 스트림에 대해 최대 직교성을 가집니다 (가장 낮은 $|\cos(\Delta_\ell, h_{\ell-1})|$). 업데이트의 랭크가 이 위상에서 정점을 찍습니다.
• 예측 상태: 후보 분포가 점진적으로 집중됩니다 (Hit@1 이 $\approx$ 56% 로 상승, Hit@5/Hit@1 비율 감소). LogitLens 는 낮게 유지되며, 이는 명시적인 최종 기저 정렬 없이 부분공간 내에서 모호성이 해소됨을 시사합니다.
• 스케일링: 이 위상은 모델 깊이에 따라 선형적으로 확장됩니다 (기울기 $\approx$ 0.6–1.0), 이는 추가 모델 용량이 주로 후보 모호성 해소에 흡수됨을 나타냅니다.

위상 3: 초점 수렴 (Focal Convergence)

• 기하학적 행동: 예측 부분공간이 언임베딩 방향으로 급격히 재배향됩니다 (RSS 하강). 업데이트의 유효 랭크가 감소하여 더 적은 방향에 집중됩니다.
• 업데이트: 업데이트는 잔여 스트림과 매우 정렬됩니다 (가장 높은 $|\cos(\Delta_\ell, h_{\ell-1})|$). 업데이트 노름이 급격히 증가합니다 (최대 10 배), 그리고 모든 계열에서 FFN 기여가 지배적입니다.
• 예측 상태: 후보 분포가 지배적인 예측으로 붕괴됩니다. LogitLens 성능이 크게 향상됩니다 (Hit@1 $\approx$ 30-50%), 이는 최종 토큰이 이제 직접 읽을 수 있게 되었음을 나타냅니다.
• 스케일링: 이 위상의 너비는 깊이에 따라 거의 일정하거나 매우 약하게 스케일링됩니다.

중요성 및 주장

본 논문은 더 깊은 LLM 의 추가 용량이 주로 초기 후보 시딩 (위상 1) 이나 최종 수렴 (위상 3) 이 아닌 **후보 모호성 해소 (위상 2)**에 활용된다고 주장합니다. 이는 심층 트랜스포머에서 다음 토큰 예측의 핵심 계산 부하가 안정적인 기하학적 부분공간 내에서 중첩된 표현을 해결하는 데 있음을 시사합니다.

이 발견은 트랜스포머 역학의 평균장 이론 (예: Bruno 등) 과 일치하면서도 독특하게 확장하여, 관찰된 기하학적 진보에 "정렬 (alignment)", "열 (heat)", "페어링 (pairing)" 위상을 매핑합니다. 이 작업은 예측 정보가 네트워크 후반부에 생성되는 것이 아니라 초기에 존재하며 중첩 상태로 숨겨져 있고, 직교 업데이트와 부분공간 집중 과정을 통해 기하학적으로 정제됨을 확립합니다.

한계점: 저자들은 본 분석이 표준 잔여 연결을 가진 밀집형 디코더 전용 아키텍처로 제한됨을 지적합니다. 해석은 상관 관계 증거에 기반하며, 더 강력한 기계적 근거를 제공하기 위해 인과적 개입 (예: 활성화 패칭) 을 향후 작업으로 제안합니다.