CLT-Forge: A Scalable Library for Cross-Layer Transcoders and Attribution Graphs

1. 문제 정의 (Problem)

거대 언어 모델 (LLM) 의 내부 작동 원리를 이해하려는 기전적 해석 가능성 (Mechanistic Interpretability) 연구에서, 최근 딕셔너리 학습과 트랜스코더 (Transcoder) 를 기반으로 한 접근법이 주목받고 있습니다. 이 방법은 모델 계산을 희소하고 해석 가능한 특징 (features) 과 그 상호작용으로 표현하여 특징 귀속 그래프 (Feature Attribution Graphs) 를 생성합니다.

그러나 기존 접근법에는 다음과 같은 심각한 한계가 존재합니다:

• 과도한 중복성과 규모: 귀속 그래프가 수백 개의 노드로 급격히 커지며, 서로 다른 층 (layer) 에서 유사한 특징이 중복되어 나타나는 문제가 발생합니다. 이는 실용적인 해석을 어렵게 만듭니다.
• 교차 층 트랜스코더 (CLT) 의 훈련 난이도: CLT 는 층 간 특징 공유를 통해 중복성을 줄이고 더 컴팩트한 그래프를 생성하지만, 파라미터 수가 기하급수적으로 증가하여 대규모 모델에서 훈련과 분석이 매우 어렵습니다.
• 오픈소스 인프라의 부재: Anthropic 등 산업계 연구소에서는 독점적인 프레임워크를 사용하지만, 오픈소스 커뮤니티에는 대규모 CLT 훈련을 지원하고 자동화된 해석 (AutoInterp), 귀속 분석, 시각화를 통합한 포괄적인 도구가 부족합니다. 기존 도구들은 분산 훈련, 메모리 효율성 (압축/양자화), 자동 해석 워크플로우가 단편화되어 있어 연구 효율성이 낮습니다.

2. 방법론 (Methodology)

이 논문은 CLT-Forge라는 오픈소스 라이브러리를 제안하여 CLT 의 엔드 - 투 - 엔드 훈련과 해석 가능성을 가능하게 합니다. 주요 기술적 구성 요소는 다음과 같습니다.

2.1. 확장 가능한 훈련 인프라 (Scalable Training Infrastructure)

• 활성화 캐싱 및 압축: 대규모 데이터셋 (수억 토큰) 에서의 훈련을 위해 모든 층의 MLP 입력/출력 활성화를 사전에 계산하여 디스크에 저장합니다. 메모리 효율성을 위해 대칭적 층별 양자화 (Symmetric per-layer quantization) 와 zstd 엔트로피 코딩을 적용합니다 (int8, int4, int2 지원). 이를 통해 저장 공간을 412 배 줄이면서도 재구성 품질은 23% 만 저하됩니다.
• GPU 샤딩 (Feature-wise Sharding): CLT 의 파라미터 수가 너무 커서 단일 GPU 에 실리지 않는 문제를 해결하기 위해, 특징 (feature) 차원으로 모델을 샤딩합니다. 각 GPU 가 특징의 하위 집합을 담당하며, 순전파 (forward pass) 종료 시 결과를 집계합니다. 이는 기존 FSDP(Fully Sharded Data Parallel) 보다 CLT 아키텍처에 더 최적화되어 있습니다.
• 훈련 목적 함수: MLP 출력 재구성 오차, 희소성 (sparsity) 규제, 그리고 죽은 특징 (dead features) 방지를 위한 정규화 항을 포함한 목적 함수를 최적화합니다.

2.2. 자동화된 해석 (Automated Interpretability / AutoInterp)

• 훈련된 CLT 의 특징을 해석하기 위해, 대규모 토큰 스트림에서 각 특징을 가장 활성화하는 시퀀스를 식별하고 통계치를 계산합니다.
• 병렬 처리: 특징 차원을 따라 병렬화하여 각 작업자가 GPU 에서 잠재 활성화를 계산하고, 상위 K 개의 활성화 시퀀스를 실시간으로 추적합니다.
• LLM 기반 설명 생성: 수집된 통계와 대표 시퀀스를 바탕으로 프롬프트를 생성하여 LLM 이 각 특징에 대한 자연어 설명을 자동으로 생성하도록 합니다.

2.3. 귀속 그래프 및 Circuit-Tracer 통합

• Circuit-Tracer 연동: CLT 를 Circuit-Tracer 라이브러리에 매핑하여 층 간 특징 간의 인과적 영향을 정량화하는 귀속 점수 (attribution score) 를 계산합니다.
• 그래프 생성 및 가지치기: 계산된 그래프를 효율적으로 가지치기 (pruning) 하고, 특정 특징이나 클러스터에 대한 개입 (intervention) 효과를 시뮬레이션할 수 있습니다.

2.4. 시각화 인터페이스

• Dash 기반 인터페이스: 생성된 귀속 그래프를 대화형으로 탐색할 수 있는 웹 기반 인터페이스를 제공합니다.
• 기능: 특징별 상위 활성화 토큰, 자동 생성 설명, 특징 클러스터링, 그리고 클러스터 수준의 개입 (intervention) 실행 등을 지원합니다. 이는 기존 도구 (예: Neuronpedia) 를 연구 워크플로우에 더 쉽게 통합할 수 있도록 확장성을 높였습니다.

3. 주요 기여 (Key Contributions)

1. 첫 번째 통합 오픈소스 라이브러리: 대규모 CLT 훈련과 사후 해석 분석을 지원하는 최초의 통일된 프레임워크를 공개했습니다.
2. 효율적인 훈련 아키텍처: 최적화된 모델 샤딩 (Feature-wise sharding) 과 압축된 활성화 캐싱을 통해 수백만 개의 특징을 가진 대규모 CLT (예: LLaMA 1B 기준 150 만 개 특징) 의 훈련을 가능하게 했습니다.
3. 자동 해석 및 귀속 분석 파이프라인: 특징 분석, 자동 설명 생성, Circuit-Tracer 를 통한 귀속 그래프 계산, 그리고 개입 실행을 하나의 통합된 워크플로우로 제공합니다.
4. 확장 가능한 시각화 도구: 연구자들이 특징 상호작용을 직관적으로 이해하고 실험할 수 있는 유연한 시각화 인터페이스를 제공합니다.

4. 결과 (Results)

• 성능: GPT-2 모델에서 기존 연구 (Hanna et al., 2025 등) 와 동등한 성능을 달성했습니다. 표준 프롬프트 (반의어, 카테고리 완성 등) 에서 총 L0=100 기준 약 0.8 의 설명 분산 (explained variance) 과 0.8 의 교체 점수 (replacement score), 0.95 의 그래프 완전성 (completeness) 을 기록했습니다.
• 효율성: LLaMA 1B 모델을 8 개의 GPU 에 분산하여 훈련할 때, 메모리 효율성을 크게 향상시키면서도 안정적인 훈련을 유지했습니다 (bfloat16 정밀도 사용).
• 저장 공간 최적화: 양자화 및 압축 기법을 통해 3 억 토큰 데이터에 대한 활성화 저장 공간을 약 20TB 에서 4TB 로 줄였습니다.
• 실증적 검증: "The opposite of 'large' is"와 같은 문장에 대한 귀속 그래프를 생성하여, 알려진 회로 (circuits) 와 일치하는 해석 가능한 특징 상호작용을 확인했습니다.

5. 의의 및 결론 (Significance)

CLT-Forge 는 기계적 해석 가능성 연구의 장벽을 낮추는 중요한 도구입니다.

• 접근성 향상: 산업계 수준의 대규모 CLT 훈련 및 분석 인프라를 오픈소스로 제공함으로써, 학계 연구자들이 대규모 모델의 내부 메커니즘을 탐구할 수 있게 했습니다.
• 재현성 및 표준화: 분산 훈련, 자동 해석, 시각화를 통합함으로써 연구의 재현성을 높이고 표준화된 워크플로우를 제공합니다.
• 미래 연구 촉진: 이 프레임워크는 더 큰 모델에 대한 CLT 연구, 주의 맵 (attention map) 기반 귀속 분석 확장, 그리고 더 효율적인 아키텍처 개발을 위한 기반을 마련했습니다.

결론적으로, CLT-Forge 는 복잡한 LLM 의 내부 작동 원리를 희소하고 해석 가능한 특징 네트워크로 변환하여 분석하는 데 있어 확장성, 효율성, 그리고 사용 편의성을 모두 갖춘 핵심 인프라로 자리 잡았습니다.