리치 베이지안 라스트 레이어: 서브샘플링 NTK 특징으로 불확실성 강화

초록

베이지안 라스트 레이어(BLL)는 마지막 층만 베이지안 처리해 효율적인 불확실성 추정이 가능하지만, 앞층의 불확실성을 무시해 에피스틱 불확실성을 과소평가한다. 본 논문은 NTK(Neural Tangent Kernel) 특징을 마지막 레이어 특징 공간에 투사하여 전체 네트워크의 변동성을 반영한 “리치 BLL(Rich‑BLL)”을 제안한다. 핵심은 비‑마지막 층의 NTK 특징을 선형 변환 행렬 A 로 근사하고, 이를 Cholesky 분해된 저차원 변환 L 으로 압축해 BLL과 동일한 O(r³) 복잡도로 posterior를 계산한다. 또한 균일 서브샘플링을 이용해 A 와 공분산 행렬을 소수의 샘플만으로 추정하고, 근사 오차에 대한 이론적 경계와 예측 분산이 기존 BLL보다 항상 보수적임을 증명한다. 실험은 UCI 회귀, 컨텍스트 밴드릿, 이미지 분류 및 OOD 탐지에서 캘리브레이션과 불확실성 추정이 크게 개선됐으며, 서브샘플링을 적용해도 비용이 크게 감소함을 보여준다.

상세 분석

본 연구는 베이지안 라스트 레이어(BLL)의 근본적인 한계인 “앞층 불확실성 무시” 문제를 NTK 이론과 결합해 해결한다. NTK는 무한 폭 신경망에서 파라미터 변동이 출력에 미치는 영향을 정확히 포착하는 커널이며, 전체 파라미터 그래디언트 ϕₚ(x)=∇θ fθ̂(x) 를 특징으로 한다. 저자들은 ϕₚ(x)를 마지막 레이어 특징 ϕᵣ(x)와 비‑마지막 레이어 특징 ϕₘ(x)로 분해하고, ϕₘ(x)≈A ϕᵣ(x) 라는 선형 근사를 제안한다. 여기서 A는 최소제곱 방식으로 학습되며, 실제 구현에서는 A=ΦₘᵀΦᵣ(ΦᵣᵀΦᵣ)⁻¹ 으로 구한다. 중요한 점은 A를 직접 저장할 필요 없이 B=