Fine-tuning MLIP foundation models: strategies for accuracy and transferability

이 논문은 다양한 화학적 벤치마크에 걸쳐 머신러닝 기반 원자 간 포텐셜(MLIP) 파운데이션 모델을 위한 일곱 가지 미세 조정 전략을 평가하며, 파운데이션 모델의 품질과 올바른 에너지 초기화와 같은 전제 조건이 매우 중요하면서도, 단일 시스템 정확도 측면에서는 단순한 미세 조정이 최적이지만 더 넓은 배포를 위한 분포 외(out-of-distribution) 강건성을 유지하는 데는 멀티헤드 리플레이(multihead replay)가 독보적이라는 점을 밝히고 있다.

핵심

당신은 암석, 금속, 소금과 같은 무기물 재료만을 사용하여 완벽한 요리를 만드는 데 수년간 공을 들인 마스터 셰프를 상상해 보세요. 이 셰프는 "파운데이션 모델(Foundation Model)"입니다. 이제 당신은 이 셰프에게 아주 적은 양의 새로운 레시피를 사용하여 섬세한 유기물 수프나 생물학적 스튜와 같은 완전히 새로운 요리를 가르치고 싶습니다.

여기서 핵심 질문은 이것입니다: 기존의 기술을 잊어버리거나 기존의 능력을 망가뜨리지 않으면서, 어떻게 이 셰프에게 새로운 요리를 가르칠 것인가?

이 논문은 이 마스터 셰프를 "미세 조정(Fine-tuning, 재학습)"하는 일곱 가지 서로 다른 방법을 테스트하는 거대한 주방 실험입니다. 연구진은 가르치는 방법 자체보다 세 가지 결정적인 "사전 조리" 단계가 더 중요하다는 것을 발견했습니다. 즉, 적절한 셰프를 고르고, 적절한 기준점을 설정하고, 불의 세기를 조절하는 것이 더 중요합니다.

다음은 그 연구 결과에 대한 쉬나한 설명입니다:

1. 세 가지 "사전 점검" (가장 중요한 부분)

새로운 레시피를 가르치기 시작하기 전에, 반드시 세 가지를 제대로 갖춰야 합니다. 만약 이 단계에서 실수한다면, 어떤 교수법도 당신을 구원할 수 없습니다.

• 적절한 셰프 선택 (파운데이션 모델의 품질):

  • 비유: 수플레 만드는 법을 가르치기 위해 물 끓이는 법밖에 모르는 셰프를 고용하지는 않을 것입니다.
  • 연구 결과: 원래 모델의 품질이 미세 조정 전략보다 더 중요합니다. 방대하고 다양한 무기물 데이터셋(예: "OMat24" 모델)으로 훈련된 모델은 더 오래되고 작은 모델보다 새로운, 기이한 화학 물질을 배우는 데 훨씬 뛰어납니다. 설령 동일한 교수법을 사용하더라도, "더 나은" 파운데이션 모델이 항상 더 나은 최종 요리를 만들어냅니다.

• 제로 포인트 설정 (원자 참조 에너지 / $E_0$):

  • 비유: 건물의 높이를 측정한다고 상상해 보세요. 만약 지표면이 아닌 지하에서부터 측정을 시작한다면, 수치가 틀려질 것이고 건물은 떠 있거나 땅에 파묻힌 것처럼 보일 것입니다. 화학에서도 개별 원자의 "무게"를 빼주어야 모델이 오직 그들의 상호작용에 대해서만 학습할 수 있습니다.
  • 연구 결과: 연구진은 이 제로 포인트를 설정할 때 스마트하고 "모델을 인지하는(model-aware)" 방식을 사용하는 것이 매우 중요하다는 것을 발견했습니다. 만약 대충 평균값으로 짐작하여 사용한다면, 모델은 불안정해집니다. 서류상으로는 좋아 보일지 몰라도(낮은 오차 점수), 실제 물리 현상(예: 강풍 테스트 중 건물이 무너지는 상황)을 시뮬레이션하려고 하면 무너져 버릴 것입니다.

• 불의 세기 조절 (하이퍼파라미터):

  • 비유: 새로운 기술을 배울 때, 너무 빨리 움직여서 넘어지게 해서도 안 되지만, 너무 느리게 움직여서 결코 끝내지 못하게 해서도 안 됩니다.
  • 연구 결과: 서로 다른 교수법에는 서로 다른 "학습률(learning rate)"이 필요합니다. 예를 들어, 모델의 아주 작은 부분만 변경하는 LoRA 방식은 매우 빠른 학습률을 감당할 수 있는 반면, 두 가지를 동시에 가르치는 방식은 매우 느리고 부드러운 속도가 필요합니다.

2. 일곱 가지 교수 전략

위의 세 가지 점검을 통과했다면, 연구진은 새로운 레시피를 가르치는 일곱 가지 방법을 테스트했습니다.

1. 나이브 미세 조정 (Naive Fine-Tuning): "그냥 계속 요리하세요." 전체 셰프를 데리고 새로운 데이터로 계속 훈련하는 방식입니다.
   • 결과: 특정한 한 가지 요리를 완벽하게 배우는 데는 매우 좋습니다. 하지만 나중에 다른 종류의 음식을 만들려고 하면, 셰프가 예전 기술을 잊어버릴 수 있습니다(이를 "파괴적 망각"이라고 합니다).
2. 레이어 프리징 (Layer Freezing): "기본은 건드리지 마세요." 셰프의 기본적인 칼질 기술은 고정해 두고, 새로운 소스를 배우는 데만 집중하게 합니다.
   • 결과: 좋긴 하지만, 때로는 너무 경직되어 있습니다. 이는 셰프가 새로운 재료에 적응하는 능력을 제한합니다.
3. LoRA (Low-Rank Adaptation): "컨닝 페이퍼를 추가하세요." 요리책 전체를 다시 쓰는 대신, 새로운 규칙만을 담은 작고 효율적인 메모지를 셰프의 앞치마에 달아주는 방식입니다.
   • 결과: 특정 작업에 대해 매우 효율적이고 정확하며, 나이브 튜닝과 유사합니다.
4. 멀티헤드 리플레이 (Multihead Replay): "두 개의 모자를 쓴 셰프." 셰프에게 두 개의 모자를 줍니다. 하나는 새 요리를 위한 모자이고, 다른 하나는 익숙한 기존 요리를 위한 모자입니다. 셰프는 두 가지를 동시에 연습합니다.
   • 결과: 안전성 면에서 승자입니다. 기존의 기술을 잊어버리는 것을 일관되게 방지하는 유일한 방법입니다. 이 방식은 셰프가 새로운 요리와 기존 요리 모두를 잘 할 수 있게 유지해 줍니다.
5. 의사 라벨 리플레이 (Pseudolabel Replay): "합성 셰프." 실제 예전 레시피를 사용하는 대신, 셰프가 예전 레시피를 어떻게 예측하는지에 대한 자신의 예측값을 사용하여 연습합니다.
   • 결과: 기존의 데이터를 직접 보유하고 있지 않아도 셰프의 기억만 있으면 되므로 유연하고 효과적입니다.
6. 리플레이 + LoRA (Replay + LoRA): 컨닝 페이퍼와 두 개의 모자를 결합한 방식입니다.
   • 결과: 좋지만, "두 개의 모자(Dual Head)"만으로도 충분한 경우가 많았습니다.

3. 핵심 요약

• 바퀴를 다시 발명하지 마세요: 만약 당신이 특정한 좁은 작업(예: 소금물 시뮬레이션)만을 위한 모델이 필요하다면, **나이브 미세 조정(Naive Fine-Tuning)**이 가장 빠르고 쉬운 방법입니다.
• 과거를 잊지 마세요: 만약 당신이 기존의 훈련 내용을 잊지 않으면서도 기이하고 새로운 상황(예: 새로운 유형의 배터리나 복잡한 생물 분자)을 다룰 수 있는 모델이 필요하다면, 반드시 **멀티헤드 리플레이(Multihead Replay)**를 사용해야 합니다. 이것이 모델을 견고하게 유지하고 "망각"으로부터 안전하게 만드는 유일한 전략입니다.
• 기교보다는 품질: 이 논문은 완벽한 미세 조정 알고리즘을 선택하는 것보다, 고품질의 파운데이션 모델을 선택하고 에너지 참조를 올바르게 설정하는 데 시간을 쓰는 것이 더 중요하다는 점을 강조합니다. 기초가 약하거나 수학적 설정이 잘못되었다면, 세상에서 가장 좋은 교수 전략도 도움이 되지 않습니다.

요약하자면: 화학 분야에서 최고의 AI를 얻으려면, 스마트한 파운데이션에서 시작하여, 수학적 규칙을 올바르게 설정하고, 만약 AI가 다재다능하면서도 망각하지 않기를 원한다면 "두 개의 모자" 방식(Multihead Replay)으로 가르치십시오.

기술 요약

기술 요약: MLIP 파운데이션 모델의 미세 조정(Fine-tuning)

문제 정의
기계 학습 기반 원자 간 포텐셜(ML-learned interatomic potential, MLIP) 파운데이션 모델은 다양한 화학 시스템 전반에 걸쳐 전이(transfer)할 수 있는 능력을 입증하며, 특정 작업에 특화된 포텐셜을 처음부터 학습시키는 자원 집약적인 과정을 피할 수 있는 워크플로우를 제공한다. 그러나 커뮤니티에는 이러한 모델을 어떻게 그리고 언제 미세 조정해야 하는지에 대한 체계적인 지침이 부족하다. 초기 보고들은 단순한(naive) 미세 조정이 종종 "파괴적 망각(catastrophic forgetting)"을 초래한다고 시사하였으며, 이에 따라 대규모 언어 모델(LLM)에서 개발된 제약 기법(예: 레이어 동결, 저차원 적응(LoRA))의 채택을 촉구했다. 본 논문은 이러한 제약이 필수적인 것인지, 아니면 초기 실패가 더 약한 파운데이션 모델, 부적절한 원자 기준 에너지($E_0$) 초기화, 또는 불안정한 학습 절차와 같은 다른 요인들 때문이었는지를 조사한다. 본 연구의 목적은 미세 조정 결과, 특히 타겟 작업 정확도와 분포 외(out-of-distribution, OOD) 강건성을 결정짓는 주요 요인들을 규명하는 것이다.

방법론
저자들은 다섯 가지의 화학적으로 다양한 벤치마크, 세 세대의 파운데이션 모델, 그리고 다섯 차수의 크기 범위를 갖는 학습 데이터셋을 대상으로 일곱 가지의 구별된 미세 조정 전략을 평가한다.

• 평가된 미세 조정 전략:

  1. 단순(Naive): 계속적인 경사 하강법을 통한 전체 파라미터 업데이트.
  2. 레이어 동결(변형 포함): 임베딩/메시지 패싱 레이어를 동결하고 리드아웃(readout)만을 학습시키거나, 임베딩과 첫 번째 메시지 패싱 레이어를 동결함.
  3. 저차원 적응(LoRA): 사전 학습된 가중치를 동결한 상태에서 스칼라 및 등변(equivariant) 선형 레이어 모두에 학습 가능한 저차원 분해를 주입함.
  4. 멀티헤드 리플레이(Multihead Replay): 별도의 리드아웃 헤드를 사용하여 타겟 데이터와 리플레이 데이터셋(사전 학습 데이터 또는 의사 라벨링된 데이터)에서 동시에 최적화를 수행함.
  5. 의사 라벨 리플레이(Pseudolabel Replay): 리플레이 소스를 사전 학습 코퍼스에서 분리하기 위해 파운데이션 모델 자체로 생성된 리플레이 라벨을 사용하는 멀티헤드 리플레이의 변형.
  6. 리플레이 + LoRA: 멀티헤드 리플레이와 LoRA를 결합함.

• 벤치마크: 본 연구는 OMat24 사전 학습 도메인(주기적 무기 벌크)으로부터 이탈 정도가 증가하는 시스템들을 다룬다:

  • 리튬 아지로다이트 전해질(inorganic periodic solid).
  • 수용성 NaCl(ionic solution).
  • 얼음 다형체(molecular solid).
  • SN2 반응(gas-phase reactive chemistry).
  • SPICE 생체 분자(organic/biomolecular conformers).

• 기술적 구현: 저자들은 MACE 코드베이스에 세 가지 새로운 기능을 구현하였다:

  • 등변 메시지 패싱 아키텍처(스칼라 및 등변 선형 레이어 모두 포함)에 적응된 LoRA.
  • 사전 학습 코퍼스에 대한 의존성을 제거하기 위한 의사 라벨링된 리플레이.
  • 사전 학습된 베이스라인을 타겟 데이터와 정렬하기 위한 모델 인식형 원자 기준 에너지($E_0$) 재추정.

• 평가 지표: 표준적인 점 단위 에너지 및 힘 오차 외에도, 본 연구는 분자 역학(MD)에서의 방사 분포 함수(RDF), Nudged Elastic Band(NEB) 반응 프로파일, MD 안정성 테스트, 그리고 단거리 반발 실패를 감지하기 위한 무작위 구조 탐색(RSS)을 포함한 동적 및 외삽적 거동을 조사한다.

주요 결과

1. 전제 조건이 전략 선택을 지배함: 본 연구는 파운데이션 모델의 품질, 올바른 $E_0$ 초기화, 그리고 잘 선택된 하이퍼파라미터가 특정 미세 조정 전략보다 더 큰 영향을 미치는 전제 조건임을 발견했다.

   • 파운데이션 품질: 새로운 파운데션 모델(예: OMat24 기반)은 고정된 미세 조정 레시피를 사용하더라도 기존 모델(MPTraj 기반)보다 OOD 전이 성능에서 일관되게 우수한 성능을 보인다.
   • $E_0$ 초기화: "평균화된" $E_0$ 값을 사용하면 오차가 크게 증가하고 MD 불안정성(예: 얼음 모델이 50 ps 이내에 붕괴됨)이 발생한다. 사전 학습된 모델의 영점(zero-point)을 타겟 데이터에 맞추는 "재추정된" $E_0$는 안정성과 전이 가능성을 위해 필수적이며, 종종 미세 조정 알고리즘의 선택보다 더 나은 결과를 낸다.
   • 하이퍼파라미터: 단순 미세 조정은 낮은 학습률과 증가된 EMA 감쇠(decay)가 필요하다. LoRA는 더 높은 학습률을 허용한다. 멀티헤드 리플레이는 상충하는 업데이트 신호를 피하기 위해 실질적으로 더 낮은 학습률을 요구한다. 가중치 감쇠(weight decay)는 파라미터가 사전 학습된 솔루션으로부터 멀어지는 것을 방지하기 위해 0으로 설정해야 한다.

2. 목적에 따른 성능:

   • 내부 분포 특화 (단일 시스템): 좁은 범위의 작업(예: SN2 장벽, 수용성 NaCl 용매화)의 경우, 대부분의 전략(단순, LoRA, 멀티헤드)이 강력한 정확도를 달ink하며, 처음부터 학습된 모델을 일관되게 능가한다. 단순 미세 조정은 단일 시스템 애플리케이션에서 가장 좋은 수렴성을 제공한다.
   • 분포 외(OOD) 강건성: 관련이 있지만 보지 못한 조성(예: 비-아지로다이트 전해질)이나 다른 화학종(예: 생체 분자)으로의 전이를 평가할 때, 멀티헤드 리플레이(원본 또는 의사 라벨 데이터를 사용)만이 일관되게 OOD 강건성을 유지하는 유일한 접근 방식이다. 이는 타겟 작업을 학습하는 동시에 사전 학습 분포에 대한 정확도를 유지함으로써 파괴적 망실을 효과적으로 방지한다.
   • 동결 및 LoRA: 파라미터 효율성 측면에서는 효과적이지만, 테스트된 시나리오에서 용매화 특징에 적응하거나 광범위한 화학적 강건성을 유지하는 데 있어 멀티헤드 리플레이에 비해 한계를 보였다.

의의 및 주장
본 논문은 MLIP에서 단순 미세 조정이 취약하다는 인식은 방법론 자체의 내재적 한계라기보다는 부적절한 설정의 결과라는 점을 주장한다. 저자들은 다음과 같이 주장한다:

• 단순 미세 조정은 파운데이션 모델이 고품질이고 $E_0$가 올바르게 재추정된다면, 단일 시스템 애플리케이션을 위한 실행 가능하고 종종 더 우수한 시작점이 될 수 있다.
• 멀티헤드 리플레이는 미세 조정 분포 외부에서의 파운데이션 모델의 거동을 보존해야 하는 광범위한 배포를 위한 필수적인 전략이다.
• 의사 라벨 리플레이는 리플레이를 위한 구조적으로 다양한 데이터셋을 사용할 수 있게 함으로써, 리플레이 소스에 대한 의존성을 제거하는 실질적인 이점을 제공한다.

본 연구는 실무자들에게 가장 강력한 파운데이션 모델을 확보하고 원자 기준 에너지 정렬을 보장하는 것이 특정 제약 미세 조정 알고리즘을 선택하는 것보다 더 중요한 설계 선택임을 확립한다. 이 연구는 미세 조정을 틈새 옵션에서 시스템별 개발을 위한 기본 시작점으로 이동시키는 체계적인 프레임워크를 제공한다.