How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

기술적 요약: 등변 신경망에서 최적화기가 학습된 해를 형성하는 방식

문제 제기
등변 신경망은 기하학적 대칭성을 아키텍처에 직접 인코딩하여, 알려진 기하학적 대칭성을 가진 데이터에 대한 자연스러운 귀납적 편향을 제공합니다. 그러나 강력한 등변성 제약은 종종 최적화를 어렵게 만들어, 임계점과 가짜 최소값을 포함하는 손실 지형을 초래합니다. 결과적으로 이러한 모델은 확장 시 제약이 덜한 대안들보다 성능이 떨어질 수 있습니다. 최근 문헌에서의 주된 대응은 근사적이거나 완화된 등변성 방법을 통해 아키텍처 제약을 완화하는 것이었습니다. 반면, 등변 모델의 해를 형성하는 데 있어 최적화기의 역할은 아직 충분히 탐구되지 않았습니다. 본 논문은 아키텍처 수정과 무관하게 최적화기 선택이 훈련 중 도달하는 해를 크게 변화시키고 성능을 향상시킬 수 있는지 조사합니다.

방법론
저자들은 Adam(Kingma & Ba, 2015) 과 뉴턴 - 슈르츠 반복을 통해 모멘텀 버퍼를 직교화하는 최근 제안된 최적화기 Muon(Jordan et al., 2024) 간의 비교 연구를 수행합니다. 평가는 세 가지 서로 다른 설정에서 이루어집니다:

1. 3D 포인트 클라우드 분류: 서로 다른 기하학적 귀납적 편향을 가진 세 가지 아키텍처를 사용하여 ModelNet40 데이터셋 (및 그 손상된 변형인 ModelNet40-C) 을 활용합니다:
   • EGNN: 명시적 $E(3)$-등변 그래프 신경망.
   • DGCNN: 동적 $k$-NN 연결성을 가진 치환 등변 그래프 CNN.
   • PointNet: 대칭 풀링을 통한 치환 불변 아키텍처.
2. 분자 학습: QM9 데이터셋을 사용하여 GotenNet($E(3)$-등변) 과 GINE(치환 등변) 을 펩타이드 및 ZINC 데이터셋에 적용합니다.
3. 기하학적 분석: 성능 차이가 발생하는 원인을 이해하기 위해 훈련된 체크포인트를 다음과 같이 분석합니다:
   • 국소 손실 기하학: 체크포인트 주변의 2D 손실 슬라이스와 3D 표면을 시각화하고, 파워 반복 및 허친슨 추정기를 통해 헤시안 요약 (최대 고유값 및 트레이스) 을 추정합니다.
   • 스펙트럼 구조: 가중치 행렬과 중간 레이어 표현 모두에 대해 안정적 순위(stable rank) 와 유효 순위(effective rank, 특이값의 섀넌 엔트로피 기반) 를 계산합니다.

주요 결과

• 성능 향상: ModelNet40 에서 Muon 은 깨끗한 설정과 손상된 설정 모두에서 모든 아키텍처 (EGNN, PointNet, DGCNN) 에 걸쳐 Adam 일관적으로 우월한 성능을 보입니다. 예를 들어, 깨끗한 ModelNet40 에서 Muon 은 EGNN 정확도를 +5.17%, PointNet 을 +2.67%, DGCNN 을 +1.96% 향상시켰습니다. 유사한 경향이 QM9 분자 데이터셋에서도 관찰되어, Muon 이 대부분의 목표 속성을 향상시켰습니다.
• 손실 지형 특성: 훈련된 체크포인트 주변의 손실 표면 시각화는 Muon 해가 Adam 에서 자주 관찰되는 불규칙한 등고선과 비교하여 더 매끄럽고 규칙적인 국소 기하학을 보임을 보여줍니다. 그러나 헤시안 추정은 더 매끄러운 기하학이 더 낮은 곡률을 의미한다는 직관과 모순됩니다. Muon 체크포인트는 실제로 Adam 체크포인트보다 더 큰 최대 고유값과 트레이스 추정값 (더 높은 곡률 표시) 을 가집니다.
• 스펙트럼 특성: Muon 은 모든 레이어와 아키텍처에 걸쳐 가중치와 중간 표현 모두에 대해 더 높은 안정적 및 유효 순위를 일관되게 산출합니다. 이는 Muon 이 스펙트럼이 몇 가지 지배적인 방향으로 집중되는 경향이 있는 경사 기반 방법과 달리, 용량을 특이 방향들 사이에서 더 균일하게 분배함을 시사합니다.

의의 및 주장
본 논문은 최적화기가 기하학적 제약과 상호작용하는 의미 있는 귀납적 편향의 원천이라고 주장합니다. 저자들은 다음과 같이 주장합니다:

1. 최적화기 선택의 중요성: 아키텍처를 수정하거나 제약을 완화하지 않고도 Adam 에서 Muon 으로 최적화기를 변경하면 등변 모델의 성능을 향상시킬 수 있습니다.
2. 구별된 해 공간: Muon 과 Adam 은 서로 다른 국소 손실 기하학과 스펙트럼 구조를 특징으로 하는 질적으로 다른 해로 수렴합니다.
3. 메커니즘 가설: Muon 의 직교화된 모멘텀 업데이트는 손실 지형의 더 유리한 영역으로의 궤적 수준의 편향을 제공하며, 업데이트의 특이 방향들을 더 균일하게 처리하는 것으로 보입니다.
4. 설계 원칙: 결과는 등변 최적화기를 위한 잠재적 설계 원칙을 시사합니다. 가중치와 표현의 스펙트럼이 몇 가지 방향으로 집중되는 것을 방지하는 업데이트 규칙은 훈련을 개선할 수 있습니다.

저자들은 스펙트럼 특성과 정확도 향상 사이의 상관관계는 일관되지만 인과 관계는 입증되지 않았음을 지적하며 겸손한 입장을 유지합니다. 그들은 최적화기 설계와 기하학적 귀납적 편향 간의 상호작용, 특히 최적화기가 등변성을 블랙박스 제약으로 취급하는 대신 대칭군의 구조를 어떻게 활용할 수 있는지에 관한 추가 조사가 필요하다고 결론지었습니다.