토큰 수준 제어 강화학습으로 LLM 해석 가능성 높이기

초록

Sparse autoencoder를 이용해 LLM 내부 표현을 해석 가능한 특징으로 분해하고, 강화학습 기반 정책이 토큰마다 하나의 특징을 선택·증폭하여 모델 출력을 조정한다. Adaptive Feature Masking으로 다양성을 유지하면서도 단일 특징 해석성을 보장한다. Gemma‑2 2B 모델에 적용해 MMLU·BBQ·GSM8K·HarmBench·XSTest에서 성능 향상과 함께 토큰 수준 개입 로그를 제공한다.

상세 분석

본 논문은 기존 SAE 기반 해석이 “어떤 특징이 활성화되는가”에 머무는 한계를 넘어, “어떤 특징을 증폭했을 때 출력이 변하는가”를 자동으로 탐색한다. 이를 위해 저자들은 Residual Stream의 현재 토큰 위치 표현을 상태로 삼고, SAE 사전학습된 디코더 가중치를 이용해 선택된 특징을 선형으로 추가하는 행동을 정의한다. 정책 네트워크는 단일 레이어의 잔차 벡터를 입력받아 MLP로 특징 선택 로그잇을 출력하고, Top‑k(본 연구에서는 k=1) 방식으로 가장 영향력 있는 특징을 선택한다. 선택된 특징은 고정된 스티어링 계수 c와 함께 디코더 행렬에 곱해 잔차에 더해지며, 이는 토큰 생성 과정에 직접적인 변화를 일으킨다. 강화학습은 PPO를 사용해 토큰 수준 보상을 최적화하고, 가치 함수 Vϕ는 정책의 한계와 보상 추정 오류를 구분하는 분석 도구로 활용된다. Adaptive Feature Masking은 각 샘플마다 활성화된 특징을 마스크에 기록해, 초기 선택에 편향되지 않도록 하면서도 단일 특징 개입이라는 해석 가능성을 유지한다. 실험에서는 Gemma‑2 2B와 Gemma Scope SAE를 결합해 다섯 가지 벤치마크에서 평균 1~2%p의 성능 향상을 달성했으며, 특히 레이어별 분석을 통해 초기 레이어는 구문적 패턴(예: 수식 기호), 후기 레이어는 의미적 추론 구조를 담당한다는 의미론적 계층 구분을 확인했다. 또한, Branch Point Tracking과 Critic Trajectory Analysis를 통해 정책이 잘못된 선택을 하는 시점과 가치 함수가 과소/과대 평가하는 상황을 시각화함으로써 메커니즘적 해석을 가능하게 했다. 이러한 접근은 사후 미세조정 없이도 모델 행동을 정밀하게 조정할 수 있는 경량화된 해석·제어 프레임워크로, 기존 정적 특성 분석과 동적 개입 탐색을 효과적으로 결합한다.