기술 요약: LLM 의 계산 밀도 추적

문제 제기

트랜스포머 기반 대규모 언어 모델 (LLM) 은 수십억 개의 파라미터와 연산을 포함하는 깊고 넓은 계산 그래프로 작동합니다. 그러나 이러한 모델이 모든 입력에 대해 전체 용량을 활용하는지, 아니면 현저히 희소한 서브네트워크에 의존하는지는 여전히 불분명합니다. 주요 미해결 질문들은 다음과 같습니다:

• 충실한 예측에 필요한 유효 계산이 균일하게 분포되어 있는지, 아니면 명확한 밀집 및 희소 서브그래프로 구성되는지 여부.
• 유효 계산의 양이 특정 입력 속성에 의존하는지 여부.
• 모델의 출력 분포를 재구성하는 데 필요한 계산 발자국을 어떻게 정확하게 정량화할 수 있는지.

이러한 질문들을 해결하는 것은 실용적 효율성 (가지치기, 라우팅, 증류) 과 이론적 이해 (기계적 해석 가능성, 언어에 대한 기호적 대립 분산적 관점의 조화, 언어적 복잡성 측정) 모두에 중요합니다.

방법론: s-Trace

저자들은 전체 모델의 출력 분포를 가장 잘 근사하는 크기 $s$의 서브그래프를 효율적으로 추정하는 s-Trace를 소개합니다.

1. 그래프 공식화

LLM 은 방향성 비순환 그래프 $G=(V, E)$로 모델링되며, 여기서:

• 노드 ($V$): 각 레이어의 중간 토큰 상태를 나타냅니다.
• 엣지 ($E$): 정보 흐름에 영향을 미치는 연산을 나타내며, 다음과 같이 분해됩니다:
  • 어텐션 엣지: 소스 토큰과 타겟 토큰 간 벡터 이동을 위한 특정 어텐션 헤드의 기여.
  • MLP 엣지: 다층 퍼셉트론 (Multi-Layer Perceptrons) 으로부터의 업데이트.
  • 잔차 엣지: 이전 레이어 상태를 전달하는 항등 연결.
    이 공식화는 개별 어텐션 헤드의 기여를 집계하는 대신 토큰 단위의 세분화를 가능하게 합니다.

2. Trace 정의

s-Trace($T_s$) 는 전체 모델의 출력 분포 $P(t_{n+1}|t)$와 서브그래프의 제한된 예측 $\hat{P}_{T_s}(t_{n+1}|t)$ 간의 거리를 최소화하는 크기 $s$ (엣지 수) 의 최적 서브그래프로 정의됩니다. 사용된 거리 척도는 총 변이 (Total Variation, TV) 거리이며, 이는 유계 특성과 저확률 꼬리가 아닌 지배적인 확률 모드에 초점을 맞추기 위해 선택되었습니다.

3. Trace 추출

대규모 모델 (7B–14B 파라미터) 에서 최적 서브그래프를 찾는 것은 계산적으로 불가능하므로, 저자들은 훈련이 필요 없는 탐욕적 최선 우선 탐색 (greedy best-first search) 전략을 사용합니다:

• 중요도 점수 매기기: 엣지는 업데이트 벡터의 L1-노름을 기반으로 점수가 매겨지며, 타겟 노드별로 정규화되어 공정한 비교를 보장합니다.
• 탐색: 마지막 레이어의 최종 토큰 노드에서 시작하여, 알고리즘은 예산 $s$에 도달할 때까지 가장 중요한 사용 가능한 엣지와 연결된 노드를 반복적으로 추가합니다.

4. 실험 설정

• 모델: 7 개의 HuggingFace 패밀리에서 가져온 10 개의 LLM 으로, 7B 에서 14B 파라미터까지 다양함 (예: Llama-3, Mistral, Qwen, OLMo, Phi-4).
• 데이터: Wikitext 의 5,000 개 시퀀스 (평균 40 단어).
• 세분화: $10^{-5}$에서 $10^0$(전체 모델) 까지 26 개의 크기 수준에서 Trace 가 추출됨.

주요 결과

1. 계산의 두 가지 명확한 단계

재구성 오차와 Trace 크기 간의 분석은 모든 모델에서 일관된 두 단계를 드러냅니다:

• 구축 단계 ($s \in [10^{-4}, 10^{-2}]$): 재구성 오차의 가파른 비선형 감쇠가 특징입니다. 이 단계에서 약 $s \approx 10^{-3}$(전체 모델의 0.1%) 에서 **최소 코어 (Minimal Core)**가 나타납니다. 이 코어는 확률 질량의 상위 1%(가장 가능성 높은 토큰) 를 회복하기에 충분하며, 입력에 따라 매우 안정적입니다.
• 정제 단계 ($s > 10^{-2}$): 최소 코어 이후 오차는 천천히 감소합니다. 이 단계는 나머지 확률 질량을 회복하기 위해 (예: 40% 에서 90% 커버리지로 이동) Trace 크기가 한 자릿수 증가해야 합니다. 추가된 엣지의 한계 효용은 낮아 정교한 정보 통합을 시사합니다.

2. Trace 의 구조적 구성

• 레이어 분포: 최소 코어는 주로 초기 레이어 노드(초기 및 초기 중간 레이어) 로 구성되어 있으며, $s=10^{-3}$에서 엣지의 약 80% 를 차지합니다. Trace 가 정제 단계로 진입함에 따라 후기 레이어가 추가된 엣지에서 점점 더 지배적이 됩니다.
• 컴포넌트 분포:
  • 전체 모델에서 어텐션 엣지는 가능한 엣지의 99% 이상을 차지합니다.
  • 최소 코어에서는 분포가 놀라울 정도로 균형 잡혀 있습니다: 약 30% 잔차, 약 20% MLP, 약 50% 어텐션 엣지.
  • 정제 단계는 어텐션 엣지 비율의 급격한 증가가 특징이며, 이는 출력 분포를 정교하게 조정하는 데 어텐션의 특정 역할을 나타냅니다.

3. 입력 의존적 희소성

계산 밀도 (오차 - 크기 곡선 아래의 면적으로 측정) 는 입력 속성과 상관관계가 있습니다:

• 모델 불확실성: 계산 밀도와 모델 출력의 섀넌 엔트로피 간에 양의 상관관계가 있습니다. 고엔트로피 (불확실한) 입력은 낮은 재구성 오차를 달성하기 위해 더 큰 Trace(더 많은 계산) 를 필요로 합니다.
• 토큰 빈도: 예측된 토큰의 빈도와 계산 밀도 간에는 음의 상관관계가 있습니다. 빈번한 토큰 (얕은 통계) 은 더 희소한 Trace 로 재구성되는 반면, 드문 토큰은 더 밀집된 서브그래프를 필요로 합니다.
• 모델 간 일관성: 이러한 패턴은 서로 다른 모델 아키텍처와 규모에 걸쳐 놀라울 정도로 일관되므로, 자원 할당은 아키텍처 세부 사항이 아닌 고유한 입력 속성에 의해 주도됨을 시사합니다.

중요성과 주장

이 논문은 LLM 의 유효 계산을 정량화하는 일반적인 방법을 제공하며, 일관된 모듈식 조직을 드러낸다고 주장합니다:

1. 희소 초기 코어: 초기 레이어 연산의 작고 안정적인 서브그래프가 가장 가능성 높은 토큰을 식별하기에 충분한 대략적인 예측을 제공합니다.
2. 밀집 후기 정제: 후속 레이어는 이 예측을 정교화하고 출력 분포의 전체 풍부함을 포착하기 위해 더 밀집되고 분산된 계산을 수행합니다.

해석 가능성에 대한 함의:

• 회로 규모: 최소 코어 (모델의 약 0.1%) 는 "충실한" 회로 발견에 대한 하한을 나타냅니다; 이보다 작은 회로는 필수 역학을 포착하지 못할 수 있습니다.
• 분포의 중요성: 이진 작업 성공 (예: top-1 정확도) 만으로 서브그래프를 평가하는 것은 불충분합니다; 전체 출력 분포의 보존이 중요합니다.
• 초희소성 너머: 기계적 해석 가능성은 초희소 회로를 찾는 것으로 제한되어서는 안 되며, 밀집된 분산 처리 단계를 분석해야 합니다.

인지 과학에 대한 함의:

• 언어적 복잡성: 계산 밀도는 입력 난이도 (불확실성과 희소성) 에 따라 변하므로 언어적 복잡성을 위한 새로운 척도로 사용될 수 있습니다.
• 효율적 처리: LLM 이 빈번한 단어에 대해 더 적은 자원을 동원한다는 발견은 인간 언어 처리의 효율성 이론과 일치합니다.
• 기호적 대 분산적: 기본 통계 (기호적 규칙과 유사) 를 처리하는 희소 초기 레이어 코어의 존재와 이를 followed by 분산 정제는 언어 능력에 대한 기호적 대 분산적 관점 사이의 잠재적 조화를 제공합니다.

저자들은 이러한 발견이 트랜스포머 아키텍처와 영어 데이터에 특화되어 있으며, Trace 식별은 선택된 추출 방법 (L1-노름 기반 탐욕적 탐색) 에 의존한다고 명시합니다. 그들은 이 작업을 단순한 회로 식별을 넘어 LLM 의 내부 역학을 특성화하기 위한 초기 단계로 위치시킵니다.