[XAI] Anthropic의 Mechanistic Interpretability

Mechanistic interpretability란?

최근 Explainable AI(XAI)에서 가장 핫한 주제 중 하나는 mechanistic interpretability 이다.

 

XAI는 여러 가지 분야가 있다. SHAP, LIME으로 대표되는 설명가능성은 입력 데이터의 feature에 따라 모델이 어떤 반응을 보이는지를 모델 종류에 상관없이 설명할 수 있는 방법론(model agnostic post-hoc explainability)이고, Grad-CAM 같은 방법론은 propagation을 이용해 모델에 들어오는 input 중 어디가 결과에 많은 영향을 미쳤는지 설명하는 방법론이다. 하지만, 이런 방법론은 굉장히 high-level한 설명이며 인간이 이해하는 concept와 연결짓기 어렵다는 한계가 있다.

 

Mechanistic interpretability란, 모델을 파라미터 단위에서 low-level한 설명을 제공하는 방법론으로서, 뉴런 단위부터 모델 내부에서 어떤 "사고과정"이 일어나는지를 다양한 수준에서 연구하는 분야이다. 즉, AI 모델의 리버스 엔지니어링이라고 이해하면 쉬울 것 같다. 2024년부터 XAI 학계에서 CBM(Concept Bottleneck Models)와 함께 가장 큰 스포트라이트를 받고 있는 분야이기도 하고, 많은 XAI 연구실에서 주목하고 있기도 하다.

 

Christopher Olah

Mechanistic Interpretability의 선구자 중 한 명으로 Christopher Olah(https://colah.github.io/)가 있다. Olah는 2018~2020년 OpenAI의 설명가능성 팀을 이끌었으며, 2021년 Anthropic을 코파운딩하며 다양한 LLM interpretability 기술들을 개발하고 있다.

 

 

그는 Distil 이라는 신생 저널의 founder이기도 하다. Distil은 2017년 OpenAI, Deepmind, Y Combinator Research의 합작으로 창간된 저널로, 딥러닝 연구에서의 interpretability를 강조하는 특이한 방향성을 가지고 있는데, Olah는 이 학술지에 feature visualization, transformer circuit을 비롯하여 아주 재미있는 글들을 많이 올려놓았다. 아래 링크에서도 확인 가능한데, 기회가 된다면 이 블로그에서도 다뤄 볼 예정이다.

 

Home - colah's blog

 

이러한 영향으로 현재 AI Safety라는 이름 하에 다양한 AI 회사들이 Mechanistic Interpretability를 연구하고 있다. 특히 Olah가 소속된 Anthropic에서는 Mechanistic Interpretability foundation model을 개발하는 것이 목표이다. https://www.anthropic.com/research/interpretability-dreams

Interpretability Dreams

 

문제 : Superposition

위에 링크를 달아 둔 Olah의 글에서도 확인할 수 있지만, Anthropic의 Mechanistic Interpretability 연구의 주요한 방향성은 superposition 문제를 푸는 것이다.

 

딥러닝 모델 특정 layer에는 유한한 차원이 존재한다. 하지만 입력하는 데이터의 concept (혹은 feature) 개수가 그 차원 수보다 많은 경우, 해당 layer에서는 concept끼리 entangling이 일어나게 된다. (대충 예시를 들자면, "비둘기집 원리"를 떠올려 보자. Layer의 차원이 비둘기집, concept가 비둘기이다.) 이를 superposition이라고 한다.

 

Superposition 문제가 왜 설명가능성에 문제가 되냐 하면, 특정 차원에 concept가 몰려있으면 개별 영향을 분석하기가 어렵기 때문이다. 더 많은 예시는 Olah의 블로그 글(https://transformer-circuits.pub/2022/toy_model/index.html)을 확인해보도록 하자.

Toy Models of Superposition

 

이렇게 하나의 dimension (혹은 하나의 뉴런)에 여러 개의 concept가 중첩되어 있는 현상을 polysemanticity라고 한다.

 

해결 : Sparse Autoencoders

수식이 싫은 분을 위해 결론을 끝부분에 요약해 놓았습니다

 

Concept의 중첩 문제를 어떻게 해결할까? 즉, 어떻게 하면 뉴런 하나당 하나의 concept만 할당할 수 있을까? 최근 연구들은 Sparse Autoencoder 구조를 많이 채택하고 있다.

Autoencoder

 

일반적인 autoencoder는 latent space의 차원수를 input보다 작게 한다. 이는 차원 축소부터, unsupervised representation learning의 기능까지 수행할 수 있게 만든다.

Sparse Autoencoder (SAE)는 latent vector의 차원 수를 input보다 넓히고, L1 loss로 sparsity constraint를 추가하여 학습시키는 구조이다.

즉 위의 loss function을 사용하는 것인데, 첫 번째 항은 input과 output의 reconstruction loss, 두 번째 항은 sparsity constraint로서 latent vector에 L1 regularization을 걸어 주었다.

 

 

얼핏 보면 위 식에서 보이는 VAE (Variational Autoencoder)의 loss term과도 굉장히 흡사해 보인다. 첫 번째 항은 마찬가지로 reconstruction loss이고, 두 번째 항은 latent vector distribution과 정규분포 (쉬운 샘플링을 위해) 간의 KL divergence loss이다. 즉 두 번째 항을 일종의 constraint로도 해석 가능한데, 이렇게 되면 SAE와 VAE는 새로운 랜덤 노이즈에서 샘플링을 하는지의 차이만 있을 뿐 굉장히 유사한 구조가 된다.

 

결론

이렇게 정의된 SAE를 training한 결과로 나오는 latent vector에서는 concept끼리의 중첩이 일어나지 않는다. 즉, 하나의 뉴런 당 하나의 컨셉이 있는 것으로 해석할 수 있으며 이는 모델을 매우 설명가능해지게 만든다.

 

예를 들어 latent를 1000차원으로 설정하면, latent에 있는 1000개의 뉴런을 각각 "컨셉 1", "컨셉 2", ... 등과 같이 이름을 붙여 어떤 컨셉이 어떤 결과에 얼마나 영향을 주었는지를 분석할 수 있게 되는 것이다. 이는 ICLR 2024에 소개되어 많은 주목을 받았으며(https://openreview.net/forum?id=F76bwRSLeK), 2024년 하반기에 구글 딥마인드 Gemma 2의 설명가능성을 위해서도 SAE를 사용한 Gemma Scope가 발표되었다(https://arxiv.org/abs/2408.05147). 

 

최근에는 해석가능성을 더더욱 높이기 위해, LLM이나 VLM을 이용해 "컨셉 n"이 아닌 인간이 이해할 수 있는 언어로 컨셉을 이름붙이는 작업을 하고 있다. (GPT-4o을 이용하면 특정 concept이 인간 언어로 무엇을 나타내는지 매우 효율적으로 라벨링할 수 있다). 또한, LLM에서 토큰 단위로 하는 설명가능성 연구를 비전 분야에 적용하면 어떨 것인가에 대한 논의도 활발히 진행되고 있다.

 

의료 도메인으로 다시 돌아오면, 의료 이미지 분류 작업은 많은 경우 가이드라인에 따라 이루어진다. 수많은 연구가 이루어져 있는 Chest PA 데이터만 해도, 실제 임상에서는 흡기가 적절한가, Airway/heart/bone이 적절한가, opacity가 있냐 등등 기준에 따라 하나씩 체크하는 과정을 거치는데, 이러한 것들이 전부 명시적 concept로 작용할 수가 있다. 만약 비전 모델이 판단하는 concept가 이러한 임상 가이드라인과 일치하게끔 모델을 잘 튜닝할 수만 있게 된다면, 사실상 인간 의사의 사고과정을 완전히 모방하는 AI 모델을 만들 수 있지 않을까 생각한다.