실용적인 깊은 이분산 회귀와 불확실성 추정

초록

본 논문은 대규모 그래프 신경망을 이용한 분자 특성 예측에서, 평균 예측 정확도는 유지하면서 입력‑의존적 불확실성을 효과적으로 추정하기 위한 사후‑학습(포스트‑핵) 방법을 제안한다. 기존의 엔드‑투‑엔드 이분산 회귀가 겪는 최적화 난이도, 표현 붕괴, 잔차 과적합 등의 네 가지 핵심 문제를 분석하고, 사전 학습된 백본의 중간 레이어를 활용해 선형 분산 헤드를 별도로 학습함으로써 이를 동시에 해결한다. 실험 결과, 제안 방법은 QM9 등 여러 분자 데이터셋에서 NLL 기준으로 최첨단 성능을 달성하면서도 평균 절대오차(MAE)는 기존 비이분산 모델과 동등하거나 더 우수하였다.

상세 분석

논문은 먼저 깊은 이분산 회귀(mean‑variance network)가 직면하는 네 가지 근본적인 문제를 체계적으로 정리한다. 첫 번째는 NLL 손실의 그래디언트 구조 때문에 예측 분산이 크게 증가하면 평균과 분산 모두의 업데이트가 급격히 감소하는 최적화 병목이다. 기존 연구들은 β‑NLL, 단계적 학습, 정규화 재파라미터화 등으로 완화하려 했지만, 대부분 추가 하이퍼파라미터 튜닝을 요구하거나 평균 정확도를 손상시킨다. 두 번째 문제는 ‘표현 붕괴’로, 평균을 최적화하는 과정에서 입력 공간의 일부 방향(특히 분산만을 설명하는 방향)이 백본의 잠재 표현에서 사라져, 분산 추정에 필요한 정보를 잃게 된다. 이는 특히 입력‑의존적 불확실성을 직접 모델링해야 하는 이분산 회귀에서 심각하다. 세 번째는 잔차 과적합이다. 과대 파라미터화된 네트워크는 훈련 데이터에 거의 완벽히 맞춰지면서 잔차가 훈련 집합에만 최적화되고, 보정용 홀드아웃 데이터가 없으면 일반화된 불확실성 추정이 불가능해진다. 네 번째는 실용성 측면으로, 대규모 그래프 신경망에 적용할 때는 평균 성능 유지, 하이퍼파라미터 최소화, 추론 시 연산 비용 절감, 구현 난이도 낮춤이 필수이다.

제안된 방법은 이러한 문제들을 한 번에 해결한다. 핵심 아이디어는 사전 학습된 백본(평균 예측기)을 고정하고, 별도의 선형 분산 헤드를 홀드아웃 데이터에만 학습하는 포스트‑핵 접근이다. 여기서 중요한 점은 마지막 레이어뿐 아니라 여러 중간 레이어의 잠재 표현 (z_l) 를 입력으로 사용한다는 점이다. 이는 표현 붕괴를 완화시키며, 각 레이어가 평균과 분산 양쪽에 기여할 수 있게 만든다. 선형 헤드는 소프트플러스 활성화와 함께 단순히 가중치 행렬 (W_l) 로 구성되어 파라미터 수가 매우 적고, 학습 시 NLL을 그대로 사용한다. 또한, 각 레이어별로 독립적인 분산 모델을 학습한 뒤 평균을 공유하는 가우시안 혼합 형태(Ensemble)로 결합함으로써 모델 앙상블 효과를 저비용으로 얻는다.

실험에서는 대규모 equivariant graph neural network(E(3)GNN) 기반 모델을 QM9, PCQM4Mv2 등 여러 분자 데이터셋에 적용하였다. 평균 절대오차(MAE)는 기존 비이분산 모델과 거의 차이가 없으며, NLL은 제안 방법이 대부분의 베이스라인(베이지안 신경망, β‑NLL, Stirn et al. 2023 등)보다 우수했다. 특히, 중간 레이어를 모두 활용한 전체 합산 방식이 단일 레이어 방식보다 일관된 캘리브레이션을 제공한다는 점이 강조된다. 또한, 추론 시 추가 연산은 선형 변환 몇 번에 불과해 실시간 응용에도 적합함을 보였다.

이 논문은 “포스트‑핵” 전략이 대규모 회귀 모델에서 불확실성 추정을 실용적으로 구현할 수 있음을 증명한다. 최적화와 표현 붕괴를 별도 학습으로 회피하고, 홀드아웃 데이터로 잔차 과적합을 방지함으로써, 복잡한 하이퍼파라미터 튜닝 없이도 높은 캘리브레이션 품질을 얻는다. 향후 연구에서는 다른 도메인(예: 물리 시뮬레이션, 의료 영상)으로 확장하거나, 비선형 분산 헤드(예: 작은 MLP)와의 비교를 통해 더 복잡한 불확실성 구조를 모델링하는 방안을 탐색할 수 있다.