Mechanistic Interpretability and Causal Feature Steering of Neural Quantum States via Sparse Autoencoders

기술 요약: 희소 오토인코더를 통한 신경 양자 상태의 기계론적 해석 가능성 및 인과적 특징 제어

문제 정의
신경 양자 상태(Neural Quantum States, NQS)는 스핀 및 페르미온 시스템 전반에서 바닥 상태 표현과 실시간 역학 모두에서 성공을 거두며, 양자 다체 파동함수를 표현하는 강력한 변분 안사츠(ansatz)로 부상했습니다. 그러나 이러한 경험적 효능에도 불구하고, NQS는 여전히 상당 부분 "블랙박스"로 남아 있습니다. NQS는 특정 물리적 관측량에 대한 명시적인 최적화 없이 오직 변분 목적 함수(예: 에너지 최소화)만을 통해 학습되므로, 내부 메커니즘을 이해하는 데 근본적인 간극이 존재합니다. 즉, NQS가 어떻게 질서 매개변수(order parameters)나 상관 함수(correlation functions)와 같은 양들을 정확하게 포착하는지는 불분명합니다. 이러한 불투명성은 수치적 벤치마크를 얻기 어려운 영역에서 모델의 신뢰도와 강건성을 저해합니다. 기존의 접근 방식은 주로 외부적인 구조적 제약이나 파라미터 분석에 집중해 왔으며, 네트워크가 내부적으로 물리적 정보를 어떻게 조직하는지에 대한 사후적(post hoc) 분석은 부족한 실정입니다.

방법론
저자들은 기계론적 해석 가능성(mechanistic interpretability)의 도구인 **희소 오토인코더(Sparse Autoencoders, SAEs)**를 사용하여 NQS의 내부 활성화(activation)를 해석 가능한 희소 특징(sparse features)으로 분해하는 체계적인 프레임워크를 제안합니다.

1. NQS 아키텍처: 본 연구는 자기회귀 트랜스포머 기반 NQS를 활용합니다. 파동함수의 진폭 $\psi(\sigma)$는 Born 확률 $p(\sigma)$와 위상 $\phi(\sigma)$로 분해됩니다. 트랜스포머는 인과적 시퀀스 모델로서 작동하며, 사이트별로 스핀 구성을 예측합니다.
2. SAE 적용: SAE는 마지막 트랜스포머 레이어의 **잔차 스트림(residual stream)**에 적용됩니다. 이는 특징 방향의 집합인 $\{f_k\}$와 희소 계수 $z_k$로 구성된 과완전 딕셔너리(overcomplete dictionary)를 학습하며, 활성화 벡터 $h$는 $h \approx \sum z_k f_k$와 같이 희소 선형 결합으로 근사됩니다.
3. 비지도 학습: 결정적으로, SAE는 잔차 스트림 활성화에 대해 전적으로 비지도 방식으로 학습됩니다. SAE는 물리적 라벨(예: 자화 값)을 받지 않으며, 오직 재구성(reconstruction)과 희소성 패널티($L_1$ 노름)만을 최적화합니다.
4. 특징-관측량 상관관계: 학습 후, 저자들은 스핀 체인에 대해 희소 계수를 "평균 풀링(mean-pooling)"하여 스칼라 특징 강도를 얻습니다. 이 값들은 동일한 샘플링된 구성으로부터 계산된 물리적 관측량(예: 자화, 교대 자화, 반-체인 상관 함수)과 상관관계를 갖습니다.
5. 인과적 개입 (제어): 인과성을 확립하기 위해 저자들은 표적 개입을 수행합니다. 특정 특징의 활성화 강도를 재조정($z_{k^*} \to \alpha z_{k^*}$)하는 동시에 다른 특징들은 고정시키고, 수정된 잔차 스트림을 재구성하여 물리적 관측량의 변화를 관찰합니다. 이때 변분 에너지가 유지되는지 모니터링함으로써 개입이 국소적으로 이루어지도록 합니다.

주요 기여 및 결과

• 해석 가능한 특징의 발견: 1D 가로 방향 이싱 모델(TFIM)에서, SAE는 자화($M_z$), 절대 자화($|M_z|$), 반-체인 상관 함수($C_{zz}$)를 포함한 물리적 관측량과 거의 완벽하게 상관관계($|r| > 0.97$)를 갖는 희소 특징을 자율적으로 발견했습니다. 이러한 발견은 SAE가 해당 물리량에 대한 사전 지식이 없음에도 불구하고 상자성(paramagnetic) 단계, 강자성(ferromagnetic) 단계, 그리고 임계점 전반에 걸쳐 이루어졌습니다.
• 인과적 제어: 저자들은 이러한 특징들이 NQS 예측에서 인과적 역할을 수행함을 입증했습니다. 특정 특징 강도($\alpha$)를 재조정하면, 변분 에너지의 변화를 거의 일으키지 않으면서(상대적 편차 $< 0.02\%$) 해당 관측량(예: 자화)을 매끄럽고 단조롭게 제어할 수 있었습니다. 이는 해당 특징들이 단순한 상관관계의 산물이 아니라, 물리적 예측을 제어하는 내부 좌표임을 확인시켜 줍니다.
• 실시간 역학: 이 프레임워크는 TFIM 하에서의 실시간 역학으로 확장되었습니다. 단일 정적 특징 방향이 전체 진화 궤적 동안 시간 의존적 자화 $M_z(t)$를 추적하는 것이 발견되었습니다. 또한, 시간에 따른 이 특징의 활성화 강도를 조절함으로써 자화 역학을 체계적으로 변화시킬 수 있었으며, 이는 NQS가 정적인 내부 좌표계에 역학적 정보를 인코딩하고 있음을 증명합니다.
• 2D 시스템으로의 일반화: 이 접근 방식은 2D 하이젠버그 반강자성 모델(비가적분 2D 시스템)에서도 검증되었습니다. SAE는 교대 자화(sublattice를 구분하는 공간적 구조를 가진 질서 매개변수)와 강하게 상관된 특징을 식별했습니다. 이 특징을 제어함으로써 교대 자화를 성공적으로 조절할 수 있었으며, 이는 본 방법론이 1D 및 가해 가능한(exactly solvable) 시스템을 넘어 적용될 수 있음을 보여줍니다.

의의 및 주장
본 논문은 NQS가 단순히 파동함수 진폭의 기능적 근사치가 아니라, 풍부하고 해석 가능한 물리적 정보의 내부 표현을 인코딩하고 있다고 주장합니다.

• 진단 도구: 이 접근 방식은 훈련된 NQS가 물리적으로 타당한 표현을 학습했는지 평가하는 방법을 제공합니다. 질서 매개변수를 인과적으로 제어하는 특징의 존재는 성공적인 학습의 증거가 되며, 반대로 특징의 부재는 모델 실패를 진단하는 지표가 될 수 있습니다.
• 개입 도구: 본 연구는 제어된 특징 수준의 조작을 통한 NQS의 제어 경로를 구축합니다. 잔차 스트림 내의 특정 스칼라 계수를 조정함으로써, 전체 네트워크를 재학습하거나 미세 조정하지 않고도 특정 물리적 관측량을 튜닝할 수 있습니다.
• 양방향 교환: 저자들은 양자 다체 물리학과 기계론적 해석 가능성 사이의 시너지 관계를 상정합니다. NQS는 해석 가능성 주장을 테스트할 수 있는 정답(ground-truth) 물리 관측량을 제공하며, 반대로 해석 가능성 도구는 변분 파동함수의 "블랙박스"를 열어 더 신뢰할 수 있고 투명한 양자 시뮬레이션을 가능하게 합니다.

저자들은 본 연구가 트랜스포머 아키텍처와 최종 레이어의 잔차 스트림에만 국한되어 있음을 언급하며, 다른 아키텍처(예: RBM, RNN)나 레이어별 셀프 어텐션 메커니즘에 대한 질문을 남겨두며 범위를 신중하게 제한하였습니다.