Beyond Point Estimates: Benchmarking Uncertainty Quantification Methods on the AION-1 Astronomical Foundation Model

기술 요약: AION-1에 대한 불확실성 정량화(Uncertainty Quantification) 벤치마킹

문제 정의
AION-1과 같은 천문 관측용 파운데이션 모델은 은하 특성 추정과 같은 다운스트림 회귀 작업에 강력한 학습된 표현(learned representations)을 제공한다. 그러나 점 예측(point predictions)만으로는 과학적 추론에 충분하지 않으며, 오차를 전파하고, 분포 외 입력(out-of-distribution inputs)을 식별하며, 견고한 모집단 수준의 추론을 보장하기 위해서는 신뢰할 수 있는 불확실성 정량화(UQ)가 필요하다. UQ는 매우 중요함에도 불구하고, 고차원 임베딩을 가진 거대 블랙박스 파운데이션 모델에 대해서는 여전히 이해가 부족한 상태이다. 기존 방법들은 공식적인 커버리지 보증(formal coverage guarantees)이 결여되어 있거나, 복잡하고 이질적인 천체 데이터의 국소적 예측 난이도에 적응하지 못하는 경우가 많다.

방법론
저자들은 동결된(frozen) AION-1 파운데이션 모델(구체적으로 3억 개의 파라미터를 가진 AION-B 변체)에 적용된 7가지 UQ 방법을 직접 비교한다. 본 연구는 Legacy Survey 광도/영상 및 DESI 스펙트럼으로부터 5가지 물리적 특성을 회귀하는 데 중점을 둔다: 적색편이($z$), 항성 질량($\log M_\star$), 항성 집단 연령($t_{age}$), 가스상 금속 함량($\log Z$), 그리고 비특이적 별 형성률($\log sSFR$). 레이블은 PROVABGS SED-fitting 사후 확률(posteriors)로부터 도출되었다.

실험 설정은 다음과 같다:

1. 기본 아키텍처: 3개 층의 MLP가 768차원의 AION-1 임베딩을 스칼라 예측값으로 매핑한다. 인코더 가중치는 동결되며, 회귀 헤드(regression head)만 학습된다.
2. 평가된 UQ 방법:
   • Conformal 방법 (분포 무관 보증):
     • VanillaSplit: 전역 분위수(global quantiles)를 사용하는 표준 분할 컨포멀 예측(constant-width intervals).
     • MADSplit: 구간 너비를 국소적 난이도에 적응시키는 MAD-정규화 변체.
     • CQR (Conformalized Quantile Regression): 분위수 네트워크를 학습시킨 후 컨포멀 보정을 수행.
     • LVD (Locally Valid and Discriminative): 축소된 특징 공간에서 커널을 학습하여 조건부 비정합성 분위수(conditional nonconformity quantiles)를 추정함. 두 가지 공간에서 테스트됨: 원시 AION-1 임베딩 상에서(LCAion) 및 256차원 MLP 은닉 활성화 함수 상에서(LCNet).
     • LVD-MAD: 스케일 적응형 구간을 위해 학습된 잔차 크기 예측기(residual-magnitude predictor)로 확장된 LVD (LCMADAion 및 LCMADNet).
   • Non-Conformal 베이스라인:
     • Deep Ensembles (DE): 5개의 독립적으로 초기화된 MLP.
     • MC Dropout: 드롭아웃이 활성화된 확률적 순전파(stochastic forward passes).

평가는 주변 커버리지(marginal coverage), 꼬리 커버리지(tail coverage, 상위/하위 10%), 최악의 예측 빈(worst-predicted bin) 커버리지, 그리고 효율성(평균 예측 구간 너비)을 측정한다.

주요 결과

• 주변 커버리지(Marginal Coverage): 모든 컨포멀 방법(LVD 변체, VanillaSplit, MADSplit, CQR)은 모든 특성에 대해 명목 목표치인 90%의 ±1 퍼센트 포인트 이내에서 주변 커버리지를 달성했다. 반면, non-conformal 베이스라인은 보정(calibrate)에 실패했다: Deep Ensembles는 체계적으로 과소 커버리지(12–38%)를 보였고, MC Dropout은 지나치게 넓은 구간과 함께 강한 과잉 커버리지(57–99%)를 나타냈다.
• 꼬리 및 최악의 예측 성능:
  • 최악의 예측 빈에서, CQR이 가장 높은 커버리지(44–66%)를 달성했으며, 그 뒤를 LVD-MAD 변체(22–42%)가 이었다.
  • VanillaSplit은 가장 어려운 영역에서 커버리지가 약 0%로 붕괴되었다.
  • 원시 AION-1 임베딩에서 작동하는 LVD 변체(LCAion, LCMADAion)가 압축된 MLP 활성화 함수를 사용하는 경우보다 일반적으로 더 우수한 성능을 보였는데, 이는 768차원 공간이 국소적 보정에 유익한 구조를 유지하고 있음을 시사한다.
  • MAD 변체(LCMADNet, LCMADAion)는 non-MAD 대응 모델보다 일관되게 우수한 성능을 보였으며, 이는 이분산성(heteroscedasticity)에 대한 스케일 적응형 정규화의 중요성을 강조한다.
• 효율성: 잘 보정된 방법들 중에서, LVD 변체들이 커버리지를 유지하면서 가장 좁은 예측 구간을 생성한다. CQR은 주변 타당성을 보장하기 위해 더 넓은 구간을 생성한다.

주요 기여

1. 체계적 벤치마킹: 본 논문은 5가지 서로 다른 은하 특성 회귀 작업에 대해 동결된 AION-1 임베딩에 적용된 7가지 UQ 방법을 직접 비교한 첫 사례를 제공한다.
2. 국소 타당성 입증: 본 연구는 컨포멀 방법이 주변 타당성을 보장하는 한편, LVD 프레임워크(특히 파운데이션 모델 임베딩에서 작동할 때)가 유한 표본에 대한 국소적 타당성(local validity)을 제공한다는 것을 입증한다. 이를 통해 예측 구간이 단순히 전역 평균에 의존하는 대신 각 은하 예측의 특정 난이도에 적응할 수 있게 한다.
3. 임베딩 공간 분석: 결과는 압축된 중간 특징(256-dim)보다 원시 고차원 임베딩(768-dim)을 활용하는 것이 특히 꼬리 영역에서 국소적 보정의 견고함을 향상시킨다는 것을 보여준다.

의의 및 주장
저자들은 컨포멀 예측, 특히 LVD 프레임워크가 파운데이션 모델 임베딩에 대한 불확실성 인지 추론을 위한 선호되는 UQ 접근 방식이 되어야 한다고 결론짓는다. 분포 무관 보증이 결핍된 non-conformal 베이스라인 및 조건부 미보정(conditional miscalibration)을 겪을 수 있는 표준 컨포멀 방법과 달리, LVD는 분포 무관 보증과 국소적 적응성 사이의 원칙적인 균형을 제공한다. 이는 개별 예측의 신뢰도가 모집단 수준의 통계만큼 중요한 과학적 추론에서 매우 중요하다. 논문은 AION-B가 사용되었으나, 임베딩에서의 LVD의 이점은 더 큰 규모의 모델(AION-L 및 AION-XL)에서도 추가적인 조사가 필요함을 언급한다.