Towards Isolated Interventions via Almost Orthogonal Features in Language Models

기술 요약: 언어 모델의 거의 직교하는 특징을 통한 고립된 개입을 향하여

문제 정의
기계론적 해석 가능성(mechanistic interpretability)의 핵심 전제는 언어 모델(LM) 내의 의미 있는 개념이 활성화 공간(activation space)의 선형 특징(linear features)에 대응한다는 것입니다. 이러한 특징들이 신뢰할 수 있는 인과적 개입(causal interventions)을 지원하려면, 하나의 특징을 조작하는 것이 다른 특징의 효과를 실질적으로 변화시켜서는 안 됩니다. 그러나 실제로는 학습된 특징들이 서로 얽혀 있는 경우가 많습니다. 이러한 "특징 간섭(feature interference)"은 특징들이 표현 공간에서 강하게 정렬되어 발생하며, 이로 인해 국소적 개입(예: 모델 스티어링)이 의도치 않은 다운스트림 효과를 일으켜 인과적 기여도를 모호하게 만들고 개입 가능성을 제한합니다. 이 문제는 모델이 보유한 차원보다 더 많은 특징을 표현하는 중첩(superposition) 현상으로 인해 더욱 악화되며, 이는 제어하기 어려운 간섭 패턴을 생성합니다.

방법론
복잡한 시스템이 서로 영향을 주고받지 않는 자율적인 모듈로 분해될 수 있다고 상정하는 독립적 인과 메커니즘(Independent Causal Mechanisms, ICM) 원리에 착안하여, 저자들은 내부 특징들이 **거의 직교(almost orthogonal)**하도록 제약하는 방안을 제안합니다. 방법론은 두 단계의 미세 조정(fine-tuning) 파이프라인을 포함합니다:

1. 직교 SAE 학습: 저자들은 언어 모델의 잔차 스트림(residual stream) 활성화에 대해 희소 오토인코더(Sparse Autoencoder, SAE)를 학습시킵니다. 표준 SAE와 달리, 이들은 디코더 행렬 $D$에 명시적인 직교 정규화 페널티를 도입합니다. 손실 함수에는 그람 행렬(Gram matrix)의 비대각 성분을 페널티로 부여하여 특징 원자(feature atoms) 간의 낮은 자기 상관성(near-orthogonality)을 강제하는 항인 $\lambda \| \text{tril}(D^\top D) \|_F^2$가 포함됩니다.
2. LM 적응: SAE가 학습되고 가중치가 고정되면, 저자들은 이를 사전 학습된 트랜스포머(구체적으로 Gemma 2 2B 모델)의 중간 레이어(26개 레이어 중 12번 레이어)에 삽입합니다. 그 후, 전체적인 성능을 유지하면서 새로운 특징 기하 구조에 모델을 적응시키기 위해 저차원 적응(Low-Rank Adaptation, LoRA)을 수행합니다.

이론적 프레임워크
본 논문은 특징 간섭와 개입 가능성 사이의 관계를 공식화합니다. 모델의 표현을 특징들의 희소 선형 결합($X = DZ$)으로 모델링하고 출력 매핑에 대한 선형 가설을 가정함으로써, 저자들은 실제 출력과 모델 출력 사이의 기대 편차에 대한 상한선을 도출합니다.

• 저자들은 사후 개입 불일치(post-intervention mismatch)가 베이스라인 항과 개별 개입의 크기에 비례하는 개입 특정 항에 의해 제어됨을 보여줍니다.
• 구체적으로, 기대 편차는 특징 딕셔너리의 **자기 상관성(self-coherence)**으로 표기되는 $\mu(D)$와 관련된 $\delta^2 (d-1)\mu(D)$ 항에 의해 상한선이 결정됩니다. 이는 이론적으로 자기 상관성을 줄이는 것(직교성을 높이는 것)이 간섭의 전파를 직접적으로 제한하여 개입의 격리성을 개선함을 입증합니다.

주요 기여

1. 이론적 상한선: 저자들은 잔차 스트림에서의 개입 가능성을 특징 간섭과 연결하는 이론적 상한선을 제공하며, 특징 딕셔너리의 자기 상관성과 모델 출력의 편차를 명시적으로 연결합니다.
2. 직교 미세 조정 파이프라인: 거의 직교하는 디코더를 가진 SAE를 중심으로 언어 모델을 미세 조정하는 파이프라인을 제시합니다. 이 방식이 수학적 추론 작업에서 비페널티(non-penalized) SAE와 비교했을 때 성능 및 해석 가능성 점수를 유사하게 유지함을 보여줍니다.
3. 고립된 개입: 저자들은 스필오버(spill-over) 효과를 명시적으로 제한함으로써 SAE 내 특정 개념에 대한 최초의 고립된 개입을 수행합니다. 이들은 개념 교체(예: 이름 "Jerry"를 "Aquaman"으로 변경)를 성공적으로 수행하면서도 모델의 추론 능력을 보존하고 다른 특징에 미치는 의도치 않은 변화를 최소화했습니다.

결과

• 직교성 대 성능: GSM8K 데이터셋에 대한 경험적 평가 결과, 직교성 페널티 $\lambda$를 높이면 수학적 추론 작업의 정확도를 크게 저하시키지 않으면서도 직교성 손실이 단조 감소(특징 분리 개선)함을 보여줍니다. 엄격한 직교성 페널티를 적용한 모델의 성능은 비페널티 베이스라인과 경쟁할 만한 수준을 유지합니다.
• 개입 효능: 첫 번째 이름(예: "Jerry"에서 "Aquaman"으로) 및 기타 개념(동물, 장소)의 치환 실험에서, 높은 직교성 페널티($\lambda \ge 10^{-6}$)를 가진 모델이 훨씬 더 나은 격리도를 보였습니다.
  • 정확한 치환: 페널티가 증가함에 따라, 모델은 치환된 이름을 생성에 포함하는 빈도가 높아졌습니다(예: $\lambda=10^{-4}$일 때 "Aquaman"을 약 70.9% 확률로 복구한 반면, 페널티가 없을 때는 약 60.1%).
  • 스필오버 감소: 직교성 페널티를 적용한 모델은 원래의 개념을 출력에서 잘못 유지하는 오류 없이 해당 개념을 제거하는 데 더 성공적이었습니다.
• 해석 가능성: (LLM이 텍스트 조각과 특징 설명을 올바르게 매칭하는 능력으로 측정된) 특징의 해석 가능성은 다양한 직교성 페널티에 관계없이 일정하게 유지되었습니다. 이는 직교성을 강제하는 것이 특징의 인간 이해 가능성을 희생하지 않음을 나타냅니다.

의의 및 주장
본 논문은 특징 얽힘이 언어 모델의 고립된 개입을 저해하는 주요 요인이라고 주장합니다. 직교성 정규화가 특징 간섭을 줄인다는 것을 입증함으로써, 저자들은 이 방식이 하위 유용성을 희생하지 않으면서도 SAE 특징의 **식별 가능성(identifiability)과 원자성(atomicity)**을 개선할 수 있는 원칙적인 경로를 제시한다고 주장합니다.

저자들은 자신의 연구를 더 신뢰할 수 있고 예측 가능한 AI 시스템을 향한 한 단계로 겸허히 규정합니다. 이들의 접근 방식은 내부 표현에 대한 정밀하고 국소적인 수정을 허용하며, 이는 특정 구성 요소가 출력에 어떻게 기여하는지 이해하는 데 매우 중요합니다. 또한, 이 방법이 모델 내부 구조에 대한 직접적인 접근과 기술적 전문성을 요구하므로 오용 가능성이 낮다고 언급하며, 모델 출력의 정밀한 스티어링과 관련된 잠재적 이중 용도(dual-use) 우려를 인정합니다. 저자들은 이 연구가 중첩이나 적대적 강건성 문제를 모두 해결한다고 주장하는 것이 아니라, 직교성을 장려하는 것이 더 강건한 모델을 위한 강력한 아키텍처적 편향(architectural bias)으로 작용할 수 있음을 시사한다고 밝힙니다.