기술 요약: 프로토타입 트랜스포머 (ProtoT)

문제 정의

GPT-4 및 LLaMA와 같은 최첨단 자기회귀 언어 모델(LM)은 다양한 도메인에서 탁월한 성능을 달ей하지만, 추론 과정의 투명성 결여라는 문제를 안고 있습니다. 이들의 내부 계산 경로는 대체로 불투명하여, 출력이 어떻게 생성되는지 이해하는 데 어려움이 있고, 환각(hallination)의 위험과 표적화된 행동 편집(behavioral editing)의 어려움을 초래합니다. 현재의 해석 가능성 방법론들은 주로 사후적(post-hoc) 방식(예: 어텐션 시각화, 프로빙, 희소 오토인코더)이며, 해석 가능성을 주요 제약 조건으로 설계되지 않은 아키텍처를 대상으로 작동합니다. 또한, 표준 셀프 어텐션 메커니즘은 시퀀스 길이에 대해 이차(quadratic) 비용을 발생시켜 확장성을 제한합니다.

방법론: 프로토타입 트랜스포머 (ProtoT)

저자들은 표준적인 이차 비용 셀프 어텐션 메커니즘을 선형 비용의 프로토타입 기반 통신 모듈로 대체한 자기회귀 LM 아키텍처인 ProtoT를 소개합니다. ProtoT는 LLaMA-3의 전체 구조(SwiGLU 피드포워드 네트워크 및 RMS pre-layernorm 포함)를 유지하면서도, "믹서(mixer)" 레이어를 근본적으로 변경합니다.

핵심 아키텍처

1. 프로토타입 기반 믹서: 셀프 어텐션 대신, ProtoT는 $R$개의 학습 가능한 파라미터 벡터인 **프로토타입(prototypes)**을 사용합니다. 이 프로토타입들은 $R$개의 독립적인 통신 채널을 정의하는 필터 역할을 합니다.
2. 쓰기 및 읽기 게이트 (Write and Read Gates):
   • 쓰기 게이트 (Write Gate): 과거 토큰($x_{j}, j < i$)으로부터 정보를 $R$개 채널로 집계합니다. 입력 토큰과 프로토타입 사이의 유사도를 계산하고, 프로토타입에 대한 소프트맥스를 사용하여 정보를 특정 채널로 라우팅합니다. 이는 "의미적 라우팅 병목(semantic routing bottleneck)"을 생성하여, 여러 의미가 채널당 하나의 타협된 표현으로 강제되도록 함으로써 의미적 전문화를 유도합니다.
   • 읽기 게이트 (Read Gate): 현재 토큰과 프로토타입 사이의 유사도에 따라 $R$개 채널로부터 정보를 검색합니다.
3. 접두사 평균 (Prefix Mean, PM) 및 시간 감쇠 (Time Discounting): 각 채널에 집계된 정보는 과거 활성화 값들의 접두사 평균으로 유지됩니다. 학습 가능한 시간 감쇠 파라미터($\beta_k$)는 지수 이동 평균(EMA)을 접두사 평균에 적용하여, 서로 다른 프로토타입이 서로 다른 시간 척도(단기 또는 장기 의존성)로 작동할 수 있게 합니다.
4. 선형 복잡도: 이 아키텍처는 시퀀스 길이에 대해 선형 계산 복잡도($O(R)$)를 달성합니다. 이는 프로토타입들이 서로 상호작용하지 않기 때문이며(Perceiver 스타일의 잠재 공간에서의 $O(R^2)$ 상호작용과 대조됨), 재귀(recurrence)를 통해 캐싱을 통한 상수 시간($O(1)$) 토큰 생성을 가능하게 합니다.
5. 알파 게이트 (Alpha Gate): 훈련 안정성을 디버깅하고 믹서가 예측에 기여하도록 보장하기 위해 각 믹서 모듈의 출력에 ReZero와 유사한 스칼라 게이트가 적용됩니다.

주요 기여

1. 아키텍처 혁신: 셀프 어텐션을 프로토타입 기반 다중 채널 통신으로 대체하여 선형 계산 복잡도를 달성하면서도 경쟁력 있는 성능을 유지하는 새로운 LM 아키텍처인 ProtoT를 도입했습니다.
2. 창발적 해석 가능성 (Emergent Interpretability): 다양한 추상화 수준에 걸쳐 프로토타입이 훈련 과정 중 분리되고 명명 가능한 개념(예: "여성", "뉴질랜드", "COVID-19")을 자동으로 학습함을 입증했습니다. 이는 명시적인 감독 없이 아키텍처의 압박에 의해 채널의 전문화를 유도함으로써 발생합니다.
3. 표적화된 행동 편집: 특정 프로토타입(예: "여성" 프로토타입을 재초기화하거나 마스킹)에 개입함으로써 모델의 추론을 정밀하게 편집할 수 있음을 보여주었습니다. 이러한 개입은 관련 토큰의 확률을 크게 변화시키면서도(예: "women"을 약 16% 감소시킴), 모델의 전반적인 퍼플렉시티(perplexity)에는 무시할 만한 저하(<1% 미만)만을 일으켰습니다.
4. 강건성 분석: ProtoT가 의미를 보존하는 어휘적 섭동(유의어, 오타)에 대해 강건함을 보이며, 이러한 강건성이 프로토타입 라우팅 메커니즘에 의해 매개됨을 보여주었습니다.

실험 결과

저자들은 고정된 백본(6개 레이어, 256 숨겨진 차원)과 FineWeb-Edu의 250M 토큰을 사용하여 ProtoT를 세 가지 베이스라인인 LLaMA(표준 셀프 어텐션), Mamba(상태 공간 모델), DeltaNet(선형 어텐션)과 비교 평가했습니다.

• 확장성: ProtoT는 모델 크기와 데이터 양에 따라 잘 확장됩니다. 대규모 설정(512 숨겨진 차원, 12개 레이어)에서 ProtoT는 29.5의 퍼플렉시티를 달성하여 DeltaNet(31.5)을 능가했으나, LLaMA(25.8)에는 뒤처졌습니다.
• 텍스트 생성: 쌍체 인간 선호도 평가(LLM-as-a-judge)에서 ProtoT는 LLaMA와 DeltaNet보다 우수한 성능을 보였으며, Mamba에 근접한 순위를 기록했습니다.
• 다운스트림 태스크: GLUE 벤치마크에서 ProtoT는 밀집(dense) 베이스라인들과 경쟁할 만한 수준을 유지하며, 종종 Mamba 및 DeltaNet을 능가하거나 대등한 성능을 보였고, 강력한 교차 도메인 강건성을 입증했습니다.
• 처리량 (Throughput): ProtoT는 선형 계산 덕분에 컨텍스트 길이가 증가함에 따라 LLaMA보다 더 잘 확장되며, 32k 토큰 지점에서 LLaMA를 추월합니다. 다만, 매우 긴 컨텍스트에서는 DeltaNet이 가장 높은 처리량을 유지합니다.
• 해석 가능성 지표: ProtoT 프로토타입은 LLaMA의 어텐션 헤드 및 희소 오토인코더(SAE) 특징에 비해 **분리도(disentanglement)**와 주요 테마의 커버리지(coverage) 측면에서 더 높은 점수를 기록했습니다.
• 개입 효능: 특정 프로토타입(예: "여성"을 위한 L9 P7)을 방해했을 때, 시퀀스의 일관성에 대한 부수적인 피해를 최소화하면서 타겟 토큰의 확률(예: "women"에 대해 -16.60%)을 크게 떨어뜨렸습니다.

의의 및 주장

이 논문은 ProtoT가 설계 단계부터 해석 가능한 자기회귀 언어 모델을 향한 유망한 단계임을 주장합니다. 프로토타입 채널을 통해 해석 가능성을 아키텍처에 직접 통합함으로써, 모델은 다음과 같은 이점을 제공합니다:

• 일관되고 명명 가능한 개념을 자동으로 포착합니다.
• 전체 시스템을 불안정하게 만들지 않고도 모델의 행동을 정밀하고 외과적으로 편집할 수 있게 합니다.
• 셀프 어텐션의 대안으로서 강건한 선형 계산 모델을 제공합니다.

저자들은 한계점도 인정했는데, 성능이 선형 베이스라인들과 비교했을 때 경쟁력은 있지만, 순수 퍼플렉시티 측면에서는 LLaMA와 같은 기존 셀프 어텐션 모델에는 현재 미치지 못한다는 점입니다. 또한, 일부 프로토타입은 다의적(polysemantic)일 수 있으므로 프로토타입 레이블을 그 자체로 정답이라기보다는 근사치로 취급해야 한다고 주의를 줍니다. 그러나 이 연구는 능력과 투명성 사이의 균형을 맞추는 것이 가능하다는 것을 보여주며, 더 안전하고 제어 가능한 언어 모델을 위한 경로를 제시합니다.