Information Routing in Atomistic Foundation Models: How Task Alignment and Equivariance Shape Linear Disentanglement

Each language version is independently generated for its own context, not a direct translation.

이 논문은 **"분자 모델이 어떻게 세상을 이해하는가?"**에 대한 흥미로운 질문에서 시작합니다.

쉽게 말해, 인공지능(AI)이 분자의 성질을 예측할 때, 그 내부에서 **"분자가 무엇으로 만들어졌는지 (조성)"**와 **"원자들이 어떻게 배열되었는지 (기하학적 구조)"**를 구분해서 기억하고 있을까요? 아니면 이 두 가지 정보가 뒤섞여 있어서, 나중에 필요한 정보를 꺼내 쓰기 어렵게 만들어 놓았을까요?

이 연구는 이 질문에 답하기 위해 **CPD(조성 탐사 분해)**라는 새로운 방법을 개발하고, 다양한 AI 모델들을 분석했습니다. 핵심 내용을 일상적인 비유로 설명해 드릴게요.

1. 핵심 비유: "레시피 vs. 요리 상태"

분자를 생각할 때 두 가지 정보가 있습니다.

조성 (Composition): "이 요리에 소금, 설탕, 계란이 얼마나 들어갔는가?" (무엇으로 만들어졌는지)
기하학 (Geometry): "그 재료들이 어떻게 섞이고 배치되었는가?" (요리의 모양과 구조)

대부분의 분자 성질 (예: 전기가 통하는지, 빛을 어떻게 반사하는지) 은 이 두 가지가 모두 영향을 줍니다. 하지만 AI 모델이 이 두 가지를 명확히 구분해서 저장했는지, 아니면 뭉개서 저장했는지는 알 수 없었습니다.

연구의 방법 (CPD):
연구진은 AI 모델의 두뇌 (내부 표현) 에서 "조성 정보 (재료 비율)"를 수학적으로 지워버린 뒤, 남은 정보로 "기하학적 정보 (구조)"를 얼마나 잘 예측할 수 있는지 확인했습니다.

만약 지우고도 구조를 잘 예측한다면? 👉 AI 는 정보를 깔끔하게 분리해 저장하고 있다. (좋음!)
만약 지우면 구조 예측이 뚝 떨어진다면? 👉 AI 는 정보를 뒤죽박죽 섞어 저장하고 있다. (나쁨)

2. 놀라운 발견: "무엇을 배웠느냐가 가장 중요하다"

연구진은 10 가지 다른 AI 모델들을 테스트했고, 놀라운 결과를 발견했습니다. 모델의 성능 차이는 어떤 구조를 썼느냐보다 **"무엇을 학습시켰느냐"**에 따라 결정되었습니다.

비유: "요리사 (AI) 가 어떤 요리를 배웠느냐가 중요함"
- HOMO-LUMO 갭 (전자적 성질) 을 배운 요리사: 이 성질은 분자의 '모양'에 매우 민감합니다. 그래서 이걸 배운 AI 는 구조 정보를 아주 선명하게 분리해서 저장했습니다. (성공!)
- 에너지 (총량) 를 배운 요리사: 에너지는 주로 '재료의 양'에 비례합니다. 모양의 영향은 작습니다. 그래서 이걸 배운 AI 는 구조 정보를 제대로 분리하지 않고, 재료 비율에 의존해 버렸습니다. (실패!)

결론: 아무리 최신 기술 (Equivariant, 텐서 곱 등) 을 쓴 모델이라도, 학습 목표 (Task) 가 맞지 않으면 정보를 깔끔하게 정리하지 못합니다. 반면, 간단한 모델이라도 학습 목표가 맞으면 정보를 아주 잘 정리합니다.

3. 흥미로운 세부 사항: "정보의 통로"

특히 MACE라는 모델은 아주 흥미로운 방식을 썼습니다.

비유: "우편함의 분류 시스템"
- MACE 는 정보를 보낼 때 **색깔 (스칼라/벡터)**별로 통로를 나누었습니다.
- 스칼라 (숫자) 통로: 분자의 '전자적 성질 (HOMO-LUMO 갭)'을 주로 다룹니다.
- 벡터 (방향) 통로: 분자의 '쌍극자 모멘트 (전하의 방향성)'를 주로 다룹니다.
- 마치 우편물이 목적지 (성질) 에 따라 자동으로 다른 색깔의 편지통으로 분류되는 것처럼, 물리 법칙 (대칭성) 에 맞춰 정보를 정리하고 있었습니다.

하지만 다른 모델 (ViSNet) 은 이런 깔끔한 분류를 하지 않아, 모든 정보가 한 통로에 뭉쳐 있었습니다.

4. 함정: "과도한 분석의 위험"

이 연구는 또 다른 중요한 경고를 보냈습니다.

비유: "잔여물을 분석할 때 너무 복잡한 도구를 쓰면 안 됨"
- 연구진은 처음에 복잡한 도구 (GBT, 트리 기반 모델) 로 남은 정보를 분석하려 했습니다. 그랬더니, 아예 정보가 없는 상태 (조성만 남은 상태) 에서도 높은 점수를 받았습니다.
- 이유: 복잡한 도구가 "지워진 정보를 다시 추측해서 만들어냈다"는 것입니다. (마치 찌꺼기를 다시 끓여서 국물인 척하는 것과 같음)
- 해결: 그래서 연구진은 **단순한 선형 도구 (Ridge 회귀)**만 사용해야 정확한 분석이 가능하다고 결론 내렸습니다.

5. 요약: 우리가 배운 교훈

학습 목표가 왕이다: AI 모델의 구조가 아무리 복잡해도, **학습시킨 목적 (Task)**이 분자의 '구조'와 관련이 깊어야 AI 는 정보를 깔끔하게 정리합니다.
데이터의 다양성: 학습 데이터가 다양하면 (예: 다양한 분자 구조를 많이 본 경우), 학습 목표가 완벽하지 않아도 어느 정도 정보를 정리하는 데 도움이 됩니다.
정보의 분리: 좋은 AI 모델은 "무엇으로 만들어졌는지"와 "어떻게 생겼는지"를 분리해서 기억합니다. 이렇게 되면 나중에 새로운 성질을 예측할 때 훨씬 쉽고 정확하게 정보를 꺼낼 수 있습니다.
분석의 진실: AI 의 두뇌를 분석할 때 너무 복잡한 도구를 쓰면, AI 가 실제로 알고 있는 게 아니라 우리가 만들어낸 착각을 볼 수 있습니다. 단순한 것이 가장 정확할 때가 많습니다.

한 줄 요약:

"AI 가 분자를 잘 이해하려면, 복잡한 구조보다 **'무엇을 배울지 (학습 목표)'**를 잘 정하는 것이 훨씬 중요하며, 정보를 깔끔하게 정리한 모델이 나중에 더 유용하게 쓰일 수 있습니다."

Each language version is independently generated for its own context, not a direct translation.

이 논문은 **원자 기반 기초 모델 (Atomistic Foundation Models)**이 분자 속성 예측을 위해 내부 표현 (representations) 을 어떻게 조직화하는지, 특히 **화학적 조성 (composition)**과 기하학적 구조 (geometry) 정보가 어떻게 분리되거나 얽혀 있는지 분석한 연구입니다. 저자는 **조성 프로브 분해 (Compositional Probe Decomposition, CPD)**라는 새로운 방법론을 제안하고, 이를 통해 다양한 분자 모델들의 표현 조직화 차이를 정량화했습니다.

다음은 논문의 상세한 기술적 요약입니다.

1. 연구 배경 및 문제 정의

문제: 분자 모델 (MACE, SchNet, PaiNN 등) 은 분자의 에너지, 힘, 전자적 성질을 예측합니다. 그러나 이러한 모델의 중간 표현이 분자의 '구성 성분 (어떤 원소가 있는지)'과 '기하학적 구조 (원자가 어떻게 배치되었는지)'를 명확히 분리하여 인코딩하는지는 불명확합니다.
난제: 기존 프로빙 (probing) 연구는 원시 표현 (raw representations) 에 직접 학습을 수행하는데, 조성 정보와 기하학적 정보는 상관관계가 있어 구분이 어렵습니다. 또한, 잔차 (residual) 에 비선형 프로브 (예: Gradient Boosted Trees) 를 적용할 경우, 제거된 조성 신호를 비선형적으로 재구성하여 과대평가된 점수를 얻는 방법론적 오류가 발생했습니다.
목표: 조성 정보를 선형적으로 제거한 후, 남은 잔차 표현에서 기하학적 정보가 얼마나 선형적으로 접근 가능한지 (linear accessibility) 를 정확히 측정하는 방법론을 확립하고, 이를 통해 모델 간 차이를 설명하는 요인을 규명하는 것입니다.

2. 방법론: 조성 프로브 분해 (CPD)

저자는 **CPD (Compositional Probe Decomposition)**라는 새로운 프로토콜을 제안했습니다.

조성 특징 정의: 각 분자에 대해 원소 비율 (C, H, N, O, F) 과 표준화된 원자 수를 벡터 $Z$ 로 정의합니다.
선형 투영 제거 (OLS Projection): 모델의 표현 $X$ $X$ 와 조성 $Z$ $Z$ 사이의 선형 관계를 Ordinary Least Squares (OLS) 회귀를 통해 추정하고, 이를 표현에서 제거하여 **기하학적 잔차 ( $X_{geom}$ $X_{g eo m}$ )**를 생성합니다.
- $X_{geom} = X - Z\hat{\beta}$
- Fold-wise 적용: 교차 검증 (Cross-Validation) 시 데이터 누설을 방지하기 위해 각 폴드별 학습 데이터만으로 투영 계수를 계산합니다.
선형 프로빙 (Ridge Regression): 생성된 잔차 $X_{geom}$ $X_{g eo m}$ 에 대해 Ridge 회귀를 수행하여 목표 속성 (예: HOMO-LUMO 갭) 을 예측하고 $R^2_{geom}$ $R_{g eo m}^{2}$ 점수를 측정합니다.
- 핵심 발견: 비선형 프로브 (GBT 등) 는 잔차에서 제거된 조성 신호를 비선형 상호작용을 통해 재구성하여 $R^2$ 를 인위적으로 높이는 경향이 있습니다. 따라서 **선형 프로브 (Ridge)**만이 잔차에 실제로 남아있는 선형적 기하학 정보를 신뢰할 수 있게 측정합니다.
검증: 구조 이성질체 (structural isomers, 조성은 같지만 구조가 다른 분자) 벤치마크를 통해 조성 제거의 정확성을 검증했습니다. 조성 성분은 무작위 수준 (50%) 의 분류 정확도를 보인 반면, 기하학적 잔차는 94.6% 의 높은 정확도를 보였습니다.

3. 주요 결과 및 발견

QM9 데이터셋의 10 개 모델 (5 가지 아키텍처 계열) 에 CPD 를 적용한 결과, **선형 접근성 그라디언트 (Linear Accessibility Gradient)**가 발견되었습니다. 조성 제거 후 접근 가능한 기하학적 정보의 양은 모델 간에 6.6 배까지 차이가 났습니다 ( $R^2_{geom}$ : 0.081 ~ 0.533). 이 차이는 다음 세 가지 요인에 의해 설명됩니다.

3.1. 요인 1: 작업 정렬 (Task Alignment) - 가장 지배적인 요인

발견: 모델이 학습한 목적 함수 (Training Objective) 가 가장 큰 영향을 미칩니다.
- HOMO-LUMO 갭 (기하학에 민감한 속성) 으로 학습된 모델 (PaiNN, ViSNet, MACE-HL) 은 $R^2_{geom}$ 이 0.44~0.53 으로 매우 높았습니다.
- 에너지 (조성에 더 의존하는 속성) 로 학습된 모델은 $R^2_{geom}$ 이 0.08~0.31 로 낮았습니다.
인과성: PaiNN 과 MACE 아키텍처 내에서 학습 목적만 변경한 실험 (Ablation) 에서, HOMO-LUMO 학습 모델이 에너지 학습 모델보다 약 0.25~0.34 만큼 높은 점수를 기록했습니다. 이는 아키텍처보다 학습 목표가 표현의 조직화를 결정함을 의미합니다.

3.2. 요인 2: 공변성 (Equivariance) 의 조건부 역할

발견: 공변성 (Equivariance) 을 가진 모델이 항상 기하학적 접근성이 높은 것은 아닙니다.
- 에너지로 학습된 MACE(공변성 있음) 는 단순 불변 모델인 SchNet 보다 오히려 기하학적 접근성이 낮았습니다.
- 공변성 아키텍처가 높은 성능을 내기 위해서는 작업 정렬된 학습 신호가 필수적입니다. PaiNN(공변성 + HOMO-LUMO 학습) 이 가장 높은 점수를 기록했습니다.

3.3. 요인 3: 데이터 다양성 (Data Diversity)

발견: 대규모 데이터 (MPTraj 등) 로 사전 학습된 모델은 QM9 만으로 학습된 모델보다 기하학적 접근성이 높았습니다.
- 예: MACE pretrained (대규모 데이터) 는 QM9 만 학습된 MACE 보다 성능이 좋았으나, 작업이 정렬된 모델 (HOMO-LUMO 학습) 에는 미치지 못했습니다. 데이터 다양성은 불일치를 부분적으로 보완하지만 완전히 해결하지는 못합니다.

3.4. 정보 라우팅 (Information Routing)

MACE 의 특성: MACE 모델 내부에서 정보가 대칭성 유형에 따라 라우팅되는 것이 발견되었습니다.
- 스칼라 채널 (L=0): HOMO-LUMO 갭 같은 스칼라 속성을 주로 인코딩 ( $R^2=0.76$ ).
- 벡터 채널 (L=1): 쌍극자 모멘트 같은 벡터 속성을 주로 인코딩 ( $R^2=0.59$ ).
- 반면, ViSNet 은 벡터 스트림에서 거의 정보를 추출하지 못해, 이 현상이 텐서 곱 (tensor product) 기반 아키텍처의 고유한 특징임을 시사합니다.

4. 주요 기여 (Contributions)

검증된 프로빙 방법론 (CPD): 구조 이성질체 벤치마크와 비선형 프로브의 인플레이션 문제를 규명하여, 잔차 표현 분석에 선형 프로브의 필수성을 입증했습니다.
3 요인 선형 접근성 그라디언트: 10 개 모델에 대한 분석을 통해 작업 정렬이 아키텍처나 데이터 다양성보다 표현 조직화에 더 결정적임을 밝혔습니다.
대칭성 기반 정보 라우팅: MACE 모델에서 스칼라/벡터 채널이 물리적 대칭성에 맞는 속성을 선택적으로 인코딩함을 발견했습니다.
강건성 (Robustness): 12 가지 다른 조건 (투영 방법, 조성 특징 정의, 개념 삭제 등) 에서 모델 순위가 완벽하게 유지됨 ( $\rho=1.0$ ) 을 확인했습니다.

5. 의의 및 시사점

모델 선택 가이드: 분자 속성 예측을 위한 사전 학습 모델을 선택할 때, 아키텍처의 복잡성 (공변성 등) 보다는 **학습 목표 (Training Objective)**가 더 중요합니다. 기하학에 민감한 속성 (전자적 성질 등) 을 예측하려면 해당 속성으로 학습된 모델을 사용해야 합니다.
데이터 다양성의 가치: 대규모 데이터로 학습된 기초 모델은 특정 작업에 맞지 않더라도 기하학적 정보를 부분적으로 접근 가능하게 만들어 주므로, 범용 모델로서의 가치가 있습니다.
방법론적 경고: 잔차 (residual) 나 개념 삭제 (concept erasure) 후 비선형 프로브를 사용하면 신호 재구성으로 인해 잘못된 결론을 내릴 수 있으므로, 선형 프로브를 사용해야 함을 강조합니다.
표현 학습 이론: 감독 학습 (작업 정렬) 이 아키텍처적 인덕티브 바이어스 (공변성) 보다 표현의 조직화에 더 강력한 영향을 미친다는 것을 보여주었습니다.

이 연구는 분자 기초 모델의 내부 작동 원리를 해석 가능하게 만들었으며, 향후 분자 발견 및 신약 개발을 위한 모델 선택과 설계에 중요한 지침을 제공합니다.