GN⁺: HN 공개: Sparse Autoencoders를 활용한 Llama 3.2 해석 가능성 연구

프로젝트 개요

• 현대의 대형 언어 모델(LLM)은 여러 특징을 동일한 뉴런에 중첩하여 개념을 인코딩하고, 각 뉴런의 활성화가 다른 뉴런의 활성화에 따라 여러 해석 가능한 의미를 갖도록 함. 이를 _중첩(superposition)_이라고 부름.
• 희소 오토인코더(Sparse Autoencoders, SAE)는 훈련된 LLM에 삽입되어 활성화를 매우 큰 희소한 잠재 공간으로 투영하여 중첩된 표현을 분리하고 해석 가능한 특징으로 변환함.
• 이 프로젝트는 Anthropic, OpenAI, Google DeepMind가 성공적으로 수행한 연구를 재현하여 해석 가능한 특징을 추출하는 것을 목표로 함.
• Llama 3.2-3B 모델을 위한 기능적이고 해석 가능한 희소 오토인코더를 생성하기 위한 전체 파이프라인을 제공함.

핵심 기능

• PyTorch로 작성된 활성화 캡처부터 SAE 훈련, 특징 해석 및 검증까지의 완전한 엔드 투 엔드 파이프라인 제공.
• 대형 언어 모델로부터 잔여 활성화를 캡처하여 SAE 훈련 데이터셋으로 사용.
• 훈련 데이터를 효율적으로 전처리하고, 다중 GPU를 활용한 대규모 분산 훈련 지원.
• SAE 훈련 중 보조 손실을 구현하여 죽은 잠재 변수를 방지하고 훈련 역학을 안정화함.
• Weights & Biases를 통한 SAE 훈련의 포괄적인 로깅, 시각화 및 체크포인트 제공.
• 해석 가능성 분석 도구를 통해 학습된 특징의 의미 분석을 지원.
• Llama 3.1/3.2의 순수 PyTorch 구현을 통해 외부 종속성 없이 일반 사용 및 결과 검증 가능.
• 텍스트 및 채팅 완성 작업을 통해 모델 행동에 대한 SAE의 영향을 검증하고, 추출된 의미적 특징의 조정 가능.

공개된 자원

1. OpenWebText 문장 데이터셋:

   • 활성화 캡처에 사용된 OpenWebText 데이터셋의 사용자 정의 버전.
   • 원본 텍스트를 유지하며, 빠른 접근을 위해 개별 문장을 파케 형식으로 저장.
   • NLTK 3.9.1의 "Punkt" 토크나이저를 사용하여 문장 분할.

2. 캡처된 Llama 3.2-3B 활성화:

   • Llama 3.2-3B의 23번째 레이어 잔여 활성화 2,500만 문장.
   • 4TB의 원시 데이터를 3.2TB로 압축하여 100개의 아카이브로 분할.

3. SAE 훈련 로그:

   • Weights & Biases를 통한 훈련, 검증 및 디버그 메트릭 시각화 로그.
   • 10번의 에포크와 10,000번의 로그된 단계 포함.

4. 훈련된 65,536 잠재 SAE 모델:

   • 10번의 에포크 후 최종 훈련된 SAE 모델.

프로젝트 구조

1. 데이터 캡처

• capture_activations.py: LLM 잔여 활성화 캡처.
• openwebtext_sentences_dataset.py: 문장 수준 처리용 사용자 정의 데이터셋.

2. SAE 훈련

• sae.py: 핵심 SAE 모델 구현.
• sae_preprocessing.py: SAE 훈련을 위한 데이터 전처리.
• sae_training.py: 분산 SAE 훈련 구현.

3. 해석 가능성

• capture_top_activating_sentences.py: 특징 활성화를 극대화하는 문장 식별.
• interpret_top_sentences_send_batches.py: 해석을 위한 배치 생성 및 전송.
• interpret_top_sentences_retrieve_batches.py: 해석 결과 수집.
• interpret_top_sentences_parse_responses.py: 해석 결과 분석.

4. 검증 및 테스트

• llama_3_inference.py: 핵심 추론 구현.
• llama_3_inference_text_completion_test.py: 텍스트 완성 테스트.
• llama_3_inference_chat_completion_test.py: 채팅 완성 테스트.
• llama_3_inference_text_completion_gradio.py: 상호작용 테스트를 위한 Gradio 인터페이스.