기술 요약: Concept-SAE

문제 정의

표준 희소 오토인코더(Sparse Autoencoders, SAEs)는 신경망 활성화를 학습된 특징 사전(dictionary of learned features)으로 분해할 수 있게 함으로써 기계론적 해석 가능성(mechanistic interpretability) 분야에서 강력한 도구로 자리 잡았습니다. 그러나 이들의 효용성은 여전히 주로 수동적입니다. 현재의 워크플로우는 연구자가 잠재 토큰을 수동으로 검사하고, 특정 활성화를 보이는 이미지를 식별하며, 시각적으로 기저의 개념을 추론해야 합니다. 이러한 접근 방식은 모델이 사용자가 명시한 개념을 우연히 인코딩하는지 여부를 검증할 수는 있지만, 사용자가 명시적으로 관심을 갖는 개념을 능동적으로 테스트, 프로빙 또는 체계적으로 평가하는 데는 한계가 있습니다.

Concept Bottleneck Models(CBMs)나 임베딩 기반 개념 벡터와 같이 이 간극을 메우려는 기존의 시도들은 상당한 단점에 직착합니다. CBM은 이진 존재 신호에 의존하여 고차원 개념을 분리하지 못하거나 제약되지 않은 토큰으로 의미적 혼입(semantic bleeding)을 일으킵니다. 외부 감독이 없는 임베딩 기반 접근 방식은 의도한 의미로부터 벗어나는 경향이 있어, 해석 가능성을 저해하고 개방형 분석을 제한합니다. 따라서 사용자가 정의한 개념과 양방향으로 능동적인 상호작용을 수행하면서도, 새로운 미제약 특징(unconstrained features)을 발견하는 SAE의 능력을 보존할 수 있는 프레임워크가 필요합니다.

방법론: Concept-SAE

저자들은 표준 SAE에 구조화되고 제어 가능하며 가역적인 인터페이스를 결합한 Concept-SAE 프레임워크를 제 제안합니다. 핵심 혁신은 활성화 서브스페이스를 두 개의 직교하는 성분으로 분해하는 **하이브리드 분리 전략(hybrid disentanglement strategy)**입니다.

1. 개념 토큰 (Concept Tokens, 지도 학습): 외부에서 지정된 의미(예: "턱수염", "눈")와 정렬된 토큰입니다. 이들은 이중 감독(dual supervision) 하에 학습됩니다:

   • 개념 존재성 (Concept Existence): 개념의 이진 존재 여부를 나타내는 스칼라 점수.
   • 공간적 국소화 (Spatial Localization): 특징 맵 내에서 개념의 공간적 위치를 나타내는 세그멘테이션 마스크.
     이 이중 감독은 토큰이 충실하고, 공간적으로 근거가 있으며, 잔차 공간(residual space)으로부터 깔끔하게 분리되도록 보장합니다.

2. 자유 토큰 (Free Tokens, 비지도 학습): 정의된 개념이 커버하지 않는 나머지 모든 정보를 포착하며, 표준 SAE 토큰처럼 작동합니다. 이들은 미제약 특징을 발견하는 SAE의 개방형 탐색 능력을 유지합니다.

학습 파이프라인

학습 과정은 4단계 파이프라인으로 구성됩니다 (논문의 그림 1 참조):

• 레이블 생성 (Label Generation): 시각-언어 모델(VLM)이 이진 존재 레이블을 생성하고, 세그멘테이션 모델이 공간 마스크를 생성합니다. 이들은 융합되어 깨끗한 정답 주석(개념이 부재할 경우 마스크를 억제함)을 생성합니다.
• 개념 토크나이저 ($T_{concept}$): 비전 모델의 특징을 개념 공간으로 매핑하여 존재 점수와 공간 마스크를 예측하는 인코더입니다. 예측값과 정답 사이의 MSE를 최소화하도록 학습되며, 희소성을 장려하기 위해 L1 페널티를 적용합니다.
• 개념 어그리게이터 ($A_{concept}$): 예측된 개념 표현으로부터 원래의 비전 모델 특징을 재구성하는 디코더입니다. 재구성 오차를 최소화하도록 학습되며, 내부 가중치 분포를 KL 발산을 통해 토크나이저와 정렬합니다.
• 자유 토크나이저 및 어그리게이터 ($T_{free}, A_{free}$): 이 모듈들은 개념 모듈과 함께 개념 잔차 특징을 재구성하기 위해 공동으로 학습됩니다. 이들은 외부 감독 없이, 오직 (개념 + 자유)의 결합된 출력이 원래의 특징을 재구성하는 오차를 최소화하는 목적 함수만을 바탕으로 학습됩니다.

주요 기여

본 논문은 세 가지 주요 기여를 주장합니다:

1. Concept-SAE 프레임워크: SAE에 고충실도, 양방향 인터페이스를 부여하여, 정의된 인간의 개념을 검증하고 조작할 수 있게 하는 새로운 아키텍처를 제공합니다.
2. 하이브리드 분리 전략: 제안된 전략이 다른 방법들보다 인터페이스 품질 면에서 우수하며, 놀라울 정도로 충실하고, 공간적으로 국소화되며, 자유 토큰 잔차 공간으로부터 깔끔하게 분리된 개념 표현을 생성함을 실험적으로 입증했습니다.
3. 진단적 유용성: 추출된 개념 점수가 모델 내부의 확실성을 정확하게 반영하며, 다음 세 가지 특정 진단 평가에 사용될 수 있음을 입 demonstrated 했습니다:
   • 탐지 (Detection): 적대적 이미지 샘플 식별.
   • 제어 가능성 (Controllability): 모델의 오예측을 수정하기 위한 통제된 역사실적 편집(counterfactual editing) 수행.
   • 안정성 (Stability): 적대적 섭동을 통한 모델의 취약점 진단.

실험 결과

저자들은 ResNet-18 및 Vision Transformer (ViT-B/32) 아키텍처를 사용하여 CelebA(성별 분류) 및 ImageNet-1k(객체 분류)에서 Concept-SAE를 평가했습니다.

• 충실도 및 국소화 (RQ1): Concept Embedding Model (CEM) 베이스라인과 비교했을 때, Concept-SAE는 유의미하게 높은 국소화 비율(Localization Ratios)(예: ResNet layer 5에서 1.472 vs 1.019)을 달성했습니다. 이는 Concept-SAE가 얼굴 개념을 배경 노이즈로부터 성공적으로 격리한 반면, CEM은 개념을 무관한 배경 콘텐츠와 얽히게 했음을 나타냅니다. 자유 토큰은 풍부한 의미 정보를 유지하였으며, 이는 개념 토큰의 도입이 개방형 발견 능력을 저하시키지 않았음을 확인시켜 줍니다.
• 적대적 탐지 (RQ2): 개념 점수의 엔트로피를 모델 확실성의 지표로 사용했습니다. 오답 및 적대적 샘플은 정답 샘플에 비해 일관되게 높은 엔트로피(모호성)를 보였습니다. 이 엔트로피를 사용하는 임계값 기반 분류기는 JPEG, Randomization, Feature Squeezing과 같은 기존 베이스라인보다 우수한 AUC 점수(예: CelebA-ViT에서 0.974)를 달성하여, 점수가 내부 프로세스를 충실히 반영함을 증명했습니다.
• 제어 가능성 (RQ3): 이 프레임워크는 **역사실적 편집(counterfactual editing)**을 가능하게 했습니다. 개념 점수를 수정하고(예: "beard"를 0.0 또는 1.0으로 설정) 어그리게이터를 통해 특징을 재구성함으로써, 저자들은 모델의 출력에 인과적으로 개입할 수 있었습니다. 예를 들어, "beard" 점수를 높임으로써 남성 이미지가 여성으로 오분류된 것을 바로잡을 수 있었습니다. 개입은 깊은 ViT 레이어와 ResNet의 얕은/깊은 레이어 모두에서 효과적이었습니다.
• 안정성 진단 (RQ4): 저자들은 깨끗한 샘 Samples와 적대적 샘플 간의 개념 점수 분포 사이의 Jensen-Shannon (JS) 거리를 사용하여 취약성을 측정했습니다. JS 거리가 높은 레이어가 더 취약한 것으로 식별되었습니다. 이러한 특정 레이어를 미세 조정(fine-tuning)하는 것이 가장 큰 개선을 가져왔으며(예: 모든 레이어를 미세 조정했을 때 ResNet-18에서 정확도 40.77% 증가), 이는 개념 표현이 모델의 취약성을 찾아내는 정밀한 진단 도구임을 입증합니다.

의의 및 주장

본 논문은 Concept-SAE를 SAE 기반의 해석 가능성을 수동적 관찰에서 능동적 심문으로 전환하는 검증된 메커니즘으로 포지셔닝합니다. 구조화된 인터페이스를 제공함으로써 연구자는 다음과 같은 작업을 수행할 수 있습니다:

• 모델이 특정 개념을 표현하는지 체계적으로 평가합니다.
• 이러한 개념들이 모델 결정에서 갖는 인과적 역할을 프로빙합니다.
• 모델의 실패와 취약점을 높은 충실도로 진단합니다.

저자들은 자신들의 접근 방식이 전통적인 SAE를 대체하는 것이 아니라 보완하는 것이며, 새로운 특징을 발견하는 능력을 유지하면서도 통제 가능하고 인간과 정렬된 해석 가능성 계층을 추가한다고 강조합니다. 결론적으로, Concept-SAE는 비전 모델을 위한 사용자 정의 개념을 평가하고 진단하는 신뢰할 수 있고, 가역적이며, 공간적으로 근거가 있는 방법을 제공합니다.