기술 요약: 프로토타입 언어 모델 (PRISM)

문제 정의

현재의 대규모 언어 모델(LLM)은 **훈련 데이터 귀속(Training Data Attribution, TDA)**이라는 근본적인 과제에 직면해 있습니다: 즉, 특정 훈련 예시가 모델의 출력에 어떤 영향을 미치는지 결정하는 문제입니다. 표준적인 밀집형(dense) 구조에서 훈련 데이터는 수십억 개의 파라미터 내에 축적된 효과를 통해 밀집되고 얽힌 경로를 거쳐 간접적으로 예측에 영향을 미칩니다. 이로 인해 단일 예시의 영향력은 분산되어 있으며, 비용이 많이 들고 근사치에 의존하며 종종 통계적으로 잘못 설정된 사후 분석(post-hoc analysis) 없이는 추적이 불가능합니다.

기존의 그래디언트 기반 TDA 방법(예: Influence Functions)은 규모가 커짐에 따라 두 가지 복합적인 문제에 직면합니다:

1. 계산 비효력성: 역 헤시안-벡터 곱(inverse Hessian-Vector Product, iHVP)의 정확한 계산은 수십억 파라미터 모델에서 불가능하므로, 충실도를 희생하고 확장성을 선택하는 근사법이 필요합니다.
2. 통계적 오설정: 이러한 방법들은 종-종히 결정론적이고 볼록(convex)한 목적 함수가 수렴할 때까지 최적화된다고 가정하지만, 이는 확률적(stochastic)인 LLM 훈련 상황에는 적용되지 않습니다. 결과적으로, 귀속 추정치는 leave-one-out 재훈련 결과와 일치하지 않을 수 있으며, 훈련 파이프라인의 하이퍼파라미터(깊이, 배치 크기, 데이터 순서)에 민 sensitive하게 반응합니다.

본 논문은 TDA를 단순한 사후 추정 문제가 아니라, 아키텍처 설계 문제로 다루어야 한다고 주장합니다.

방법론: PRISM 아키텍처

저자들은 **해석 가능한 시퀀스 모델링을 위한 프로토타입(Prototypes for Interpretable Sequence Modeling, PRISM)**을 소개합니다. 이 언어 모델 아키텍처는 밀집된 출력 헤드 대신, 학습된 프로토타입의 희소하고 비음수적인 혼합(sparse, non-negative mixture)을 통해 예측을 형성합니다.

핵심 아키텍처

PRISM은 표준 트랜스포머 디코더 백본을 유지하면서, 마지막 밀집 출력 경로를 **희소 프로토타입 재구성(sparse prototype reconstruction)**으로 대체합니다:

1. 프로토타입 뱅크(Prototype Bank): $K$개의 학습된 프로토타입 벡터 $P = [p_1, \dots, p_K] \in \mathbb{R}^{d \times K}$ 세트를 유지합니다.
2. 희소 활성화(Sparse Activation): 주어진 은닉 상태 $z_t$에 대해 모든 프로토타입과의 코사인 유사도를 계산하고, 비음수성을 강제하기 위해 ReLU를 적용한 후, 상위 $k$개의 활성 프로토타입($k \ll K$)을 선택합니다.
3. 재구성(Reconstruction): 은닉 상태는 프로토타입 재구성 $\hat{z}_t$와 잔차(residual) $r_t$로 분해됩니다:
   $$z_t = \hat{z}_t + r_t = \sum_{i \in \mathcal{K}_t} a_{t,i} p_i + r_t$$
   여기서 $\mathcal{K}_t$는 활성 프로토타입의 집합이며, $a_{t,i}$는 그들의 활성화 값입니다.
4. 로짓 형성(Logit Formation): 다음 토큰 로짓은 가산적으로 형성됩니다:
   $$\ell_t = W z_t = W r_t + \sum_{i \in \mathcal{K}_t} a_{t,i} (W p_i)$$
   여기서 $W p_i$는 각 프로토타입의 고정된 "토큰-로짓 시그니처(token-logit signature)"를 나타냅니다.

훈련 목적 함수

PRISM은 언어 모델링의 충실도와 프로토타입의 해석 가능성 사이의 균형을 맞추기 위해 네 가지 손실 항(loss terms)과 함께 엔드투엔드(end-to-end)로 훈련됩니다:

• 교차 엔트로피 ($L_{CE}$): 표준적인 다음 토큰 예측 품질을 보장합니다.
• 재구성 손실 ($L_{REC}$): 잔차를 억제하여 프로토타입이 은닉 상태의 대부분을 설명하도록 유도합니다.
• 클러스터링 손실 ($L_{R1}, L_{R2}$): 이 대칭적인 목적 함수들은 프로토타입을 일관된 훈련 데이터의 이웃에 고정합니다. $L_{R1}$은 프로토타입을 배치 내 가장 가까운 토큰 표현으로 끌어당기고, $L_{R2}$는 토큰 상태를 가장 가까운 프로토타입으로 끌어당깁니다. 이는 프로토타입이 임의의 희소 딕셔너리 방향으로 표류하지 않고 데이터 매니폴드 위에 놓이도록 보장합니다.

이론적 통찰: 곡률 국소화(Curvature Localization)

이 논문의 주요 이론적 기여는 프로토타입 부공간(subspace)에서의 손실 헤시안(loss Hessian) 분석입니다. 저자들은 클러스터링 목적 함수가 **곡률을 국소화(localize curvature)**함을 보여줍니다:

• 헤시안은 전역적 공동 사용 항(프로토타입이 어떻게 함께 사용되는지 포착)과 로컬 블록 대각 항(각 프로토타입 주변의 기하학적 구조를 포착)으로 분리됩니다.
• 더 강력한 클러스터링은 프로토타입 간의 결합(coupling)을 줄여 헤시안의 조건수(condition number)를 개선합니다.
• 이러한 구조는 확장 가능한 TDA를 가능하게 합니다. 곡률이 국소화되어 있기 때문에, 전체 파라미터 공간의 그래디언트 곱을 구할 필요 없이 캐싱된 희소 통계량을 사용하여 영향력 점수를 빠르게 계산할 수 있습니다.

주요 기여

1. 아키텍처 (PRISM): 다음 토근 예측을 위한 최초의 LLM 규모 프로토타입 언어 모델 제품군입니다. 이는 모든 예측에 대해 훈련에 기반한 희소한 프로토타입 혼합을 노출함으로써 사례 기반 추론(case-based reasoning)을 자기회귀(autoregressive) 모델링으로 확장합니다.
2. 이론 및 귀속(Attribution): 희소 프로토타입 구조와 클러스터링 목적 함수가 손실 지형(loss landscape)의 곡률를 국소화함을 입증합니다. 이는 구조화된 헤시안을 생성하여, 동일한 메모리를 소비하면서도 밀집형 그래디언트 기반 베이스라인(예: EK-FAC)보다 약 500배 빠른 훈련 데이터 귀속을 가능하게 합니다.
3. 경험적 성능 및 제어:
   • 성능: FineWeb 및 Nemotron-CC 코퍼스(최대 50B 토큰)에서 130M에서 1.6B 파라미터에 이르는 모델을 통해 성능을 검증했습니다.
   • 정확도: 1.6B 파라미터 규모에서 10B 토큰으로 훈련된 PRISM은 다운스트림 벤치마크에서 매칭된 밀집형 GPT 베이스라인과 대등하거나 최대 2.5% 이내의 차이를 유지했습니다.
   • 제어 가능성: 프로토타입 컨트롤러는 활성 프로토타입의 가중치를 재조정함으로써 다운스트림 정확도를 약 3포인트 개선할 수 있으며, 이 수정 사항은 특정 훈련 이웃으로 추적 가능합니다.
   • 정렬(Alignment): 타겟팅된 프로토타입 억제를 통해 추론 시점에 파인튜닝 없이도 바람직하지 않은 행동(예: NSFW 콘텐츠)을 제거할 수 있으며, 생성 품질의 측정 가능한 손실 없이 이를 수행할 수 있습니다.

실험 결과

규모 및 성능

• 규모: PRISM은 FineWeb 및 Nemotron-CC 코퍼스에서 130M에서 1.6B 파라미터까지 확장되었습니다.
• 정확도: 1.6B 파라미터에서 10B 토큰으로 훈련된 PRISM은 밀집형 GPT 베이스라인과 유사한 성능을 보였습니다. 20B~50B 토큰을 사용하면 밀집형 베이스라인과의 격차를 좁혀 평균 다운스트림 정확도에서 2.5포인트 이내의 차이를 달야냈습니다.
• 효율성: 프로토타입 레이어는 최소한의 오버헤드(예: XL 규모에서 1.7% 파라미터 증가)를 추가하며, 밀집형 GPT 모델의 97.3%에 달하는 훈련 처리량(throughput)을 달성합니다.
• 접지(Grounding): 프로토타입 커버리지 지표($\bar{R}_1, R_2$)가 높게(0.94–1.00) 유지되었으며, 이는 프로토타입이 임의의 방향으로 표류하지 않고 실제 훈련 컨텍스트에 잘 고정되어 있음을 나타냅니다.

귀속 효율성

• 속도: 캐싱된 프로토타입 귀속은 동일한 GPU 메모리 제약 조건 하에서 124M 모델에 대해 EK-FAC보다 470배 빠른 속도를 기록했습니다.
• 신호 보존: 캐싱된 점수는 전체 프로토타입 공간 그래디언트와 비교했을 때 상위 랭크된 훈련 시퀀스의 약 80%를 보존하였으며, 이는 희소하고 국소적인 기하학적 구조가 필요한 귀속 신호를 유지하고 있음을 입증합니다.

제어 및 스티어링(Steering)

• 컨트롤러: 활성 프로토타입에 대한 희소 선형 컨트롤러는 4개의 비이진 다중 선택 과제에서 평균 +2.72포인트의 정확도 향상을 보였으며, 이는 SAE 기반 컨트롤러 베이스라인(+0.32포인트)을 크게 상회하는 결과입니다.
• 안전 정렬: 86개의 식별된 NSFW 관련 프로토타입(뱅크의 0.52%)을 추론 시점에 억제함으로써, 50,000개의 생성물에 대한 평균 NSFW 점수를 0.668에서 0.154로 낮추었으며, 텍스트 품질은 변하지 않았습니다. 가장 어려운 케이스(기본 NSFW 점수 = 2)에서도 점수는 0.144로 떨어졌습니다.

의의 및 주장

본 논문은 PRISM이 해석 가능한 언어 모델이 어떻게 설계되어야 하는지에 대한 패러다임의 변화를 보여준다고 주장합니다. 밀집된 파라미터로부터 훈련 데이터의 흔적을 복구하기 위해 까다로운 사후 분석에 의존하는 대신, PRISM은 예측 경로 내에 직접적으로 희소하고 훈련에 기반한 좌표를 노출합니다.

저자들은 모델이 유사한 예측 성능을 달end할 때, 계산 과정이 더 쉽게 귀속, 감사 및 제어될 수 있는 아키텍처가 더 바람직하다고 주장합니다. PRISM은 다음과 같은 방식으로 이를 달성합니다:

1. 훈련 데이터 귀속을 모델 아키텍처의 고유하고 계산 비용이 저렴한 속성으로 만듭니다.
2. 파인튜닝 없이 모델 행동에 대한 직접적인 개입(스티어링, 억제)을 가능하게 합니다.
3. 특정 테스트 타임 행동을 특정 훈련 이웃으로 추적할 수 있는 "모델 일반화의 과학"을 제공합니다.

이 연구는 "설계에 의한 해석 가능성(interpretable by design)" 모델링이 대규모 환경에서도 실현 가능하다는 것을 시사하며, 표준적인 밀집형 LLM에서 발견되는 효과의 얽힘(entanglement)에 대한 실용적인 대안을 제시합니다.