Expander attention as exchange-correlation

본 논문은 H₂ 및 H₄와 같은 강상관 계에 대해 높은 정확도를 달성하면서도 기존 기계학습 기반 접근법의 불리한 계산적 확장성을 극복하는, 익스팬더 그래프 트랜스포머를 기반으로 한 선형적으로 확장되는 비국소 교환 - 상관 함수를 소개한다.

핵심

혼잡한 방에서 사람들이 어떻게 행동할지 예측하려 한다고 상상해 보세요. 양자화학 세계에서는 이러한 "사람들"이 전자이고, "방"은 분자입니다.

수십 년 동안 과학자들은 이 행동을 예측하기 위해 **밀도 범함수 이론 (DFT)**이라는 도구를 사용해 왔습니다. 이는 빠르고 일반적으로 충분히 정확하기 때문에 이 분야의 "일꾼"으로 불립니다. 그러나 DFT 에는 맹점이 있습니다. DFT 는 전자를 매끄럽고 평균적인 군중처럼 취급하여, 전자가 매우 가까워지거나 "스트레스"를 받을 때 발생하는 혼란스러운 개별 상호작용을 무시합니다 (이를 강한 상관관계라고 하는 상태).

이를 해결하기 위해 DFT 는 교환 - 상관 (XC) 범함수라는 수학적 "패치"를 사용합니다. 이는 컴퓨터가 그 messy 한 개별 상호작용을 어떻게 처리해야 하는지 알려주는 규칙집과 같습니다. 문제는 그 정확한 규칙집을 아무도 모른다는 점입니다. 과학자들은 이를 추측 (근사) 해야 합니다.

문제: "비싼" 해결책

최근 연구자들은 완벽한 규칙집을 학습하기 위해 **머신러닝 (ML)**을 사용하려 했습니다. 이러한 ML 모델은 전통적인 규칙이 실패하는 messy 한 "강한 상관관계" 상황 (예: 수소 분자가 분리될 때) 을 처리하는 데 탁월합니다.

그러나 함정이 하나 있었습니다: 비용입니다.
이전 ML 모델들은 군중의 역학을 이해하기 위해 방 안의 모든 사람을 서로 소개하는 것과 같았습니다. 방이 커질수록 (원자가 늘어날수록) 이를 수행하는 데 걸리는 시간은 기하급수적으로 증가합니다. 이는 너무 느리고 비싸서 대규모 시스템에는 쓸모가 없습니다. 조각을 하나 추가할 때마다 이동 횟수가 두 배가 되는 퍼즐을 풀려고 하는 것과 같습니다.

해결책: "Exphormer"

이 논문의 저자인 Karim K. Alaa El-Din 과 옥스퍼드의 동료들은 이 규칙집을 구축하는 새로운 방법을 제안했습니다. 그들은 이를 Exphormer-XC라고 부릅니다.

그 작동 원리에 대한 간단한 비유는 다음과 같습니다:

1. 그리드: 분자가 단순히 몇 개의 원자가 아니라, 3D 이미지의 픽셀처럼 작은 점들의 거대한 3D 그리드라고 상상해 보세요.
2. 이전 방식: 이전 ML 모델들은 서로가 서로에게 미치는 영향을 보기 위해 모든 픽셀을 모든 다른 픽셀에 연결하려 했습니다. 이것이 바로 "비싼" 부분입니다.
3. 새로운 방식 (Exphormer): 모든 사람을 서로 연결하는 대신, 그들은 Expander Graph라는 수학 개념을 사용하여 스마트 네트워크를 구축했습니다.
   • 지역 친구들: 각 점은 즉시 이웃과 연결됩니다 (옆에 서 있는 사람들과 대화하는 것처럼).
   • "마법" 연결: 그들은 몇 가지 특수한 무작위 장거리 연결을 추가합니다 (방 안의 다른 모든 사람에 대해 조금씩 아는 "슈퍼 커넥터"처럼).
   • 결과: 이로 인해 모든 사람을 서로 소개할 필요 없이 정보가 방 전체로 빠르게 이동하는 네트워크가 생성됩니다. 이는 "큰 그림" 효과를 포착하면서도 복잡성을 낮게 유지합니다 (선형 확장).

그들이 테스트한 것

그들은 이 새로운 "규칙집"을 두 가지 매우 어려운 시나리오에서 테스트했습니다:

1. 수소 해리 곡선: 두 개의 수소 원자를 분리하여 끊어질 때까지 당기는 상황을 상상해 보세요. 전통적인 물리 모델은 여기서 완전히 실패하여 잘못된 에너지를 예측합니다. Exphormer 모델은 이를 정확히 예측하여 물리 계산의 "골드 스탠다드"와 거의 완벽하게 일치했습니다.
2. 평면 H4 (정사각형 수소): 이는 네 개의 수소 원자로 이루어진 정사각형입니다. 전자가 너무 혼란스러워 (축퇴) 있어 컴퓨터에게는 악몽과 같습니다. 가장 첨단 슈퍼컴퓨터 방법조차 종종 충돌하거나 잘못된 답을 내놓습니다.
   • Exphormer 모델은 전통적인 방법보다 이 시스템의 에너지를 훨씬 더 잘 예측했습니다.
   • 참고: 모델은 정사각형의 가장 혼란스러운 부분에서 "집중력을 유지하는" 데 (수렴 문제) 일부 어려움을 겪었습니다. 이는 시스템이 너무 불안정했기 때문일 가능성이 높지만, 그럼에도 불구하고 다른 모든 방법보다 우수한 성능을 보였습니다.

결론

이 논문은 다음과 같은 양자 화학용 최초의 머신러닝 모델을 구축했다고 주장합니다:

• 정확함: 전자가 이상하게 행동하는 "messy"한 상황 (강한 상관관계) 을 처리할 수 있습니다.
• 저렴함: 효율적으로 확장되므로 분자가 커질수록 기하급수적으로 느려지지 않습니다.

그들은 이를 이전에 연구하기에는 너무 비쌌던 더 크고 복잡한 시스템에 대해 고정밀 양자 시뮬레이션을 가능하게 하는 길로 규정합니다. 그들은 아직 약물 개발이나 의학적 응용에 이를 테스트하지 않았으며, 수학이 이러한 특정 수소 시스템에서 작동함을 증명하는 데만 집중했습니다.

기술 요약

기술적 요약: 익스팬더 어텐션을 교환 - 상관 함수의 관점에서

문제 제기
코른 - 샴 밀도 범함수 이론 (DFT) 은 정확도와 계산 비용 간의 균형으로 인해 전자 구조 계산의 표준으로 자리 잡았습니다. 그러나 그 실용적 유용성은 알려지지 않은 교환 - 상관 (XC) 범함수에 대한 근사에 의존합니다. 많은 밀도 범함수 근사 (DFA) 가 존재하지만, 수소 분해 곡선이나 평면 H4 와 같은 강하게 상관된 시스템에서는 어려움을 겪으며, 종종 올바른 에너지학을 포착하지 못합니다. 기계 학습 (ML) 기반 DFA 는 비국소적 상호작용을 학습함으로써 이러한 한계를 해결할 수 있는 유망한 대안으로 부상했습니다. 그러나 지속적인 병목 현상이 남아 있습니다: 강상관을 포착할 수 있는 고정밀 ML 범함수는 일반적으로 불리한 계산 스케일링 (예: $O(N^2)$ 또는 $O(N^4)$) 으로 인해 대규모 응용 분야에 적용하기에는 비용이 너무 많이 듭니다.

방법론
저자들은 익스팬더 그래프 트랜스포머 안사츠 (ansatz) 에 기반한 선형 스케일링 비국소 XC 근사인 Exphormer-XC를 제안합니다. 방법론은 다음과 같은 주요 구성 요소를 포함합니다:

1. 계산 그리드 위의 그래프 구성: 분자 그래프 (노드가 원자핵인 경우) 를 사용하는 대신, 이 접근법은 DFT 에서 사용하는 계산 전자 그리드 (Becke 그리드) 위에 직접 그래프를 구성합니다. 그래프 $G$는 그리드 점을 나타내는 정점 $V_{grid}$와 소수의 가상의 전역 노드 $V_{global}$로 구성됩니다.
2. 엣지 정의: 그래프의 선형 스케일링을 보장하면서도 연결성을 유지하기 위해 그래프 엣지는 세 가지 범주로 정의됩니다:
   • 로컬 엣지 ($E_{local}$): 해버사인 거리를 기반으로 Lebedev 쉘 내에서 가장 가까운 방사형 이웃과 각도 이웃을 연결합니다.
   • 익스팬더 엣지 ($E_{exp}$): 단순화된 프리드만 (Friedman) 방식을 사용하여 희소하면서도 고도로 연결된 그래프 구조를 생성합니다. 이를 통해 그래프는 선형 스케일링 엣지 수를 유지하면서도 큰 스펙트럼 갭 (라마누잔 기준) 을 갖게 되어 그리드 전체에 효율적인 정보 전파가 가능해집니다.
   • 전역 엣지 ($E_{global}$): 고정된 소수의 전역 저장소 노드를 모든 그리드 노드에 연결합니다.
3. 신경 아키텍처: 다층 멀티헤드 트랜스포머가 그래프를 처리합니다. 입력 노드 특징에는 전자 밀도 ($n$) 와 스핀 분극 ($\zeta$) 이 포함됩니다. 엣지 특징에는 유클리드 거리와 엣지 유형 (로컬, 익스팬더, 또는 전역) 이 포함됩니다.
4. XC 범함수 공식화: 트랜스포머는 기본 로컬 XC 에너지 밀도 $\epsilon_{XC}$에 적용되는 향상 인자 $F_{exp}$를 출력합니다. 최종 범함수는 $\tilde{\epsilon}_{XC} = \epsilon_{XC}(1 + \beta F_{exp})$이며, 여기서 $\beta$는 기본 DFA 에서의 매끄러운 전환을 보장하기 위해 0 으로 초기화되는 학습 가능한 매개변수입니다.
5. 학습 프레임워크: 모델은 FCI(Full Configuration Interaction) 데이터를 기준 (ground truth) 으로 사용하여 미분 가능한 KS 솔버 (DQC 패키지 확장) 내에서 자기 일관적으로 학습됩니다.

주요 결과
이 논문은 두 가지 벤치마크 강상관 시스템에서 Exphormer-XC 를 평가합니다:

• 수소 분해 곡선: 이 모델은 준국소 및 하이브리드 DFA 가 실패하는 영역인 H2 분자의 올바른 분해 곡선을 성공적으로 복원합니다. 다양한 기하구조 (스케일링 인자 $S=1$부터 $4.5$까지) 에 대해 학습함으로써, 모델은 보간 영역에서 1 kcal/mol 미만의 평균 절대 오차 (MAE) 를 달성합니다.
• 어블레이션 연구: 저자들은 아키텍처의 모든 구성 요소가 중요함을 보여줍니다. 구체적으로:
  • 순수 로컬 모델 (NN-LDA) 과 표준 그래프 컨볼루션은 곡선을 포착하지 못합니다.
  • 익스팬더 엣지나 거리 임베딩을 제거하면 성능이 크게 저하됩니다.
  • 전역 노드는 정확도 임계값에 도달하는 데 필수적이지는 않지만, 이를 제외하면 학습 수렴이 크게 지연됩니다 (~21%).
• 평면 H4 시스템: 이 모델은 강한 정적 상관과 준퇴화 (near-degeneracy) 로 알려진 정사각형 구성 근처의 평면 H4 에 적용됩니다.
  • 표준 DFA(예: PBE) 는 날카로운 에너지 꼭짓점을 잘못 예측하는 반면, FCI 는 포물선형 장벽을 예측합니다.
  • Exphormer-XC(비제한) 는 올바른 포물선 형태를 포착하며, 다른 DFA 들보다 FCI 에 더 가까운 에너지를 보여줍니다.
  • 한계: 모델은 준퇴화로 인해 정사각형 구성 근처에서 수렴 문제 (싱글렛과 트리플렛 상태 간의 확률적 점프) 를 보입니다. 저자들은 모델이 두 상태의 에너지학을 포착하지만, 현재 사용 중인 미분 가능한 솔버는 표준 FCI 코드에는 존재하지만 현재 미분 가능 프레임워크에는 아직 없는 대칭성 깨짐을 명시적으로 강제하여 계산을 안정화할 수 없다고 지적합니다.

의의 및 주장
이 논문은 수소 분해 곡선을 정확하게 포착할 수 있는 최초의 선형 스케일링 ML-DFA 를 제시한다고 주장합니다. 주요 기여는 Exphormer-XC아키텍처로, 이전 ML 범함수의 $O(N^2)$ 또는 그 이상의 스케일링을 강상관 시스템에 필요한 비국소성을 유지하면서 선형 스케일링($O(N)$) 으로 개선합니다.

저자들은 이 접근법이 어려운 상관 시스템에 대해 정확하면서도 대규모 적용에 충분히 계산 비용이 적게 드는 ML 범함수로 가는 길을 제시한다고 주장합니다. 단순한 그래프 토폴로지는 수렴하지 않거나 필요한 표현력이 부족하기 때문에, 익스팬더 그래프 구성이 이러한 균형을 달성하는 데 필수적임을 강조합니다. 현재 작업은 특정 테스트 시스템 (H2 및 H4) 으로 제한되며 명시적 대칭성 깨짐 없이 퇴화 영역에서 수렴 문제를 겪고 있지만, 그 결과들은 선형 스케일링 비국소 ML 범함수가 이전 고정밀 방법의 나쁜 스케일링에 대한 실현 가능한 대안임을 시사합니다.