기술 요약: 다층 트랜스포머 개입을 위한 잔차화된 희소 오토인코더

문제 제기

희소 오토인코더 (SAE) 는 기계적 해석 가능성의 주요 도구로, 트랜스포머의 은닉 상태를 희소하고 해석 가능한 특징으로 분해합니다. 그러나 표준 관행은 각 계층의 원시 활성화 (raw activations) 를 독립적으로 학습합니다. 이 접근법은 트랜스포머의 근본적인 구조적 속성인 잔차 스트림이 깊이에 걸쳐 강하게 결합되어 있다는 사실을 무시합니다. 어떤 계층의 활성화 중 상당 부분은 이전 계층에서 전달되거나 가볍게 변환된 정보로 구성됩니다.

원시 활성화에 독립적인 SAE 를 학습시키는 것은 다층 개입에 두 가지 실용적인 문제를 야기합니다:

1. 용량 중복성: 계층별 사전 (dictionary) 들이 동일한 예측 가능하고 전달된 정보를 재학습함으로써 용량을 낭비합니다.
2. 예측 불가능한 상호작용: 여러 계층이 동시에 교체될 때 (예: 특징 제거 또는 재구성을 통해), 계층 간 상호작용은 가산적이지 않습니다. 동일한 정보가 여러 사전에 표현되어 있다면, 동시 교체는 특징이 중복되거나 상쇄되거나 활성화가 학습 분포 밖으로 벗어나게 만들어 누적 오류와 모델 성능 (교차 엔트로피) 저하를 초래할 수 있습니다.

방법론: 잔차화된 희소 오토인코더 (ReSAE)

저자들은 계층별 정보를 예측 가능한 계층 간 구조로부터 분리하기 위해 잔차화된 희소 오토인코더 (ReSAE) 를 제안합니다. 이 방법은 세 가지 주요 단계를 포함합니다:

1. 쌍별 아핀 잔차화

선택된 계층 부분집합 $L = (\ell_1, \ell_2, \dots, \ell_M)$에 대해, 이 방법은 연속된 선택된 계층 간 아핀 맵을 적합시킵니다. 쌍 $(\ell_m, \ell_{m+1})$에 대해, 아핀 변환 $A_m h^{(\ell_m)} + c_m$이 리지 정규화된 일반 최소제곱법을 통해 이전 계층의 활성화로부터 후속 계층의 활성화를 예측하도록 학습됩니다.
잔차화된 활성화 $r^{(\ell_{m+1})}$는 이 아핀 맵으로 설명되지 않는 후속 계층 활성화의 성분으로 정의됩니다:
$$r^{(\ell_{m+1})} = h^{(\ell_{m+1})} - (A_m h^{(\ell_m)} + c_m)$$
첫 번째 앵커 계층은 경험적 평균으로 중심화됩니다.

2. 독립적 계층별 SAE 학습

각 SAE 는 원시 활성화가 아닌 잔차에 대해 학습됩니다.

• 입력: 잔차 활성화는 (블록별 RMS 정규화를 통해) 정규화되어 표준 BatchTopK SAE 로 입력됩니다.
• 목적: SAE 는 잔차 목표의 재구성 손실을 최소화합니다.
• 독립성: 그라디언트는 서로 다른 블록의 SAE 간에 흐르지 않습니다. 각 SAE 는 특정 계층에서 도입된 "새로운" 정보만 학습합니다.

3. 원본 공간 재구성과 온라인 개입

표준 프로토콜을 사용하여 ReSAE 를 평가하기 위해, 재구성된 잔차는 적합된 아핀 체인을 통해 원본 활성화 공간으로 매핑됩니다.

• 재구성: $\hat{h}^{(\ell_{m+1})} = A_m \hat{h}^{(\ell_m)} + c_m + S^{-1}_{m+1} \hat{z}^{(m+1)}$, 여기서 $\hat{z}$는 잔차의 SAE 재구성입니다.
• 온라인 개입: 후크를 사용한 순전파 (forward pass) 동안, 계층 $\ell_m$의 잔차 입력은 현재 유입 활성화에서 이전에 재구성된 계층 $\hat{h}^{(\ell_{m-1})}$의 예측 기여분을 빼서 계산됩니다. 이는 상류 계층이 수정되었더라도 계층 $\ell_m$의 SAE 가 학습 목표와 유사한 입력 분포를 보도록 보장합니다.

주요 기여

1. ReSAE 아키텍처: 이전 선택된 계층에서 선형적으로 예측 불가능한 활성화 성분에 대해 SAE 를 학습하고, 재구성을 아핀 체인을 통해 구성하는 학습 패러다임의 도입.
2. 개선된 다층 개입: 원시 SAE 에 비해 더 적은 원시 분산을 재구성함에도 불구하고, ReSAE 가 동시 다층 교체 하에서 트랜스포머 교차 엔트로피를 더 충실하게 회복함을 입증.
3. 감소된 중복성: 잔차화가 더 직교적인 디코더 방향을 생성하여 용량이 중복된 계층 간 구조가 아닌 계층 지역 정보로 향함을 나타내는 증거.
4. 포괄적 평가: 재구성 품질, 교차 엔트로피 저하, 디코더 유사성, 희소 프로빙, 표적 프로브 교란 (TPP), 그리고 허위 상관관계 제거 (SCR) 를 다루는 Pythia-1.4B 와 Gemma-2-9B 에 대한 광범위한 벤치마크.

결과

실험은 Pythia-1.4B(계층 0, 6, 12, 18) 와 Gemma-2-9B(다양한 계층 간격) 에서 수행되었습니다.

• 재구성 대 기능성: ReSAE 는 선형적으로 예측 가능한 구조를 의도적으로 폐기하므로 원시 SAE 에 비해 원본 활성화 공간에서 더 낮은 설명된 분산 (EV) 을 보입니다. 그러나 이 낮은 EV 는 더 낮은 디코더 코사인 유사성(더 적은 중복성) 과 상관관계가 있습니다.
• 교차 엔트로피 회복: 다층 순차적-온라인 교체(오류가 전파됨) 하에서 ReSAE 는 특히 더 높은 희소성 수준 ($k \ge 160$) 에서 모델의 교차 엔트로피 회복에 있어 원시 SAE 를 크게 능가합니다. 교사 강제 설정 (오류 전파 없음) 에서 ReSAE 는 일관되게 최고의 교차 엔트로피 회복을 달성합니다.
• 과상호작용: 원시 SAE 는 여러 계층이 교체될 때 큰 부정적인 "과상호작용"(비가산적 저하) 을 보입니다. ReSAE 는 이 지표를 0 에 가깝게 줄여 계층별 교체들이 더 예측 가능하게 구성됨을 나타냅니다.
• 특징 품질:
  • 희소 프로빙 및 TPP: ReSAE 는 일반적으로 희소 프로빙 정확도와 표적 프로브 교란에서 원시 SAE 를 능가하며, 학습된 특징이 선형 판독과 선택적 제거에 더 유용함을 시사합니다.
  • 허위 상관관계 제거 (SCR): ReSAE 는 여기서 균일하게 원시 SAE 를 능가하지는 않습니다. 일부 설정 (특히 Pythia 에 대한 계층별 공동 제거) 에서 원시 SAE 가 더 높은 점수를 받습니다. 이는 잔차화된 특징이 작업 신호를 손상시키지 않고 특정 허위 속성을 제거하는 데 덜 선택적일 수 있음을 시사합니다.
• 규모 및 간격: ReSAE 의 이점은 서로 다른 계층 간격을 가진 더 큰 Gemma-2-9B 모델에서도 지속되며, 디코더 중복성의 일관된 감소와 교차 엔트로피 회복의 개선이 관찰됩니다.

중요성 및 주장

이 논문은 트랜스포머 잔차 스트림을 고립된 계층들의 집합이 아닌 결합된 객체로 취급하는 것이 효과적인 다층 개입에 필수적이라고 주장합니다.

• 용량 할당: 선형적으로 예측 가능한 계층 간 구조를 제거함으로써, ReSAE 는 희소 사전이 해당 특정 계층에서 모델의 하류 계산과 가장 관련 있는 활성화 성분들에 집중하도록 강제합니다.
• 개입 안정성: 이 방법은 다층 개입을 위한 더 안정적인 기본값을 제공하여 표준 SAE 배포를 괴롭히는 누적 오류와 예측 불가능한 상호작용을 줄입니다.
• 절충: 저자들은 ReSAE 가 모든 해석 가능성 작업 (예: SCR 결과의 혼재) 에 대한 보편적인 개선은 아니라고 인정합니다. 그러나 다층 교체와 특징 기반 개입이라는 특정 목표에 대해서는, 더 낮은 원시 분산과 더 높은 기능적 충실도 및 감소된 중복성 사이의 절충이 유익합니다.

이 연구는 아핀 회귀를 통한 예측 가능한 구조 제거가 다층 설정에서 SAE 를 위한 유용한 전처리 단계임을 시사하지만, 최적의 예측자 (아핀 대 비선형) 와 계층 선택은 여전히 열린 설계 공간으로 남아 있습니다.