Dissecting Jet-Tagger Through Mechanistic Interpretability

본 논문은 입자 변환기 제트 태거에 기계적 해석 가능성 기법을 적용하여 에너지 상관관계 기반 표현과 특정 소스-릴레이-출력 아키텍처에 의존하는 희소 6-헤드 회로가 모델의 전체 분류 성능을 회복할 수 있음을 밝히고, 경사 하강법이 물리적으로 의미 있는 제트 하구조 특징을 자연스럽게 발견함을 보여줍니다.

핵심

"Dissecting Jet-Tagger Through Mechanistic Interpretability"라는 논문에 대한 설명을 일상적인 언어와 비유를 사용하여 번역한 것입니다.

큰 그림: 블랙박스 열기

상당히 숙련된 탐정 (입자 가속기에서 생성된 입자 '제트'라는 혼란스러운 범죄 현장을 분석하여 '상위 쿼크'라는 용의자인지, 아니면 단순한 배경 잡음 (QCD) 인지를 판단하도록 훈련된 컴퓨터 프로그램인 Particle Transformer) 이 있다고 상상해 보세요.

오랫동안 우리는 이 탐정이 사건을 해결하는 데 놀라울 정도로 능숙하다는 것을 알았지만, 어떻게 작동하는지는 알지 못했습니다. 그것은 '블랙박스'였습니다. 이 논문은 탐정의 두뇌를 열어, 정확히 어떤 뉴런들이 활성화되는지 매핑하고, 그들이 판결에 도달하는 데 사용하는 단계별 논리를 설명하는 법의학 팀을 고용하는 것과 같습니다.

탐정의 두뇌: 전문가 팀

연구자들은 이 탐정이 사건을 해결할 때 두뇌 전체를 사용하는 것이 아니라, 16 명 중 단 6 명의 전문가로 구성된 작고 효율적인 팀에 의존하여 작업의 97% 를 수행한다는 사실을 발견했습니다. 연구자들은 이 팀을 **'회로 (Circuit)'**라고 부릅니다.

다음은 릴레이 경기 비유를 사용하여 이 6 인 팀이 어떻게 작동하는지 설명한 것입니다:

1. 스카우트 (주요 소스): 뇌의 첫 번째 레이어에 있는 한 명의 전문가가 스카우트 역할을 합니다. 이 사람은 직접 '나쁜 놈들'을 찾지 않습니다. 대신, 그들은 '배경 잡음' (부드러운 충돌 입자) 을 찾기 위해 군중을 스캔합니다. 잡음을 이해함으로써 그들은 나머지 팀원들을 위한 무대를 마련합니다. 그들은 가장 중요한 사람입니다; 만약 그들을 제거하면 팀은 사건을 해결할 능력의 거의 대부분을 잃게 됩니다.
2. 두 번째 스카우트 (보조 소스): 첫 번째 레이어의 또 다른 한 명의 전문가가 스카우트를 돕습니다. 그들은 스카우트와 매우 유사하지만 약간 다른 세부 사항에 초점을 맞춥니다.
3. 릴레이 주자들 (중간 레이어): 중간 레이어에 있는 세 명의 전문가가 주자 역할을 합니다. 그들은 스카우트들로부터 정보를 받아 특정한 것을 찾습니다: 무겁고 에너지가 높은 입자 쌍을. 입자 물리학의 세계에서 상위 쿼크는 'W 보손'으로 붕괴된 후 두 개의 무거운 입자로 분열됩니다. 이 주자들은 이러한 무거운 쌍을 포착하는 데 능숙한 전문가들입니다.
   • 중요한 발견: 탐정은 3 부분 구조인 '상위 쿼크'를 찾아야 한다고 생각되지만, 실제로 이 주자들은 2 부분 구조인 'W 보손'만 찾고 있습니다. 논문은 탐정이 다음과 같은 단축경을 찾아냈다고 제안합니다: "만약 내가 무거운 2 부분 W 보손을 찾을 수 있다면, 그것이 상위 쿼크일 가능성이 매우 높다는 것을 확신할 수 있다." 이는 전체 범죄 현장을 재구성하려는 시도 대신 살인 무기를 찾아서 살인을 해결하는 것과 같습니다.
4. 판사 (읽기 출력): 마지막 레이어에 있는 한 명의 전문가가 판사 역할을 합니다. 그들은 직접 입자를 보지 않습니다. 대신, 그들은 릴레이 주자들로부터 보고서를 받아 요약한 후 최종 결정을 내립니다: '유죄' (상위 쿼크) 또는 '무죄' (배경).

'아하!' 순간: 새로운 아이디어가 아니라 새로운 언어일 뿐

논문에서 가장 놀라운 발견 중 하나는 탐정이 언제 결정을 내리는지와 관련이 있습니다.

보통 우리는 탐정이 레이어별로 단서를 수집하다가 마지막 순간에 갑자기 "누가 했는지 알았다!"라고 외친다고 생각합니다. 그러나 연구자들은 탐정이 실제로 거의 즉시 (첫 번째 레이어 스캔 후) 정답을 안다는 사실을 발견했습니다.

그렇다면 마지막 단계가 왜 그렇게 극적으로 보일까요?

• 비유: 탐정이 첫 번째 레이어에서 비밀 코드 (다른 언어) 로 작성된 답을 가지고 있다고 상상해 보세요. 마지막 단계는 '생각'하거나 '새로운 단서를 찾는' 것이 아니라, 단순히 그 비밀 코드를 마지막 판사가 읽을 수 있도록 번역하는 것입니다.
• 논문은 이를 **'기저 회전 (Basis Rotation)'**이라고 부릅니다. 정보는 이미 거기에 있었을 뿐, 최종 출력이 이해할 수 있도록 올바른 방향으로 회전시키기만 하면 되었습니다.

탐정은 실제로 무엇을 배웠을까요?

연구자들은 탐정이 어떤 종류의 '물리학'을 배웠는지도 확인했습니다. 그들은 탐정의 내부 메모를 인간 전문가들이 사용하는 표준 물리학 공식과 비교했습니다.

• 결과: 탐정은 인간들이 보통 사용하는 복잡한 3 부분 공식을 무시했습니다. 대신, 그것은 자연스럽게 더 간단하고 2 부분인 공식 ( Energy Correlators라고 함) 을 발견하고 선호했습니다.
• 교훈: 컴퓨터는 인간이 "W 보손을 찾아라!"라고 말해줄 필요가 없었습니다. 그것은 스스로 무거운 2 부분 붕괴를 찾는 것이 퍼즐을 푸는 가장 쉽고 신뢰할 수 있는 방법임을 깨달았습니다. 게임에서 이기려고 노력함으로써 의미 있는 물리적 진리를 재발견한 것입니다.

요약

이 논문은 고에너지 물리학에 사용되는 복잡한 현대 AI 를 역공학하여 그 안에 단순하고 논리적인 회로가 있음을 발견할 수 있음을 증명합니다.

1. 효율적입니다: 6 개의 '뉴런'으로 구성된 작은 팀이 거의 모든 작업을 수행합니다.
2. 논리적입니다: 팀은 명확한 경로를 따릅니다: 잡음 스카우트 $\rightarrow$ 무거운 쌍 릴레이 $\rightarrow$ 결과 판정.
3. 똑똑합니다: AI 는 더 큰 문제 (상위 쿼크 찾기) 를 해결하는 가장 좋은 방법이 더 간단한 하위 문제 (2 부분 W 보손 찾기) 를 해결하는 것이라고 스스로 깨달았습니다.
4. 번역입니다: AI 의 마지막 단계는 새로운 발견이 아니라, 초기의 비밀 지식을 최종 답으로 번역하는 것일 뿐입니다.

저자들은 챗봇과 같은 AI 언어 모델을 이해하기 위해 사용하는 도구들이 입자 물리학의 AI 를 이해하는 데도 완벽하게 작동하며, 이러한 기계들이 스스로 깊은 물리적 진리를 배울 수 있음을 드러낸다고 결론지었습니다.

기술 요약

기술적 요약: 기계적 해석성을 통한 Jet-Tagger 해부

문제 제기
딥러닝 아키텍처, 특히 파티클 트랜스포머 (ParT) 는 하드론성 탑 쿼크 붕괴를 QCD 배경 제트와 구별하는 것과 같은 제트 태깅 작업에서 최첨단 성능을 달성했습니다. 그러나 이러한 모델이 분류 결정을 내리는 내부 계산 메커니즘은 여전히 불투명합니다. 이전 연구들은 사후 귀속 방법 (예: Shapley 값, saliency map) 이나 어텐션 시각화를 활용했으나, 이러한 접근법들은 어떤 입력이 중요한지 식별할 수는 있지만 네트워크가 이를 어떻게 결합하는지 설명하거나 행동을 담당하는 최소 인과 서브네트워크 (회로) 를 분리해 내지는 못했습니다. 본 논문은 자연어 모델을 위해 처음 개발된 기계적 해석성의 전체 도구를 제트 물리 분류기에 적용함으로써 이 격차를 해소하는 것을 목표로 합니다.

방법론
저자들은 탑 쿼크 태깅 참조 데이터셋의 하위 집합 (신호: $t \to Wb \to q\bar{q}b$; 배경: 경쿼크/글루온) 에 대해 작은 파티클 트랜스포머 (4 개의 파티클 어텐션 레이어, 레이어당 4 개 헤드, 약 130 만 개 파라미터) 를 훈련시켰습니다. 분석에는 개입 및 프로빙 기법 세트를 활용했습니다:

1. Zero Ablation: 개별 어텐션 헤드의 출력을 0 으로 설정하여 평균 로그 확률 차이 (mean logit difference) 의 감소 정도를 측정함으로써 구조적 중요성을 평가합니다.
2. Path Patching: "깨끗한" 입력에서 특정 헤드의 출력을 "손상된" 입력으로 치환하는 인과적 개입 (배치 내 파티클 교체 또는 전체 제트 순열 사용) 입니다. 이는 헤드 간의 직접적 효과와 경로 효과 (정보 흐름) 를 분리합니다.
3. Logit Lens 및 레이어별 프로브: 표준 logit lens 는 중간 표현을 최종 훈련된 분류 헤드를 통해 투사합니다. 기저 불일치를 해결하기 위해 저자들은 각 레이어의 표현에 대해 레이어별 로지스틱 회귀 프로브를 훈련하여 클래스 정보의 진정한 선형 접근성을 결정했습니다.
4. 선형 프로빙: 잔류 스트림 (residual stream) 에서 다양한 깊이의 고전적 제트 서브구조 관측량 (예: $N$-subjettiness, Energy Correlation Functions) 을 예측하기 위해 Ridge 회귀 모델을 훈련하여 내부 표현의 물리적 내용을 특성화합니다.

주요 기여 및 결과

• 희소 6-헤드 회로의 식별: Zero ablation 과 path patching 을 통해 저자들은 전체 모델의 AUC 의 97.3% 를 회복하는 6 개 어텐션 헤드로 구성된 최소 회로를 식별했습니다. 이 회로는 무작위로 샘플링된 6 개 헤드의 부분집합보다 훨씬 성능이 우수하며 (무작위 기준 분포의 96 백분위수에 위치함) 무작위 베이스라인보다 훨씬 뛰어납니다.
• 인과적 구조 (Source-Relay-Readout): 회로는 명확한 인과적 위계를 보입니다:
  • 주요 소스 ($L0H1$): 첫 번째 파티클 어텐션 레이어의 단일 헤드로, 주요 인과적 소스 역할을 합니다. 이 헤드 단독으로 전체 모델 AUC 의 88.6% 를 회복하며 path patching 에서 "초과 회복 (super-recovery)"을 보입니다. 이는 부드럽고 콜리너 (collinear) 인 파티클 쌍을 선호적으로 어텐션하여 문맥적 정규화를 제공합니다.
  • 2 차 소스 ($L0H2$): $L0H1$ 과 유사한 표현 공간을 가지지만 구별되는 인과적 역할을 하는 같은 레이어의 헤드로, 보완적인 신호를 기여합니다.
  • 릴레이 헤드 ($L1H0, L1H1, L1H3$): 두 번째 레이어의 헤드 군집으로, 하드한 쌍별 서브구조 (높은 불변 질량, 높은 $k_T$) 를 선택적으로 어텐션합니다. 이들의 기능은 소스 헤드로부터의 상류 신호에 조건부입니다.
  • 리드아웃 헤드 ($L3H3$): 릴레이된 신호를 집계하는 네 번째 레이어의 단일 헤드입니다.
• 기저 회전 vs 정보 획득: 표준 logit lens 분석은 클래스 정보가 첫 번째 클래스 어텐션 블록 ($Cls0$) 에서만 나타나며 AUC 가 0.111 에서 0.973 으로 급증한다고 시사했습니다. 그러나 레이어별 훈련 프로브는 클래스 변별 신호가 첫 번째 파티클 어텐션 레이어 ($L1$) 에서 이미 AUC $\approx$ 0.97로 선형적으로 접근 가능함을 드러냈습니다. 따라서 $Cls0$ 에서의 극적인 점프는 새로운 정보 생성이 아니라, 잠재 신호를 최종 분류 헤드의 기저에 정렬시키는 기저 회전으로 해석됩니다.
• 물리적 내용: 3-프론트보다 우선하는 2-프론트 인코딩: 잔류 스트림을 고전적 관측량에 대해 선형 프로빙한 결과, 모델은 3-프론트 탑 태깅 작업임에도 불구하고 2-프론트 에너지 상관자 관측량 (예: $D^{(\beta=1)}_2$) 을 3-프론트 관측량 (예: $C^{(\beta=1)}_3, N^{(\beta=1)}_3$) 보다 선호적으로 인코딩하는 것으로 나타났습니다. 모델은 명시적 감독 없이도 문제를 암묵적으로 인수분해하여 전체 3-바디 위상보다는 하드론성 $W$-보손 붕괴 (2-프론트 서브구조) 를 식별하는 데 집중합니다. 이 선호도는 제트 질량에 대해 잔류화 (residualizing) 한 후에도 지속됩니다.
• 에너지 상관자 기저: 잔류 스트림은 $N$-subjettiness 기저보다 에너지 상관자 기저와 더 선호적으로 정렬되어 있으며, 이는 모델이 명시적 감독 없이 경사 하강을 통해 물리적으로 의미 있는 구조를 재발견했음을 시사합니다.

의의 및 주장
본 논문은 자연어 모델을 위해 개발된 기계적 해석성 방법이 제트 물리 분류기에 성공적으로 이전될 수 있음을 입증한다고 주장합니다. 발견된 바에 따르면:

1. 경사 하강은 아키텍처에 명시적 물리 제약이 없음에도 불구하고 제트 태깅의 물리적으로 의미 있는 측면 (특히 탑 붕괴를 $W$-보손 하위 문제로 인수분해하는 것) 을 재발견할 수 있습니다.
2. 여기서 식별된 소스 - 릴레이 - 리드아웃 회로 구조는 언어 모델의 특정 구조와 구별되는 물리 도메인 트랜스포머의 특징적인 패턴일 수 있습니다.
3. 이 연구는 운동학적으로 좁은 물리 데이터셋에서 path patching 을 위한 매니폴드 내 (on-manifold) 손상 전략의 필요성을 강조합니다. 매니폴드 외 (Gaussian) 손상은 표준 회복 점수 공식과 구조적 비호환성을 초래할 수 있기 때문입니다.

저자들은 그들의 발견의 보편성에 대해 겸손하게 언급하며, 특정 6-헤드 회로는 작은 모델에서 식별되었으며 더 큰 모델은 더 풍부한 회로 구조를 가질 수 있다고 지적합니다. 또한 선형 프로브는 비선형 인코딩이 감지되지 않기 때문에 정보 내용의 하한선을 제공한다는 점도 인정합니다.