[논문 리뷰] Scaling Spherical CNNs

Carlos Esteves, Jean-Jacques Slotine외 1명|arXiv (Cornell University)|2023. 06. 08.

Machine Learning and Data Classification인용 수 6

한 줄 요약

이 논문은 하드웨어 가속기 최적화된 핵심 연산, 활성화 함수 및 정규화와 같은 모델 구성 요소 재설계, 그리고 애플리케이션 특화 입력 표현 방식을 통해 구면 컨볼루션 신경망(Spherical CNNs)을 확장 가능한 프레임워크로 제안한다. 이 방법은 QM9 분자의 물리적 성질 예측 벤치마크에서 최신 기술 수준의 성능을 달성하고 weahter 예측 분야에서도 경쟁력 있는 결과를 보이며, 과학적 응용 분야에서 대규모로 확장 가능한 구면 CNN이 될 수 있음을 입증한다.

ABSTRACT

Spherical CNNs generalize CNNs to functions on the sphere, by using spherical convolutions as the main linear operation. The most accurate and efficient way to compute spherical convolutions is in the spectral domain (via the convolution theorem), which is still costlier than the usual planar convolutions. For this reason, applications of spherical CNNs have so far been limited to small problems that can be approached with low model capacity. In this work, we show how spherical CNNs can be scaled for much larger problems. To achieve this, we make critical improvements including novel variants of common model components, an implementation of core operations to exploit hardware accelerator characteristics, and application-specific input representations that exploit the properties of our model. Experiments show our larger spherical CNNs reach state-of-the-art on several targets of the QM9 molecular benchmark, which was previously dominated by equivariant graph neural networks, and achieve competitive performance on multiple weather forecasting tasks. Our code is available at https://github.com/google-research/spherical-cnn.

연구 동기 및 목표

• 기존의 구면 CNN의 확장성에 한계가 있다는 점을 해결하기 위해, 높은 계산 비용으로 인해 저해상도, 얕은 모델에 국한되어 왔던 문제를 해결한다.
• 분자의 물리적 성질 예측 및 기상 예측과 같은 과학 분야에서 대규모 응용을 가능하게 하기 위해 구면 CNN을 활용한다.
• 스펙트럼 도메인의 구면 컨볼루션에서 발생하는 계산 병목 현상을 하드웨어 인식 구현과 아키텍처 혁신을 통해 극복한다.
• 실세계 과학 벤치마크에서 그래프 신경망 및 트랜스포머와 같은 최신 기술 모델을 능가하거나 동등하게 성능을 내는 구면 CNN이 가능함을 입증한다.

제안 방법

• 더 높은 표현력과 효율성을 갖춘 새로운 종류의 구면 CNN을 설계하여, 새로운 활성화 함수, 정규화 레이어, 잔차 블록을 포함한다.
• TPU 아키텍처에 맞게 최적화된 스핀 무게가 부여된 구면 조화 함수 변환을 구현하여 스펙트럼 도메인 컨볼루션의 속도를 향상시킨다.
• 3차원 원자 기하학적 구조와 회전 대칭성을 더 잘 캡처하는 새로운 분자의 구면 표현 방식을 도입한다.
• 공간 풀링 대신 스펙트럼 풀링을 적용하여 고주파 정보를 유지하고 모델 일반화 성능을 향상시킨다.
• 분리 가능한 컨볼루션과 제어된 채널 스케일링을 사용하여 기상 예측 작업에서 성능을 유지하면서 FLOPs를 줄인다.
• 정확도와 효율성에 기여하는 각 아키텍처 구성 요소의 기여도를 검증하기 위해 체계적인 아블레이션 연구를 수행한다.

Figure 1: Previous spherical CNNs were limited to low resolutions and relatively shallow models. In this work, we scale spherical CNNs by one order of magnitude and show that they can be competitive and even outperform state-of-the-art graph neural networks and transformers on scientific application

실험 결과

연구 질문

• RQ1 높은 해상도와 더 큰 모델 용량을 갖는다 하더라도, 과도한 계산 비용 없이도 구면 CNN을 확장할 수 있는가?
• RQ2 활성화 함수, 정규화, 풀링 레이어 등의 아키텍처 수정이 구면 CNN의 성능과 효율성에 어떤 영향을 미치는가?
• RQ3 기존에 그래프 신경망과 트랜스포머가 지배해 온 QM9 벤치마크에서 구면 CNN이 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ4 애플리케이션 특화 입력 표현 방식이 분자의 예측 및 기상 예측 작업에서 모델 성능 향상에 얼마나 기여하는가?
• RQ5 대규모 구면 학습 작업에서 입력 해상도, 특징 맵 크기, 모델 깊이가 예측 정확도에 어떤 영향을 미치는가?

주요 결과

• 제안된 구면 CNN 모델은 QM9의 원자화합 에너지 회귀 과제에서 테스트 MAE 15.25 meV를 기록하여 이전의 구면 CNN을 능가하고 새로운 최신 기술 수준을 수립했다.
• 크기 활성화 함수 기반 모델 대비 16.1% 향상되었고, 이전 최신 기술 수준 대비 25.1% 향상된 성능을 기록했다.
• 기상 예측 분야에서는 500 hPa에서 지위항력에 대해 RMSE 34.93 m²/s², 850 hPa에서 온도에 대해 0.62 K를 기록하여 기존 기반 모델과 동등하거나 이를 초월했다.
• 스펙트럼 풀링은 공간 풀링 대비 QM9에서 1.5% 향상된 성능을 보였고, 단계 붕괴 활성화 함수는 게이트드 비선형성 대비 1.5% 향상된 성능을 기록했다.
• 기상 예측에서 입력 해상도를 256×256에서 64×64로 낮추면 RMSE가 100% 증가하여 고해상도 입력이 모델 정확도에 매우 중요함을 입증했다.
• 32개의 TPUv4 가속기에서 학습한 모델는 4일 만에 최신 기술 수준의 성능을 달성하여, 높은 계산 비용에도 불구하고 실현 가능함을 보였다.

Figure 2: Our efficient residual block contains spin-weighted spherical Fourier transforms (FT) and inverses (IFT), multiplication with filter coefficients ( $*K$ ), activation ( $\sigma$ ) and spectral batch normalization (BN). The residual connection happens in Fourier space. Optionally, spectral