Machine-learned, finite temperature Fermi-operator expansions suitable for GPUs and AI-hardware

본 논문은 SP2 방식에 기반하여 전자 구조 계산을 심층 신경망 아키텍처로 매핑하는 기계학습 강화 유한 온도 재귀적 페르미 연산자 전개 방법을 제시하며, 최적화된 행렬 - 행렬 곱셈과 아핀 재스케일링을 활용하여 명시적 대각화와 모델 재학습을 회피함으로써 GPU 및 AI 하드웨어에서 수배의 속도 향상을 가능하게 한다.

핵심

이 논문은 간단한 언어와 일상적인 비유를 사용하여 설명합니다.

큰 그림: 원자를 시뮬레이션하는 더 빠른 방법

방 (물질) 안에서 사람 (전자) 무리가 어떻게 움직이고 상호작용할지 예측하려고 한다고 상상해 보세요. 양자 물리학의 세계에서는 이것이 매우 어렵습니다. 정확한 답을 얻으려면 보통 '대각화 (diagonalization)'라는 거대하고 복잡한 퍼즐을 풀어야 합니다.

대각화를 생각해보면, 올바른 순서를 찾기 위해 모든 책의 모든 페이지를 읽으며 백만 권의 책을 분류하려는 것과 같습니다. 이는 정확하지만, 특히 방이 커질수록 시간이 매우 오래 걸립니다.

이 논문의 저자들은 **단순한 방법 (shortcut)**을 개발했습니다. 모든 페이지를 읽는 대신, 책을 거의 즉시 분류하는 방법을 배우는 '스마트 추측' 기계를 만들었습니다. 그들은 이를 **기계 학습 페르미 연산자 전개 (Machine-Learned Fermi-operator expansion)**라고 부릅니다.

문제: 뜨겁고 차가운 군중

과거에는 이러한 단순한 방법들이 '군중'이 매우 차가울 때 (절대 영온) 만 잘 작동했습니다. 차가운 군중에서는 모두가 매우 예측 가능한 줄에 서서 가만히 있습니다. 수학은 단순합니다. 줄에 있거나 없거나 둘 중 하나일 뿐입니다.

그러나 현실 세계에서는 종종 상황이 '뜨겁습니다'. 전자가 뜨거워지면 불안정해집니다. 줄에 서 있던 사람들이 나가고, 기다리던 사람들이 들어올 수 있습니다. 이는 사람들이 부분적으로 안에 있고 부분적으로 밖에 있는 '흐릿한' 경계를 만들어냅니다.

이전 단순한 방법들은 여기서 실패했습니다. 너무 경직되어 있었기 때문입니다. 그들은 뜨겁고 흐릿한 군중을 처리할 수 없었습니다.

해결책: 신경망에 '압축 (Squash)'을 가르치기

저자들은 차가운 군중을 분류하는 데 사용된 수학이 **심층 신경망 (Deep Neural Network)**의 구조와 정확히 일치한다는 사실을 깨달았습니다. (얼굴을 인식하거나 시를 쓰는 데 사용되는 AI 유형입니다.)

• 구식 방법 (SP2): 기계가 숫자를 받아 제곱 ($x^2$) 하거나 특정 뺄셈 ($2x - x^2$) 을 수행한다고 상상해 보세요. 이 과정을 반복하여 숫자들을 '압축'하여 0 또는 1 이 되도록 만듭니다. 이는 차가운 군중에게는 훌륭하게 작동합니다.
• 신식 방법 (MLSP2): 저자들은 이 기계에 '두뇌'를 부여했습니다. 고정된 규칙을 사용하는 대신, 기계 학습을 통해 기계를 훈련시켰습니다. 그들은 기계가 내부의 조절 장치 (계수) 를 스스로 조정하여 '흐릿한' 뜨거운 군중을 완벽하게 처리할 수 있도록 가르쳤습니다.

이렇게 생각해보세요:

• 구식 기계: '예' 또는 '아니오'만 찍는 경직된 도장.
• 신식 기계: 군중이 얼마나 뜨거운지에 따라 그 사이를 부드럽고 완벽한 곡선으로 만드는 '예'와 '아니오'를 정확히 어떻게 성형할지 배우는 유연한 3D 프린터.

마법 같은 트릭: 하나의 모델이 많은 온도에 적합함

보통 시뮬레이션의 온도를 변경하면 AI 모델을 처음부터 다시 훈련시켜야 합니다. 이는 영원히 걸리는 일입니다.

저자들은 **아핀 리스케일링 (Affine Rescaling)**이라는 교묘한 트릭을 발견했습니다.
도시 지도가 있다고 상상해 보세요. 확대하거나 축소하고 싶다면 도시 전체를 다시 그릴 필요가 없습니다. 지도를 늘이거나 줄이기만 하면 됩니다.

저자들은 특정 '확대/축소 수준' (특정 온도와 화학 퍼텐셜) 에 대해 AI 모델을 단 한 번만 훈련시킬 수 있음을 발견했습니다. 그런 다음, 일정 범위 내의 다른 온도에 대해서는 입력 데이터 (해밀토니안 행렬) 를 모델에 넣기 전에 단순히 '늘려주기'만 하면 됩니다. 모델은 다시 배우지 않아도 됩니다. 데이터가 약간 다른 규모로 보일 뿐이며, 올바른 답을 내놓습니다.

이는 온도가 끊임없이 변하는 시뮬레이션 (예: 화학 반응) 을 AI 를 다시 훈련시키지 않고 실행할 수 있음을 의미합니다.

하드웨어: 과학을 위한 AI 칩 사용

이 논문은 이 방법이 특히 GPU(그래픽 처리 장치) 와 텐서 코어(AI 를 위해 설계된 칩) 와 같은 현대 컴퓨터 칩에 맞춰 구체적으로 구축되었음을 강조합니다.

• 비유: 전통적인 대각화는 숙련된 목수가 가구 하나하나를 손으로 정교하게 조각하는 것과 같습니다. 정밀하지만 느립니다.
• 신식 방법: 이는 고속 3D 프린터를 사용하는 것과 같습니다. AI 칩의 특정 아키텍처를 활용하여 대규모 계산 (행렬 곱셈) 을 놀라울 정도로 빠르게 수행합니다.

저자들은 Nvidia RTX 6000 Ada GPU 에서 이를 테스트했습니다. 그들은 이 방법이 오늘날 과학자들이 사용하는 표준적이고 최적화된 방법보다 9 배에서 16 배까지 빠르며, 여전히 높은 정확도를 유지한다는 사실을 발견했습니다.

결과 요약

1. 속도: 특히 현대 AI 하드웨어에서 물질 내 전자의 행동을 계산하는 속도가 대폭 향상되었습니다 (최대 16 배).
2. 정확도: 이전의 단순한 방법들이 제대로 처리하지 못했던 '뜨거운' 전자 (분수 점유) 를 극도로 정밀하게 모델링할 수 있습니다.
3. 효율성: 모델을 한 번 훈련시키고 수학적 트릭을 사용하여 입력을 재조정함으로써, 시뮬레이션 중 온도가 변할 때마다 모델을 다시 훈련시킬 필요가 없습니다.
4. '마법' 같은 대각화 부재: 그들은 느리고 무거운 대각화 수학을 완전히 피하고, 대신 AI 칩이 매우 좋아하는 빠르고 반복적인 곱셈 단계에 의존합니다.

요약하자면, 저자들은 느리고 경직된 수학적 과정을 현대 컴퓨터 칩에서 매우 효율적으로 실행되는 빠르고 유연하며 AI 기반의 도구로 변환하여, 과학자들이 이전보다 훨씬 빠르게 복잡한 물질을 시뮬레이션할 수 있게 했습니다.

기술 요약

기술 요약: 기계학습 기반 유한 온도 페르미 연산자 전개

문제 제기
전자 구조 계산, 특히 코언-샴 밀도 범함수 이론 (KS-DFT) 내에서는 해밀토니안 행렬을 대각화하여 고유값 문제를 풀기 위한 입방 (cubic) 스케일링 비용으로 인해 계산적 한계에 직면해 있습니다. 2 차 스펙트럴 프로젝션 (SP2) 방법과 같은 재귀적 페르미 연산자 전개 방식은 대각화 없이 직접 밀도 행렬을 계산할 수 있는 방법을 제공하지만, 기존 효율적 구현은 영 (zero) 전자 온도로 제한됩니다. 영 온도에서 밀도 행렬은 멱등성 (occupations 이 엄격하게 0 또는 1) 을 가집니다. 그러나 금속이나 고온 전자 상태를 가진 시스템과 같은 많은 물리적 시스템은 축퇴된 고유상태나 열적 확산을 정확하게 모델링하기 위해 분수 궤도 점유를 필요로 합니다.

이전에는 유한 온도로 SP2 를 일반화하기 위해 재귀를 절단하여 열적 확산을 도입하려는 시도가 있었습니다. 그러나 이러한 절단된 전개는 본질적으로 근사적이어서, 특히 정확도가 중요한 화학적 전위 근처에서 정확한 페르미 함수를 재현하지 못합니다. 체비셰프 전개나 파데 근사법과 같은 대안적 방법들은 깁스 진동을 억제하기 위해 과도하게 높은 다항식 차수를 요구하거나, 반복적인 선형 시스템 풀이로 인해 상당한 계산 오버헤드를 초래합니다.

방법론
저자들은 재귀적 SP2 방법을 유한 온도로 일반화하는 프레임워크를 제안하며, 이를 위해 그 대수적 구조를 심층 신경망 (DNN) 아키텍처에 매핑합니다. 핵심 통찰력은 재귀적 SP2 업데이트가 신경망의 레이어와 유사하다는 점입니다. 전개 계수를 학습 가능한 가중치와 편향으로 취급함으로써, 저자들은 임의의 온도에서 분수 점유를 가진 페르미 분포 함수를 근사할 수 있는 머신러닝 모델을 구축합니다.

주요 방법론적 구성 요소는 다음과 같습니다:

1. 신경망 아키텍처:

   • MLSP2 (기계학습 SP2): SP2 의 일반화로, 2 차 업데이트 규칙 ($X^2$ 또는 $2X-X^2$) 을 누적항이 있는 학습 가능한 2 차 다항식 ($ax^2 + bx + c$) 으로 대체합니다. 이를 통해 모델은 절단된 계단 함수가 아닌 페르미 함수의 정확한 열적 확산을 근사할 수 있습니다.
   • Max-SP2: '스킵 연결 (skip connections)'을 포함하는 더 표현력 있는 아키텍처로, 각 레이어는 모든 이전 레이어의 선형 결합의 제곱입니다.
   • Skip-SP2: 표현력과 메모리 사용량을 균형 있게 조절하기 위해 최근 레이어와 누적항의 유한 메모리를 사용하는 Max-SP2 의 압축 버전입니다.

2. 엔트로피 근사:
   저자들은 또한 전자 자유 에너지를 계산하는 데 필요한 전자 엔트로피 함수 $s(x)$를 근사하기 위한 재귀적 방식을 개발합니다. 이들은 페르미 함수와 그 여집합의 곱인 $f(x)(1-f(x))$를 스케일링한 것을 초기 추측값으로 활용하며, 이는 화학적 전위에서 실제 엔트로피의 2 차 도함수와 일치하도록 훈련된 재귀적 2 차 전개를 통해 정제됩니다.

3. 훈련 및 최적화:
   모델은 전체 행렬 대신 단위 구간 $[0, 1]$ 내의 스칼라 입력에 대해 훈련되며, 측지선 가속 (geodesic acceleration) 을 적용한 레번버그 - 마쿼트 알고리즘을 사용합니다. 훈련 데이터는 화학적 전위 근처의 최대 오차를 최소화하기 위해 페르미 함수의 도함수에 비례하는 가중치로 샘플링됩니다.

4. 아핀 재스케일링 및 전이성:
   중요한 혁신은 시뮬레이션 매개변수가 변경될 때 재훈련의 필요성을 제거하기 위한 아핀 재스케일링의 사용입니다. 해밀토니안 ($H'$), 화학적 전위 ($\mu'$), 그리고 역온도 ($\beta'$) 를 정규화함으로써, 특정 매개변수 $(\beta_0, \mu_0)$에서 훈련된 단일 모델은 다른 매개변수의 넓은 '유효 영역'에 적용될 수 있습니다. 이는 훈련 조건과 일치하도록 입력 해밀토니안을 재스케일링함으로써 달성되며, 이를 통해 시뮬레이션 중 다양한 온도와 화학적 전위에 걸쳐 동일한 가중치 세트를 사용할 수 있습니다.

5. 하드웨어 구현:
   알고리즘은 현대 GPU 와 AI 하드웨어 (특히 NVIDIA Tensor Cores) 에 최적화되어 있습니다. 저자들은 행렬 제곱 연산을 효율적으로 수행하기 위해 혼합 정밀도 산술 (FP16/FP32) 을 활용하며, 해밀토니안의 대칭성을 이용하여 필요한 곱셈과 데이터 전송 횟수를 줄입니다.

주요 결과

• 정확도: MLSP2 모델은 페르미 함수 근사에 대해 $10^{-7}$ 수준의 오차를 달성하여, 절단된 SP2 방법 (오차 약 $10^{-2}$) 보다 훨씬 우수하며 많은 영역에서 이중 정밀도 대각화의 정밀도와 일치합니다.
• 성능: NVIDIA RTX 6000 Ada GPU 에서 MLSP2 접근법은 중간 크기 행렬에 대해 이중 정밀도 대각화 (cuSOLVER 사용) 대비 16 배, 더 큰 행렬에 대해 9 배의 속도 향상을 보입니다. 단일 정밀도 대각화와 비교하더라도 MLSP2 는 우수한 안정성과 정확도를 유지하면서 2 배에서 5 배의 속도 향상을 제공합니다.
• 확장성: 이 방법은 명시적 대각화를 피하고 고도로 최적화된 행렬 - 행렬 곱셈 커널에만 의존합니다. 목표 정확도를 달성하는 데 필요한 레이어 수는 역온도 ($\beta$) 에 대해 로그적으로 스케일링되어 저온에서도 효율적인 계산을 가능하게 합니다.

의의 및 주장
이 논문은 대각화의 계산 병목 현상을 피하는 유한 온도 전자 구조 계산을 위한 견고하고 일반화 가능한 솔루션을 제공한다고 주장합니다. 머신러닝을 통해 SP2 를 일반화함으로써, 저자들은 전통적 방법의 일부 비용으로 분수 점유를 가진 시스템의 밀도 행렬 계산을 가능하게 합니다.

의의는 화학적 전위와 전자 온도가 시간 단계 간에 변동하는 동적 유한 온도 시뮬레이션 (양자 분자 동역학 등) 을 수행할 수 있는 능력에 있습니다. 아핀 재스케일링 전략은 단일 사전 훈련된 모델을 시뮬레이션 전체에 걸쳐 재훈련 없이 재사용할 수 있도록 하여 대규모 응용에 실용적으로 만듭니다. 또한, 이 접근법은 현대 AI 하드웨어 (Tensor Cores) 의 성능 특성을 활용하도록 특별히 맞춤화되어, 높은 수치적 정확도를 유지하면서 벤더 최적화 대각화 루틴보다 상당한 속도 향상을 제공합니다.