[AAAI-25] 3차원 좌표를 활용한 Denoising Distillation 모델 연구

최근 다양한 분야에서 AI 연구가 활발히 진행되고 있습니다. 특히 화학 분야와 생물학 및 제약 분야에서의 약진이 눈에 두드러집니다. 매일 뉴스에서는 쉴 새 없이 인공지능을 활용해 구조 예측을 한 사례, 원하는 물성을 만족하는 분자를 설계한 사례 등이 봇물 터지듯 쏟아져 나오고 있습니다. 이러한 융합 연구는 새로운 물질이나 시약을 개발할 때 활용되는 경우가 많으며, 기존의 방법으로 예측 혹은 시뮬레이션하기 어려웠던 과제를 비교적 정확하고 빠르게 수행할 수 있게 해줍니다. 이로 인해 연구에 소요되는 시간이 줄어들고, AI모델의 결과를 토대로 새로운 영감을 얻어낼 수도 있게 됩니다.

LG AI연구원은 LG 계열사와의 협력을 통해 소재 개발에 관한 다양한 연구 과제를 진행해왔습니다. 그 성과 중 하나로 이번에 LG생활건강과 함께 한 공동 연구에서 물질의 용해도와 안전성을 개선한 화장품 효능 소재를 개발하는 데 성공했습니다. 이 연구에는 LG AI연구원의 신물질 발굴 특화 AI 모델인 EXAONE Discovery의 역할이 컸습니다. 연구자의 경험과 논문에 의존하면 후보 물질을 선정하는 데만 2년 가까이 소요되지만, EXAONE Discovery는 단 하루면 충분했습니다.

이번 화장품 효능 소재 개발의 기반이 된 “3D Denoisers are Good 2D Teachers: Molecular Pretraining via Denoising and Cross-Modal Distillation^[1]” 연구는 AAAI (The Association for the Advancement of Artificial Intelligence)-25에서도 Oral paper로 선정되어 주목을 받았습니다. 그 연구 과정과 성과를 공개합니다.

AI와 기초과학 분야의 만남

소재 개발에 관한 과제를 진행하다 보면 서로 다른 과제임에도 불구하고 전체를 관통하는 공통된 지향점이 있다는 것을 알 수 있습니다.

1. 새로운 분포의 Data, 즉 OOD(Out-of-Distribution)에서의 예측력을 보장할 수 있는가?

2. 적은 수의 Data로 새로운 물성에 대한 모델을 개발할 수 있는가?

사실 위의 두 가지 질문은 비단 과학과 AI가 융합하는 경우가 아니더라도 AI 모델 개발 시 항상 고려하는 근본적인 질문들에 해당합니다. 이 문제들을 매우 효과적으로 극복한 연구 분야가 LLM입니다. LLM 연구는 대량의 데이터를 활용해 거대한 모델, 일명 Foundation 모델을 만들고 이를 이용해 주어진 명령을 수행합니다. 데이터가 막대하게 들어가다 보니 OOD를 맞닥뜨릴 일도 적습니다. 또한 막대한 데이터를 활용해 기 학습된 모델이 존재하기 때문에, 비교적 적은 수의 데이터를 가지고도 세부 조정을 통해 새로운 기능을 수행하는 모델 구현이 가능합니다.

LLM 분야가 이룬 눈부신 발전은 화학 및 생물학 분야 연구에서도 본받을 점이 많습니다. 특히 세부적인 요소보다는 LLM 연구가 추진한 전체 전략을 눈여겨볼 필요가 있습니다. 다만, 화학 및 생물학 분야에 AI를 적용하는 것이 생각만큼 단순하지만은 않습니다. 수천 년간 인류가 발전시켜 온 기초과학 분야는 극단적으로 간단한 수식을 이용해 자연의 섭리를 설명하려는 노력에서 비롯되기 때문에 근본적인 연구 철학 자체가 AI 분야와 상충하는 경우가 많기 때문입니다.

대표적으로, AI 모델은 근본적으로 ‘귀납적’ 접근을 전제로 합니다. 다량의 데이터를 확보하고, 이를 분석하여 설명 방법을 찾아내는 것입니다. 이 일련의 과정을 자동화하는 것이 바로 우리가 부르는 ‘AI 모델 학습’입니다. 하지만 기초과학 분야 연구는 귀납적 접근 방식으로 진행되지 않는 경우가 많습니다. 특히 화학이나 물리학 같은 분야는 더더욱 그렇습니다. 자연현상을 설명할 수 있는 방정식을 찾거나 유도하는 것으로부터 연구가 시작되며, 이 방정식을 해석하면서 새로운 물성과 특징을 찾아내는 ‘연역적’ 방법이 우위에 있기 때문입니다. 그렇기 때문에 이 두 분야의 융합은 Vision 혹은 LLM과는 다른 영역에 있으며, 아직 앞선 두 분야만큼의 파급효과를 불러오지 못하고 있습니다.

AI와 기초과학과의 융합은 AI 연구 방법론과 과학적 방법론의 장점을 조화롭게 결합하는 것이 중요합니다. AI에서 사용하는 귀납적 방법론은 많은 수의 데이터에서 규칙성을 읽어내는 만큼 방대한 기능으로의 확장을 쉽게 도모할 수 있게 만들어 주며, 전통적인 과학적 연구 결과는 모델의 정확성을 확보하는 데 활용될 수 있습니다. 이런 철학을 바탕으로 LG AI연구원 MI(Materials Intelligence) Lab은 “3D Denoisers are Good 2D Teachers: Molecular Pretraining via Denoising and Cross-Modal Distillation^[1]”을 발표했습니다. 본 연구는 3차원 좌표를 활용한 Denoising Distillation 연구로, AAAI-25에서 Oral Accept 되었습니다.

LG AI연구원의 3D Denoisers are Good 2D Teachers: Molecular Pretraining via Denoising and Cross-Modal Distillation^[1]

최근 AI를 활용한 물성 예측 모델 연구 양상은 마치 전국시대를 방불케 할 정도로 가지각색입니다. 분자 기반의 데이터를 사용하는 것은 동일하지만, 연구마다 입력 형태가 통일되지 않은 문제가 있습니다. 크게 보면 2차원 Graph 형태로 시작하는 모델과 3차원 정보를 활용하는 모델로 나눠지는데, 일반적으로 후자가 대부분의 물성 예측 분야에서 우월한 성능을 지닙니다. 데이터가 내포하고 있는 정보의 양이 풍부하기 때문입니다. 다만 분자의 3차원 정보는 AI 모델을 효과적으로 학습할 수 있는 수준의 정확성과 수를 확보하는 것이 물리적으로 불가능에 가까운 단점이 있어, 많은 경우 2차원 Graph 형태를 선호하게 됩니다.

2차원 Graph 형태는 분자의 원자 구조를 Graph 형태로 Embedding 하여 활용하는 것으로 분자의 구성만 알면 얼마든지 만들 수 있다는 장점이 있습니다. 하지만 2차원으로 표현할 수 없는 분자 구조 정보가 배제되는 문제가 있으며, 이 2차원 Graph를 Latent Vector로 변환하는 방법도 제각각입니다.

이러한 특징으로 인해 이 분야의 연구는 크게 두 방향으로 나뉘게 됩니다. 첫 번째는 대중에게 공개된 대량의 3차원 벤치마크 데이터(QM9 혹은 PCQM4Mv2 같은 데이터셋)를 활용하여 커다란 모델을 만들어내는 연구입니다. 두 번째는 구체적인 Down Stream을 결정한 상태에서 비교적 소량의 2차원 데이터를 활용해 특정 물성을 잘 예측하게 만드는 연구입니다. 두 가지 모두 각각의 장점이 명확하기 때문에 본 연구에서는 이 두 가지 방법의 장점을 융합하는 것을 목표로 삼았습니다.

일반적으로 계열사와 함께 진행하는 과제는 예측해야 하는 물성과 분자의 구조 모두 공개되어 있는 경우가 드뭅니다. 또한 비슷한 형태 혹은 비슷한 물성을 찾아내는 것조차 쉽지 않은 경우가 많습니다. 따라서 타 과제에서 개발한 모델을 기반으로 튜닝하는 것이 근본적으로 불가능해지기도 합니다.

이러한 문제를 타파하기 위해 분자의 일반적인 성질을 내포한 Latent를 산출할 수 있는 모델 개발을 1차 목표로 잡고, 이를 위해 공개된 데이터셋 중 3차원 좌표 정보를 제공하는 데이터셋을 활용했습니다. 널리 알려진 Teacher-student 모델 구조로 디자인하여 대량의 3차원 공개 데이터셋을 활용해 Teacher 모델을 학습하고, 추후에 Teacher 모델을 모사하는 Student 모델을 학습하는 전략을 적용했습니다. 이때 Student 모델은 2차원 Graph를 활용하기 때문에 3차원 정보를 잘 녹일 수 있는 2차원 정보를 추출하는 방향으로 학습합니다. 또 한 가지 장점은 향후 진행하는 과제에서 요구하는 분자의 원자구성만 알면 바로 예측할 수 있어 활용도가 높아진다는 점입니다. 이렇게 학습된 Student 모델은 추후에 원하는 소규모 물성 데이터에 미세 조정을 거쳐 최종 예측 모델로 거듭나게 됩니다.

이미지 1. Comparison between D&D and existing molecular pretraining frameworks^[1]

Teacher 모델과 Student 모델이 서로 다른 Input 형태를 활용한다는 점에서 이미 통상적인 Distillation 모델과의 차이점을 확인할 수 있습니다. Teacher 모델을 학습하는 Pretraining 알고리즘을 더 자세히 들여다보면 차이점을 극명하게 느낄 수 있습니다. Teacher 모델을 학습하는 알고리즘은 3차원 좌표의 Denoising을 목표로 합니다. 간단히 풀어보자면 주어진 분자의 구성원자의 3차원 좌표 정보를 받게 되면, 이 좌표 정보에 직접적으로 Noise를 주고 모델에게 이 Noise를 직접적으로 예측하게 하는 목표를 설정하는 것입니다. 이 일련의 과정은 아주 단순해 보이지만 과학적으로는 상당히 의미가 있습니다.

이미지 2. Illustration of our D&D framework. First we pretrain a 3D conformer encoding module by denoising perturbed conformers^[1]

통상적으로 분자를 구성하는 구성원소의 분포는 Maxwell-Boltzmann 분포를 따른다고 가정합니다. 이 Maxwell-Boltzmann 분포는 그 형태가 Exponential에 음수의 제곱형태에 비례하기 때문에 일반적으로 고려하는 Gaussian Noise의 형태와 유사합니다. 즉 구성원소의 기준좌표에서 변위를 Gaussian으로 가정하고 Noise를 생성한 후 이 Noise를 예측하는 모델을 개발하게 된다면, 이 Loss는 구성 원소 간에 작용하는 힘을 나타내는 물리량으로 해석할 수 있습니다. 이 때 Gaussian Noise는 구성 원소 간의 힘을 단순조화 운동자로 간략화하여 표현하게 됩니다. 이는 계산화학 분야에서도 구성원소의 초기좌표를 계산할 때 사용하는 공신력 있는 방법론입니다. 이와 같은 알고리즘을 학습 시에 포함시킴으로써 기초과학의 지식을 AI 모델에 성공적으로 융합할 수 있었으며, 다양한 실험에서도 그 유효성이 증명되었습니다.

이미지 3. Average performance and standard deviations for OGB and Curated datasets:
12 tasks on the left are from OGB, and 4 tasks on the right are from manually curated data.
Best results are in bold.^[1]

본 연구는 현재 주류로 연구되고 있는 Multi-modal, Distillation 및 대용량 데이터를 이용한 Foundation 모델 학습이라는 방법론에 기초과학분야의 지식인 Force-field Optimization을 효과적으로 섞어 시너지를 극대화하는 알고리즘이라는 점에서 가치가 돋보입니다. 이와 더불어 본 연구팀은 다방면의 전문 지식을 토대로 한 미분기하를 활용한 Transfer-learning모델, Group 수를 자동으로 설정해 분류해 주는 Graph-pooling 연구 등 여러 갈래로 확장하여 연구를 진행하고 있습니다. 특히 고전적인 방법인 Force-field에서 벗어나 양자적 정보를 내재할 수 있는 연구나 Teacher-Student Distillation 학습에 미분기하학의 이론을 가미한 알고리즘의 확장 등 수많은 연구가 기획 및 진행되고 있어 귀추가 주목됩니다.

기초과학분야와 AI의 융합은 연구 및 제조의 효율화에 직접적으로 기여할 수 있다는 점에서 향후 이 분야의 새로운 지평을 열 수 있을 것으로 기대됩니다. 본 연구를 수행한 LG AI연구원 MI Lab은 AI 전공자를 비롯해, 다양한 기초과학 분야의 전문 지식을 갖춘 연구원들로 구성된 조직입니다. 이러한 역량을 기반으로 전통적인 인공지능 방법론에 과학적 지식을 효과적으로 융합할 수 있는 연구를 진행하고 있습니다. 앞으로도 AI와 과학의 융합 분야에서 주도적이고 선도적인 역할을 수행해 나가고, 다양한 연구적 성과 창출은 물론 실제 비즈니스 현장에 적용하는 사례 창출에도 앞장서겠습니다.