기술 요약: 유리형 희소 오토인코더 (Rational Sparse Autoencoder, RSAE)

문제 정의

희소 오토인코더(Sparse Autoencoders, SAEs)는 대규모 언어 모델(LLMs)의 내부 활성화를 과잉 결정된(overcomplete) 특징 사전(feature dictionary)의 희소한 선형 결합으로 분해하는 기계론적 해석 가능성(mechanistic interpretability)의 표준 도구이다. 그러나 기존의 SAE 제품군은 고정되고 매끄럽지 않은 인코더 비선형성, 구체적으로 ReLU($\ell_1$ 정규화 포함), TopK, JumpReLU에 의존한다. 이러한 고정된 활성화 프리미티브는 다음과 같은 특정 병리 현상을 유발한다:

• ReLU SAE: 활성화된 특징의 크기 수축(magnitude shrinkage)과 지속적인 "죽은" 잠재 변수(dead latents) 문제를 겪는다.
• TopK SAE: 소프트한 패널티를 하드한 카디널리티 제약으로 대체하여, 비활성화된 특징을 통한 그래디언트 흐름을 끊고 보조적인 부활 손실(revival losses)을 필요로 한다.
• JumpReLU SAE: 학습 가능한 임계값을 도입하지만, 역전파를 위해 연속적 완화 서로게이트(continuous-relaxation surrogates)를 필요로 한다.

이러한 하드코딩된 메커니즘은 관찰된 전-활성화(pre-activation) 기하학에 적응하는 모델의 능력을 제한하며, 재구성 충실도(reconstruction fidelity)와 희소성 사이의 트레이드오프를 왜곡할 가능성이 있다.

방법론

저자들은 고정된 인코더 활성화를 학습 가능한 유리 함수(rational function)로 대체하는 **유리형 희소 오토인코더(Rational Sparse Autoencoder, RSAE)**를 제안한다. RSAE는 표준 SAE 구조(아핀 전-활성화 레이어, 희소 활성화 블록, 선형 디코더)를 유지하면서, 요소별로 적용되는 학습 가능한 유리 함수 $\phi(\cdot)$로 활성화 함수를 수정한다.

아키텍처

RSAE 활성화는 다음과 같이 정의된다:
$$\phi(h) = C_{out} \cdot r_{(a,b)}\left(\frac{h}{C_{in}}\right)$$
여기서 $h$는 전-활성화 벡터이며, $C_{in}$과 $C_{out}$은 학습 가능한 스케일링 파라미터이고, $r_{(a,b)}(t) = P(t)/Q(t)$는 학습 가능한 계수 $(a, b)$를 갖는 유리 함수이다. 스케일링 파라미터는 전-활성화 분포를 유리 함수의 설계 구간인 $[-1, 1]$로 매핑하고, 출력을 기대되는 특징 크기로 다시 매핑한다.

이론적 토대

본 논문은 현재의 SAE에서 사용되는 피스와이즈-어파인(piecewise-affine) 게이트와 유리 함수 사이의 표현력의 비대칭성을 입증한다:

1. 상한 (유리 근사): 저자들은 Zolotarev sign 함수를 사용하여, 학습 가능한 유리 활성화가 타겟 정확도 $\epsilon$에 대해 다항 로그 크기(polylogarithmic size)로 ReLU, JumpReLU, 그리고 공급된 임계값 TopK 게이트를 근사할 수 있음을 증명한다 (구체적으로 불연속 게이트와 마진 $\delta$에 대해 $O(\log(1/\epsilon)\log(1/\delta))$).
2. 하한 (피스와이즈-어파인 한계): 반대로, 특정 $O(1)$-파라미터 유리 타겟 맵을 스칼라 출력 단일 레이어 피스와이즈-어파인 인코더(ReLU/JumpReLU/TopK)로 근사하려면 $\Omega(\epsilon^{-1/2})$의 활성화된 좌표가 필요함을 보여준다. 이는 유리 활성화가 SAE 인코더를 위한 더 효율적인 표현 클래스를 제공함을 시사한다.

훈련 파이프라인

RSAE는 사전 학습된 "교사(teacher)" SAE를 업그레이드하도록 설계된 2단계 파이프라인을 통해 구축된다:

1. 초기화:
   • Remez Exchange: $L_\infty$ 오차를 최소화하기 위해 완화된 Remez exchange 알고리즘을 사용하여 합성 밀집 그리드 상에서 교사의 활성화 프리미티브(ReLU, JumpReLU 또는 TopK)에 유리 함수를 피팅한다.
   • 보정(Calibration): 교사의 실제 전-활성화 분포와 유리 출력 사이의 제곱 오차를 최소화함으로써 유리 계수와 스케일링 파라미터($C_{in}, C_{out}$)를 조정한다.
   • 가중치 상속: 인코더 및 디코더 가중치를 교사로부터 그대로 복사한다.
2. 미세 조정(Fine-Tuning): 모든 파라미터(가중치, 편향, 유리 계수)의 잠금을 해제하고 표준 $\ell_1$-정규화된 재구성 목적 함수 하에서 공동 최적화한다.

주요 기여

1. RSAE 아키텍처: 얕은 SAE에 맞춤화된 근사 이론에 기반하여 학습 가능한 유리 활성화를 갖는 희소 오토인코더를 도입하였다.
2. 이론적 표현력 격차: 유리 활성화가 고정된 SAE 게이트를 압축적으로 근사하는 반면, 피스와이즈-어파인 인코더는 단순한 유리 타겟을 근사하기 위해 훨씬 더 많은 활성화된 좌표가 필요함을 증명하였다.
3. 2단계 훈련 알고리즘: 무시할 만한 오버헤드로 사전 학습된 SAE를 교체할 수 있는 실용적인 초기화 및 미세 조정 절차를 제공한다.
4. 경험적 검증: 세 가지 오픈 웨이트 언어 모델(GPT-2 small, Pythia-160m, Gemma-2-2B)과 세 가지 모든 베이스라인 활성화 제품군에 대한 포괄적인 평가를 수행하였다.

결과

저자들은 잔차 스트림(residual-stream) 활성화에 대해 RSAE를 사전 학습된 베이스라인과 비교 평가하였다. 결과는 미세 조정 후 엄격한 개선을 나타낸다:

• 재구성 지표: RSAE는 테스트된 모든 구성에서 교사 모델에 비해 일관되게 낮은 평균 제곱 오차(MSE), 낮은 $\ell_0$(희소성), 그리고 더 높은 "살아있는(alive)" 특징 비율을 달enc한다.
• 다운스트림 지표: RSAE는 잔차 스트림을 가로챌 때 감소된 교차 엔트로피 저하($\Delta CE$)와 높은 손실 회복률($LR$)을 보여줌으로써 다운스트림 동작을 개선한다.
• 해석 가능성: Pythia-160m에 대한 희소 프로빙 실험은 RSAE가 특징 수준의 해석 가능성을 보존함을 보여준다. 8개의 이진 분류 작업에 대한 전체 딕셔너리 프로브 정확도는 교사와 유사한 수준을 유지하며, 유의미한 저하는 발생하지 않았다.
• 효율성: 업그레이드는 오토인코더당 단 몇 개의 스칼라 파라미터만을 추가한다. 전체 파이프라인(초기화 + 2,000 단계의 미세 조정)은 단일 소비자용 GPU(NVIDIA RTX 5090)에서 몇 분 만에 실행된다.
• 강건성: 이점은 서로 다른 호스트 모델, 베이스라인 활성화 제품군, 그리고 광범위한 베이스라인 희소성 수준($\ell_0 \in [51, 683]$) 전반에서 일관되게 나타난다.

의의 및 주장

본 논문은 RSAE가 해석 가능성을 희생하지 않으면서 동일한 희소도에서 모델의 충실도를 엄격히 개선하는, 사전 학습된 SAE를 위한 "드롭인(drop-in)" 수정 방식이라고 주장한다. 고정된 비선형성을 학습 가능한 유리 함수로 교체함으로써, RSAE는 호스트 모델의 전-활성화에 특화된 기하학에 적응하여 하드코딩된 활성화 프리미티브의 한계를 극복한다.

저자들은 이러한 개선이 무시할 만한 계산 비용(단일 GPU에서 몇 분)과 최소한의 파라미터 오버헤드로 달성된다는 점을 강조한다. 저자들은 RSAE를 기존 SAE를 향상시킬 수 있는 보편적인 업그레이드 전략으로 위치시키며, 이는 재구성 품질과 다운스트위 성능을 모두 높이는 동시에 기계론적 해석 가능성에 필요한 단일 의미적(monosemantic) 특징을 유지할 수 있다. 논문은 평가가 세 가지 모델과 세 가지 활성화 제품군을 다루었지만, 이 접근 방식이 일반적이며 프런티어 모델이나 다른 아키텍처 변형(예: Gated, BatchTopK, Matryoshka SAEs)으로 확장될 수 있다고 언급한다.