대규모 선형화 라플라스 근사 위한 대체 신경 커널 학습

초록

본 논문은 사전 학습된 대형 신경망에 대해 Linearized Laplace Approximation(LLA)의 커널을 직접 계산하지 않고, 작은 차원 임베딩을 출력하는 대체 신경망을 학습시켜 NTK를 근사한다. Jacobian‑vector product만을 이용해 효율적으로 학습하며, 학습된 커널에 의도적인 편향을 주어 OOD 탐지 성능을 크게 향상시킨다. 실험 결과는 기존 LLA 변형들과 비교해 예측 정확도·캘리브레이션·OOD 검출 모두에서 경쟁력 있거나 우수함을 보여준다.

상세 분석

이 연구는 LLA가 요구하는 거대한 Jacobian 행렬을 직접 계산·저장하는 비용을 회피하기 위해 “Surrogate Neural Kernel”(SNK)이라는 아이디어를 도입한다. 원본 네트워크 fθ∗(x)의 Jacobian Jθ∗(x)∈ℝC×P는 파라미터 수 P가 수백만에 달하면 메모리와 연산량이 폭발한다. 저자는 대신 차원 m≪P인 임베딩 gϕ(x)∈ℝC×m을 출력하는 작은 네트워크를 학습시켜, gϕ(x)gϕ(x′)ᵀ가 NTK K_NTK(x,x′)=Jθ∗(x)Jθ∗(x′)ᵀ를 근사하도록 한다. 핵심은 Jacobian‑vector product(JVP)만을 이용해 손실 L(ϕ)=E‖gϕ(x)gϕ(x′)ᵀ−Jθ∗(x)v·(Jθ∗(x′)v)ᵀ‖²를 최소화하는 것이다. 무작위 벡터 v는 Rademacher 분포를 사용해 분산을 낮추고, 자동 미분을 통해 JVP를 효율적으로 얻는다.

학습 과정에서 “context points”라 불리는 보조 데이터셋을 추가로 사용한다. 이는 원본 훈련 데이터와 구분된 입력으로, gϕ가 훈련 영역뿐 아니라 인접한 OOD 영역에서도 NTK 구조를 유지하도록 강제한다. 이후 커널 행렬을 구성할 때 훈련‑컨텍스트 간 교차 공분산을 0으로 강제하는 편향(bias) 기법을 적용한다. 이는 블록 대각 형태의 공분산을 만들며, 훈련 데이터와 미지의 영역 사이에 사전 분포로 되돌아가는 효과를 부여해 OOD 샘플에 대해 높은 불확실성을 반환한다.

실험에서는 FMNIST를 학습한 CNN을 대상으로, MNIST 테스트를 컨텍스트로, KMNIST를 OOD로 설정하였다. 비교 대상은 LLLA, VALLA, FMGP, MFVI, SNGP 등 최신 베이지안 근사법이다. ScaLLA(편향 없음)는 NLL·ECE에서 기존 방법과 동등하거나 약간 우수했으며, 편향된 버전은 OOD 검출 AUC‑ROC를 크게 끌어올리면서도 인‑디스트리뷰션 성능을 크게 손상시키지 않았다. 특히 LLLA는 OOD 검출은 최고였지만 인‑디스트리뷰션 정확도와 캘리브레이션이 떨어지는 반면, ScaLLA는 전반적인 트레이드오프를 가장 잘 맞춘다.

제한점으로는 OOD 편향을 위해 선택하는 컨텍스트 데이터가 실제 배포 환경과 얼마나 일치하느냐에 따라 성능이 좌우된다는 점이다. OOD 샘플이 사전에 알려지지 않은 경우 적절한 컨텍스트를 찾는 것이 어려울 수 있다. 또한, surrogate 네트워크 자체가 추가 파라미터를 요구하지만, m이 작아 전체 메모리 비용은 여전히 기존 LLA보다 훨씬 낮다.

요약하면, 이 논문은 JVP 기반의 효율적인 학습 프레임워크와 커널 편향 기법을 결합해, 대규모 사전 학습 모델에서도 실용적인 베이지안 불확실성 추정과 OOD 탐지를 가능하게 만든다.