기술 요약: 제어 강화 학습 (CRL)

1. 문제 정의

희소 오토인코더 (SAE) 는 언어 모델의 활성화 (activations) 를 해석 가능하고 인간이 읽을 수 있는 특징으로 분해하는 강력한 도구로 부상했습니다. 그러나 기존 분석 방법들은 대부분 **정적 (static)**입니다. 즉, 특정 입력에서 어떤 특징이 활성화되는지는 식별하지만, 어떤 특징이 증폭될 때 실제로 모델의 출력을 변화시키는지는 밝히지 못합니다. 현재 활성화 엔지니어링 기법들은 특징 선택을 위해 대조적 데이터셋에 의존하거나 계수 최적화를 위해 광범위한 활성화 저장을 필요로 하며, 생성 품질에 기반하여 조향 (steering) 전략을 조정할 수 있는 적응형 피드백 메커니즘이 부족합니다.

본 논문은 정적 특징 식별과 동적이며 해석 가능한 개입 사이의 격차를 해소합니다. 핵심 과제는 모델을 조향하기 위해 SAE 특징을 선택할 뿐만 아니라, 왜 특정 특징이 선택되었는지와 어떻게 그것이 출력에 영향을 미쳤는지를 설명하는 토큰 수준의 개입 로그를 제공하는 프레임워크를 개발하는 것입니다. 이를 통해 조향을 기계적 해석 도구로 변환합니다.

2. 방법론: 제어 강화 학습 (CRL)

저자들은 트랜스포머 표현의 제어를 SAE 특징에 대한 마르코프 결정 과정 (MDP) 으로 간주하는 제어 강화 학습 (CRL) 프레임워크를 제안합니다.

2.1. MDP 공식화

• 상태 ($S$): 관측치는 현재 토큰 위치의 특정 층 $\ell$에서의 잔류 스트림 활성화 ($x \in \mathbb{R}^d$) 입니다.
• 행동 ($A$): 하나 이상의 SAE 특징을 증폭하도록 선택하는 이진 벡터입니다. 정책은 해석 가능성을 위해 일반적으로 $k=1$인 상위-$k$개 특징을 선택합니다.
• 전이 ($P$): 선택된 특징 개입이 적용된 트랜스포머의 순전파 (forward pass) 에 의해 결정되는 결정론적 과정입니다.
• 보상 ($R$): 최종 출력 품질에 기반한 작업별 보상 (예: MMLU 의 정확성, 안전 작업에서의 거부 등) 입니다.

2.2. CRL-Token 및 특징 조향

각 생성 단계 $t$에서 프레임워크는 토큰 수준의 개입을 수행합니다:
$$\tilde{x}_t = x_t + c \cdot a_t W_{dec}$$
여기서 $a_t$는 희소 이진 선택 벡터, $c$는 조향 계수, $W_{dec}$는 SAE 디코더입니다. 정책 네트워크 $\pi_\theta$는 잔류 스트림 활성화를 특징 선택 로짓 (logits) 으로 매핑하며, 2 층 MLP 로 구현됩니다.

2.3. 적응형 특징 마스킹 (AFM)

정책이 자주 활성화되는 특징의 좁은 집합으로 수렴하는 것을 방지하기 위해 CRL 은 적응형 특징 마스킹을 도입합니다.

• 샘플별 불리언 마스크가 특징 선택 옵션을 제한합니다.
• 마스크는 특징의 작은 부분집합으로 시작하여 생성이 진행됨에 따라 점진적으로 확장되며, 이전 단계에서 자연스럽게 활성화된 특징들을 포함합니다.
• 이 메커니즘은 단일 특징 개입의 해석 가능성을 유지하면서 다양한 특징 발견을 장려합니다.

2.4. 최적화

정책과 상태 값을 추정하는 크리틱 네트워크는 **근접 정책 최적화 (PPO)**를 사용하여 훈련됩니다.

• 조향 계수 ($c$): 올바르게 예측된 훈련 샘플들의 활성화 크기 평균에 의해 적응적으로 결정되며, 층 간 수동 튜닝의 필요성을 제거합니다.
• 크리틱 네트워크: 가치 함수 $V_\phi(s)$를 추정하는 2 층 MLP 로, 분석 중 정책의 한계와 가치 추정 오류를 분리할 수 있게 합니다.

3. 주요 기여

논문은 네 가지 주요 기여를 제시합니다:

1. 토큰 수준 해석 가능 개입 로그: 각 토큰에서 선택된 정확한 SAE 특징과 그것이 출력에 미친 영향을 식별하는 로그를 생성하는 프레임워크로, 세밀한 행동 분석을 가능하게 합니다.
2. 진단 도구:
   • 분기점 추적 (Branch Point Tracking): 동일한 접두어가 특징 선택에 따라 다른 결과를 산출하는 중요한 의사결정 지점을 식별합니다.
   • 크리틱 궤적 분석: 학습된 가치 함수를 사용하여 정책 실패 (특징 선택) 와 크리틱 실패 (보상 추정) 를 구분합니다.
3. 적응형 특징 마스킹 (AFM): 단일 특징 해석성을 희생하지 않으면서 다양한 특징 발견을 보장하기 위해 탐색과 활용을 균형 있게 조절하는 메커니즘입니다.
4. 특징 의미론의 경험적 특성화: 초기 층은 구문 구조를 인코딩하고 후기 층은 추상적 의미론을 포착하며, 다양한 작업이 고유한 병목 패턴을 보임을 입증합니다.

4. 실험 결과

프레임워크는 Gemma Scope 및 LLaMA Scope SAE 를 사용하여 Gemma 2 2B(및 LLaMA-3.1 8B)에서 MMLU, BBQ, GSM8K, HarmBench, XSTest의 다섯 가지 벤치마크를 통해 평가되었습니다.

• 성능 향상: CRL 은 모든 벤치마크에서 개선을 달성했습니다. 주목할 만한 향상은 다음과 같습니다:
  • HarmBench: +7.66 퍼센트 포인트 (41.46% $\to$ 49.12%).
  • BBQ (모호한 경우): +5.69 퍼센트 포인트.
  • MMLU: +3.31 퍼센트 포인트.
  • GSM8K: +1.03 퍼센트 포인트.
• 상호 보완성: CRL 은 감독 미세 조정 (SFT) 모델 위에 적용될 때 추가적인 이점을 제공합니다 (예: SFT + CRL 로 MMLU 정확도가 55.2% 에서 56.1% 로 증가).
• 층별 의미론:
  • Layer 10: 특징은 구체적인 구문적 속성 (예: "수학적 표기법", "대명사 참조") 을 포착합니다.
  • Layer 20: 특징은 추상적, 문맥적 의미론 (예: "논리적 유도 구조", "수학적 증명") 을 포착합니다.
  • 최적 개입 층은 작업에 따라 다릅니다 (예: BBQ 의 경우 Layer 5, MMLU 의 경우 Layer 24).
• 크리틱 분석:
  • BBQ: 명확한 샘플 구별과 함께 효과적인 정책 - 크리틱 조율을 보여줍니다.
  • MMLU: 수정된 샘플과 오도된 샘플이 거의 구별되지 않는 "크리틱 병목" 현상을 보이며, 이는 크리틱이 복잡한 추론 작업에 대한 가치를 추정하는 데 어려움을 겪음을 시사합니다.
  • HarmBench 대 XSTest: HarmBench 는 정답/오답 샘플 간 명확한 격차를 보이지만, XSTest 는 역전된 추정 패턴을 보여 보상 추정에서의 작업 의존적 어려움을 나타냅니다.

5. 중요성 및 주장

본 논문은 CRL 을 기계적 해석 가능성 분야에서 중요한 진전으로 위치시킵니다. 그 주요 중요성은 정적 특징 분석(무엇이 활성화되는지 관찰) 에서 동적 개입 프로브(무엇이 모델을 변화시키는지 관찰) 로의 패러다임 전환에 있습니다.

• 해석 가능성: 단일 특징 선택을 강제하고 개입을 로그로 기록함으로써 CRL 은 모델 행동에 대한 "기계적" 관점을 제공하며, "수학적 연산"과 같은 특징이 단어 문제에서 "아기"와 같은 정량적 단위로 어휘적으로 일반화될 수 있음을 드러냅니다.
• 진단 능력: 이 프레임워크는 모델과 SAE 모두를 위한 진단 도구 역할을 합니다. AFM 이 있음에도 불구하고 정책이 의미론적으로 일관성 없는 특징으로 수렴한다면, 이는 제어 방법의 실패가 아니라 SAE 의 결함 (예: 누락된 개념 또는 다의성) 을 나타냅니다.
• 실용성: CRL 은 낮은 추론 오버헤드 (8–9% 지연 시간) 로 작동하며 기본 모델 가중치를 재훈련할 필요가 없어 경량화된 사후 (post-hoc) 향상 메커니즘을 제공합니다.

저자들은 강화 학습과 결합된 학습된 특징 조향이 언어 모델을 이해하고 제어하기 위한 실용적인 도구를 확립하며, 정적 분석을 동적이고 작업별 통찰력으로 보완한다고 결론지었습니다.