SAE를 활용한 예측: 해석 가능한 특징을 통한 LLM의 도메인 간 일반화 성능 예측 (학습 없이)

SAE as a Crystal Ball: Interpretable Features Predict Cross-domain Transferability of LLMs without Training

최근 몇 년 동안, 사전 학습된 대규모 언어 모델(LLM)은 다양한 작업에서 놀라운 성공을 거두었습니다. 자기 지도 학습의 중요한 역할 외에도, 이러한 모델의 효과는 다운스트림 애플리케이션에서의 성능에 결정적으로 영향을 미치는 사후 학습 과정에 달려 있습니다. 이 과정은 모델을 작업별 데이터와 목표에 맞게 조정하지만, 불가피하게 모델의 변화를 초래하며, 이러한 변화가 다양한 도메인에서 어떻게 전파되는지는 아직 제대로 이해되지 않습니다. 이러한 문제를 해결하기 위해, 우리는 희소 오토인코더(SAE)를 활용하여 사후 학습 과정에서의 일반화 성능을 예측하는 새로운 지표인 SAE 기반 일반화 성능 점수(STS)를 제안합니다. 지도 학습을 예시로 들어, STS는 SAE 표현에서 변화된 차원을 식별하고, 이러한 차원과 다운스트림 도메인 간의 상관관계를 계산하여, 사후 학습 과정에 앞서 일반화 성능을 신뢰성 있게 예측할 수 있습니다. 다양한 모델과 도메인을 대상으로 한 광범위한 실험 결과, STS는 지도 학습의 일반화 성능을 정확하게 예측하며, 실제 성능 변화에 대한 Pearson 상관 계수가 0.7 이상인 것을 확인했습니다. 또한, 우리는 STS를 강화 학습으로 확장하기 위한 초기 단계를 수행했습니다. 우리는 STS가 LLM의 사후 학습 전략을 안내하는 해석 가능한 도구로 활용될 수 있다고 믿습니다. 관련 코드는 https://github.com/PKU-ML/STS에서 확인할 수 있습니다.

In recent years, pre-trained large language models have achieved remarkable success across diverse tasks. Besides the pivotal role of self-supervised pre-training, their effectiveness in downstream applications also depends critically on the post-training process, which adapts models to task-specific data and objectives. However, this process inevitably introduces model shifts that can influence performance in different domains, and how such shifts transfer remains poorly understood. To open up the black box, we propose the SAE-based Transferability Score (STS), a new metric that leverages sparse autoencoders (SAEs) to forecast post-training transferability. Taking supervised fine-tuning as an example, STS identifies shifted dimensions in SAE representations and calculates their correlations with downstream domains, enabling reliable estimation of transferability \textit{before} fine-tuning. Extensive experiments across multiple models and domains show that STS accurately predicts the transferability of supervised fine-tuning, achieving Pearson correlation coefficients above 0.7 with actual performance changes. Beyond this, we take an initial step toward extending STS to reinforcement learning. We believe that STS can serve as an {\color{black} interpretable} tool for guiding post-training strategies in LLMs. Code is available at https://github.com/PKU-ML/STS.