Scaling Laws and Symmetry, Evidence from Neural Force Fields

문제 제기
본 논문은 원자 시스템의 양자 역학적 특성 (특히 퍼텐셜 에너지와 원자 힘) 을 예측하도록 설계된 딥러닝 모델인 신경망 원자간 퍼텐셜 (NNIPs) 의 스케일링 행동을 다룹니다. 자연어 처리 및 컴퓨터 비전 분야의 최근 문헌은 성능과 데이터/파라미터/연산량 간의 스케일링 법칙 (멱법칙 관계) 이 대체로 아키텍처와 무관하며, 모델이 스케일링됨에 따라 대칭성과 같은 필요한 귀납적 편향을 스스로 학습할 수 있음을 시사합니다. 그러나 기하학적 도메인에서는 이러한 관점이 이의제기를 받고 있습니다. 저자들은 명시적인 아키텍처 공변성 (회전 및 치환 대칭성 강제) 이 NNIP 의 스케일링 법칙에서 뚜렷한 이점을 제공하는지, 아니면 충분한 연산량이 주어지면 더 단순한 비공변 모델이 동등한 성능을 달성할 수 있는지 조사합니다.

방법론
저자들은 OpenMol 중성 분자 데이터셋 (약 3,400 만 개의 훈련 샘플, 약 9.2 × 10⁸ 토큰) 에 대해 포괄적인 실증 연구를 수행했습니다. 대칭성 제약의 정도가 다양한 네 가지 아키텍처 계열을 비교합니다:

1. 비제약 MPNN: 대칭성 제약 없이 기하학적 특징 (상대적 위치) 을 처리하는 기본 메시지 전달 신경망 (Message Passing Neural Network).
2. 불변 스칼라 (GemNet-OC): 불변 특징 (거리, 각도, 이면각) 을 사용하지만 가장자리 기반 메시지 전달을 통해 공변 함수를 근사합니다. 4-바디, 텐서 차수 $\ell=0$ 아키텍처로 분류됩니다.
3. 카르테시안 벡터 (EGNN): 벡터 채널 (텐서 차수 $\ell=1$) 을 사용하는 $E(n)$-공변 GNN.
4. 고차 구면 텐서 (eSEN): 회전 군의 고차 기약 표현 ($\ell \ge 2$) 을 활용하는 공변 네트워크로, 텐서 곱을 희석하기 위해 프레임 정렬을 사용합니다.

이 연구는 이론적 스케일링 법칙 문헌과 정렬하기 위해 단일 에포크 훈련 체제를 채택했으며, 학습률 스케줄 아티팩트를 완화하기 위해 스케줄러가 없는 AdamW 옵티마이저를 사용했습니다. 스케일링 법칙은 세 가지 지표를 기준으로 적합되었습니다:

• 연산량: 이론적 FLOPs ($C$) 및 실제 훈련 시간 (GPU 시간, $H$).
• 데이터: 훈련 토큰 수 ($D$).
• 파라미터: 모델 크기 ($N$).

저자들은 또한 대칭성 손실 정규화 (비공변 모델에서 공변성 편차를 패널티화), 데이터 증강을 통한 다중 에포크 훈련, 그리고 테스트 시간 그룹 평균화의 효과를 조사했습니다.

주요 기여

1. 아키텍처 의존적 스케일링 지수: 본 논문은 스케일링 지수가 아키텍처 간에 일정하지 않음을 보여줍니다. 공변성의 "정도"가 증가함에 따라 (비제약에서 저차, 고차 순으로) 데이터 ($\beta$) 및 파라미터 ($\alpha$) 에 대한 멱법칙 지수가 크게 증가합니다.
2. 우수한 공변 모델의 스케일링: 공변 아키텍처, 특히 고차 텐서 표현을 가진 모델 (eSEN 등) 은 더 가파른 스케일링 곡선을 보입니다. 이는 연산량과 데이터가 스케일링됨에 따라 공변 모델과 비공변 모델 간의 성능 격차가 벌어짐을 의미하며, 모델이 나중에 단순히 대칭성을 "학습"할 수 있다는 관념과 모순됩니다.
3. 연산 최적 할당: 연구 결과에 따르면, 연산 최적 훈련을 위해서는 모든 아키텍처에서 모델 크기 ($N$) 와 데이터셋 크기 ($D$) 가 언어 모델링 (Chinchilla 스케일링) 의 발견과 유사하게 함께 스케일링되어야 합니다 ($N \propto D$). 그러나 비례 상수와 결과적인 손실 감소는 아키텍처의 대칭성 편향에 따라 다릅니다.
4. 대칭성 손실의 비효율성: 비제약 모델에서 손실 항 (정규화) 을 통해 대칭성을 강제하는 것은 아키텍처에 공변성을 내장하는 것과 동일한 스케일링 이점을 제공하지 않습니다. 이는 데이터 효율성을 약간 향상시키지만, 고유한 공변 모델의 스케일링 지수와는 일치하지 않습니다.
5. 다중 에포크 및 증강 통찰: 저데이터, 다중 에포크 설정에서 비제약 모델은 과적합을 방지하고 멱법칙 스케일링을 회복하기 위해 데이터 증강이 필요합니다. 그러나 증강이 있더라도 비제약 모델은 공변 모델의 스케일링 지수와 일치하지 않습니다.

결과

• 스케일링 법칙: 검증 손실은 멱법칙 $L \propto C^{-\gamma}$를 따릅니다. 지수 $\gamma$는 아키텍처 복잡도에 따라 증가합니다:
  • 비제약 MPNN: $\gamma \approx 0.14$
  • EGNN: $\gamma \approx 0.17$
  • GemNet-OC: $\gamma \approx 0.25$
  • eSEN (고차): $\gamma \approx 0.40$
• 데이터 및 파라미터 스케일링:
  • 데이터 스케일링 지수 ($\beta$) 는 0.31 (비제약) 에서 0.75 (eSEN) 까지입니다.
  • 파라미터 스케일링 지수 ($\alpha$) 는 0.28 (비제약) 에서 0.82 (eSEN) 까지입니다.
• 대칭성 손실: 비제약 모델에 대칭성 손실 항을 추가하면 데이터 지수 ($\beta$) 는 약간 증가하지만 파라미터 지수 ($\alpha$) 는 감소하여, 비제약 기준선에 비해 연산 최적 프런티어 기울기에 대한 순이익이 없습니다.
• 깊이: 공변 모델의 경우 최적 네트워크 깊이는 회전 표현의 차수가 증가함에 따라 증가하는 반면, 비제약 모델은 더 높은 깊이에서 과소평화 (over-smoothing) 에 시달립니다.

의의 및 주장
본 논문은 귀납적 편향을 데이터에서 학습해야 한다는 "쓴 교훈 (bitter lesson)" 가설과 반대로, 기하학적 작업의 스케일링에는 명시적인 아키텍처 대칭성이 결정적이라고 주장합니다. 저자들은 대칭성이 단순히 데이터 축소 기술이 아니라 작업의 본질적 난이도와 그 스케일링 법칙을 근본적으로 변화시킨다고 주장합니다.

주요 의의는 고차 공변 표현이 더 나은 스케일링 지수로 이어진다는 발견에 있습니다. 이는 대규모 NNIP 의 경우, 더 단순한 비공변 모델을 스케일링하는 것보다 복잡하고 대칭성을 인지하는 아키텍처 (eSEN 등) 에 투자하는 것이 더 효과적임을 시사합니다. 저자들은 대칭성과 같은 근본적인 귀납적 편향은 모델이 발견하도록 맡기는 것이 아니라 아키텍처에 인코딩되어야 한다고 결론지으며, 이는 스케일링 궤적 자체를 변화시키기 때문입니다.

이 논문은 단일 에포크 훈련에 대한 초점, 사용된 특정 데이터셋 (중성 분자), 그리고 최근 다른 연구에서 사용된 노이즈 제거 사전 훈련 전략의 배제와 같은 한계를 지적하며 범위에 대해 겸손한 태도를 보입니다. 저자는 대칭성이 왜 스케일링 지수를 변화시키는지 설명하기 위한 향후 이론적 연구를 요청하며, 이러한 연구를 더 다양한 분자 유형과 다중 에포크 체제로 확장할 것을 제안합니다.