Learning the Kohn-Sham map with neural operators for quasi-linear scaling density functional theory

60년이 넘는 시간 동안 과학자들은 원자들이 어떻게 서로 결합하여 우리 주변의 물질을 형성하는지 이해하기 위해 강력한 일련의 규칙에 의존해 왔습니다. 밀도 범함수 이론(density functional theory)으로 알려진 이 프레임워크는 원자를 둘러싼 전자 구름을 개별 입자의 집합이 아니라 매끄럽고 연속적인 밀도로 취급합니다. 이러한 접근 방식은 연구자들이 새로운 의약품부터 첨단 배터리에 이르기까지 모든 것의 특성을 예측할 수 있게 해줍니다. 하지만 여기에는 함정이 있습니다. 정확한 결과를 얻으려면 이 이론은 시스템 내의 모든 단일 전자의 거동을 계산하기 위해 복잡하고 반복적인 계산을 수행해야 합니다. 시스템이 커질수록 계산에 소요되는 시간은 폭발적으로 증가하여, 적당한 크기의 금속 조각이나 커다란 생물학적 분자를 시뮬레이션하는 것이 가장 강력한 슈퍼컴퓨터로도 불가능해집니다. 이러한 병목 현상은 많은 중요한 과학적 질문들을 손이 닿지 않는 곳에 머물게 했으며, 연구자들이 정확성과 규모 사이에서 하나를 선택하도록 강요했습니다.

캘리포니아 공과대학교(Caltech)의 한 연구팀은 컴퓨터에게 정확도를 잃지 않으면서 계산의 가장 비용이 많이 드는 부분을 건너뛰도록 가르침으로써 이 장벽을 깨는 방법을 찾아냈습니다. 최종 답을 직접 예측하거나 종종 실패하는 어려운 수학적 지름길을 배우는 대신, 그들은 인공지능에게 모든 계산 주기마다 발생하는 특정 반복 단계를 수행하도록 가르쳤습니다. 이 중간 단계에 집중함으로써, 그들은 단 하나의 그래픽 카드로 수만 개의 원자가 있는 시스템을 시뮬레이션할 수 있는 방법을 만들어냈는데, 이는 이전에는 수천 개의 프로세서를 필요로 했던 작업입니다. 그 결과, 가장 단순한 근사치만큼 빠르면서도 가장 엄격한 방법의 높은 정밀도를 유지하는 새로운 물질 모델링 방식이 탄생했습니다.

문제의 핵심은 이 이론이 전자를 다루는 방식에 있습니다. 물질의 안정적인 상태를 찾기 위해 컴퓨터는 전자 밀도 지도가 변하지 않을 때까지 반복적으로 조정해야 합니다. 전통적인 방식에서 이 조정은 거대한 고유값 문제(eigenvalue problem)를 푸는 것을 필요로 하는데, 이는 본질적으로 시스템 내 모든 원자의 전자 이동 패턴을 찾는 과정입니다. 이 연산은 계산 비용이 전자의 수에 따라 세제곱으로 증가할 만큼 매우 무겁습니다. 즉, 시스템의 크기를 두 배로 늘리면 계산은 8배 더 어려워집니다. 수십 년 동안 과학자들은 개별 전자 패턴을 추적할 필요를 없애는 방식인 궤도 없는 밀도 범함수 이론(orbital-free density functional theory)을 통해 이를 우회하려고 노력해 왔습니다. 그러나 머신 러닝을 사용하여 이러한 지름길을 학습하려는 이전의 시도들은 어려움을 겪었습니다. 어떤 이들은 전자 밀도와 에너지 사이의 관계를 학습하려 했으나, 이 관계는 수학적으로 불안정하고 오류가 발생하기 쉬웠습니다. 또 다른 이들은 원자의 배치로부터 최종 전자 밀도를 직접 예측하려 했지만, 이 모델들은 훈련된 것보다 더 크거나 화학적으로 다른 시스템을 예측하라는 요청을 받았을 때 실패했습니다.

연구진은 핵심이 최종 목적지를 예측하는 것이 아니라 여정을 배우는 데 있다는 것을 깨달았습니다. 전통적인 계산의 매 단계에서 컴퓨터는 전자들에 작용하는 힘의 지도를 가져와서 전자가 발견될 가능성이 높은 새로운 위치의 지도를 생성합니다. 이 순방향 단계는 수학적으로 안정적이고 명확하게 정의되어 있습니다. 연구팀은 공간적 패턴을 이해하도록 설계된 유형의 인공지능인 특화된 신경망을 사용하여 정확히 이 변환을 학습시켰습니다. 그들은 8,504개의 유기 화합물, 절연체, 금속을 포함한 다양한 분자와 고체 재료로부터 추출한 수천 개의 힘의 지도와 그에 대응하는 전자 밀도 지도의 사례를 네트워크에 입력했습니다. 네트워크는 힘으로부터 새로운 전자 밀도를 직접 예측하는 법을 배웠으며, 이를 통해 느리고 무거운 계산을 빠르고 학습된 예측으로 효과적으로 대체했습니다.

이 접근 방식이 독특한 점은 다양한 시스템의 복잡성을 처리하는 방식에 있습니다. 연구진은 네트워크가 새로운 크기에 대해 다시 훈련될 필요 없이 작은 분자와 거대한 결정 모두에 동일한 학습 규칙을 적용할 수 있도록 '도메인 불변성(domain-invariant)'을 갖추도록 설계했습니다. 또한 네트워크가 공간의 근본적인 대칭성을 준수하도록 하여, 분자를 회전시키거나 이동시켜도 물리적 특성이 변하지 않는다는 것을 이해하게 했습니다. 학습하지 못한 시스템에 대해 테스트했을 때, 이 모델은 놀라울 정도로 견고함을 증명했습니다. 훈련 데이터에 포함되지 않은 원소와 크기를 가진 약물 유사 분자 세트에서, 새로운 방식은 높은 정확도를 유지한 반면, 이전의 직접 예측 모델들은 분자가 커짐에 따라 오차가 급격히 치솟았습니다. 이 방식의 자기 일관적(self-consistent) 특성은 마치 항해사가 전체 경로를 추측하는 대신 자주 자신의 위치를 확인하는 것처럼, 매 단계마다 스스로의 실수를 바로잡을 수 있게 해줍니다.

이 방법의 위력은 연구진이 결정과 같은 주기적 시스템과 금속의 확장 결함(extended defects)에 적용했을 때 분명해졌습니다. 결정에 대한 표준 계산에서 컴퓨터는 결정의 반복되는 패턴을 설명하는 수학적 공간의 여러 지점에서 무거운 계산을 반복해야 합니다. 새로운 방식은 이러한 지점들을 개별적으로 샘플링할 필요 없이 단위 격자(unit cell) 상에서 전자 밀도를 진화시킴으로써 이 반복을 완전히 건너뜁니다. 이를 통해 연구진은 단 하나의 모델로 금속 시스템과 반도체의 계산을 수렴시킬 수 있었으며, 전통적인 방식과 동일한 정확도로 전자 스펙트럼과 구조적 특성을 재현해 냈습니다. 더욱이, 이 모델은 서로 다른 이론적 근사법 사이에서도 놀라운 전이 능력을 보여주었습니다. 한 종류의 계산으로 훈련된 모델은 입력되는 힘을 교체하는 것만으로도 재훈련 없이 더 정확한 다른 근사법과 함께 사용할 수 있었습니다.

최종 시험은 연구진이 금속 결정의 결함 중 하나이자 재료가 어떻게 변형되고 파괴되는지 이해하는 데 중요한 역할을 하는 마그네슘 전위(magnesium dislocation)에 이 방법을 적용했을 때 이루어졌습니다. 이러한 결함은 정확한 모델링을 위해 수천 개의 원자를 시뮬레이션해야 하는 장거리 탄성장을 생성합니다. 이러한 시스템을 높은 정밀도로 시뮬레이션하려는 이전의 시도들은 수천 개의 그래픽 처리 장치(GPU)를 갖춘 슈퍼컴퓨터를 필요로 했으며 실행에 며칠이 걸렸습니다. 현대적인 단일 그래픽 카드에서 실행된 이 새로운 방식은 8,250개의 마그네슘 원자와 82,500개의 원자가 전자를 포함하는 시스템에 대해 성공적으로 계산을 수렴시켰습니다. 계산 완료까지 걸리는 시간은 시스템의 크기에 따라 거의 선형적으로 증가하여, 계산상의 병목 현상이 제거되었음을 확인시켜 주었습니다. 결과로 얻은 전자 밀도는 고정밀 참조 모델과 구별할 수 없을 정도였으며, 이를 통해 연구진은 금속의 강도를 결정하는 다양한 결함 유형 간의 미세한 에너지 차이를 규명할 수 있었습니다.

이 연구는 머신 러닝이 전체 물리 이론을 대체하는 것이 아니라, 가장 비용이 많이 드는 반복적인 연산을 가속화함으로써 전자 구조 계산의 범위를 확장할 수 있음을 보여줍니다. 힘으로부터 밀도로 가는 순방향 지도를 학습함으로써, 연구진은 안정적이고, 다양한 유형의 물질로 전이 가능하며, 이전에는 접근 불가능했던 규모를 다룰 수 있는 도구를 만들어냈습니다. 이 방법은 계산이 실패할 경우 모델의 신뢰 범위를 벗어났음을 알리는 직접적인 진단 기능을 제공하는데, 이는 직접 예측 모델이 갖지 못한 특징입니다. 현재 모델은 전자 밀도를 예측하지만, 총 에너지와 다른 특성들은 여전히 단 한 번의 최종 단계를 통해 얻을 수 있으므로 전체 이론의 정확성을 보존합니다. 이 접근 방식은 복잡한 재료, 생물학적 시스템, 그리고 화학 반응을 과거에는 이론적으로 불가능했던 수준의 상세함과 규모로 시뮬레이션할 수 있는 문을 열어줍니다.