PR402 - 10: Bayesian Deep Learning via Subnetwork Inference

 오늘 가져온 논문은 Bayesian Deep Learning via Subnetwork inference(Daxberger et al, ICML 2021 accepted)입니다. 저번에 ICML 2021 accepted를 쭉 훑어보면서 읽어볼 논문들을 쭉 정리해놨는데 하나하나 읽고 리뷰할 만한 것들을 리뷰하다보니 계속 ICML에서 나오네요. 

Bayesian 관련 내용이 계속 나오는 걸 보니 확실히 Bayesian 딥러닝이 핫한 이슈인 것 같습니다. 물론 제가 관심 있어서 ICML 중에서도 그것들만 추린 것도 있긴 한데, 그럼에도 이렇게 끝없이 나온다는 자체만으로도요...!

Link - https://arxiv.org/abs/2010.14689

Intro:

Bayesian 딥러닝이 핫하지만, 현존하는 방법론들로 deep하고 wide한 NN으로 Bayesian 방법론들을 scale하기 위해서는 몇 가지 '빡센' 가정들을 필요로 합니다. 이러한 가정들은 inferred posterior의 representation power를 제한하고 uncertainty estimates의 퀄리티에 악영향을 끼칩니다. 이를 최근에 제안된 2가지 thesis를 바탕으로 해결해보고자 노력합니다. 첫째, NN의 accuracy는 small subnetwork로도 거의 재현된다. 둘째, low-dimensional subspace에서의 inference가 더 정확한 uncertainty quantification을 할 수 있다. 

위 두 가지 주장을 바탕으로 원래의 문제점인 'overparameterized NN이 posterior의 퀄리티를 낮춘다'라는 현상에서 overparameterized를 삭제하여(subspace로 대체하여) posterior와 uncertainty estimates의 퀄리티를 높이고자 하면서 원래의 full NN이 가지던 accuracy를 유지할 수 있지 않을까? 하는 생각이 들게 됩니다. 그렇다면 이제 문제는 과연 subspace를 어떻게 선택해야 '정확도'와 'posterior & uncertainty 퀄리티' 둘 다 잡을 수 있을까? 하는 것으로 귀결됩니다.

Contents:

Subnetwork의 선택은 다음 그림과 같은 순서로 이루어집니다:

우선, (a) Full NN을 구합니다. 구해서 각 weight를 고정합니다. (point estimation)

(b) subnetwork를 선택합니다. 이때 subnetwork를 선택하는 방법은 Full NN이 뱉어내는 posterior와 subnetwork가 뱉어내는 posterior의 KL-Divergence를 최소화하는 subnetwork를 선택합니다. 

(c) 선택한 subnetwork를 통해 Bayesian Inference를 시행합니다. 사실 (c)와 (d)는 거의 세트로 시행됩니다. (c)가 Jacobian과 Hessian의 근사, (d)가 근사된 Jacobian과 Hessian을 이용하여 Bayesian Inference를 하는 것이기 때문에 사실상 (c)와 (d)를 묶어서 Bayesian Inference라고 보는 편이 좋을 것 같습니다.

(c)에서의 근사는 GGN(generalized Gaussian Newton) matrix를 통해 이루어집니다. 대략 느낌은 SWAG에서 weights의 mean과 variance를 근사할 때처럼 Jacobian과 Hessian 또한 근사시켜서 사용하는 것과 같습니다. 

Jn과 Hn은 각각 n번째 데이터를 대상으로 일계미분값과 이계미분값을 계산한 값입니다. 원래 (weight의 차원^2) 크기의 Hessian을 구해야 하는데, 위처럼 구하면 (weight의 차원 * output size) * (output size ^2) * (output size * weight의 차원) 이 되기 때문에 계산량이 현저히 줄어들어 근사가 가능합니다. 이렇게 구한 근사를 바탕으로 Bayesian Inference를 수행합니다.

위와 같이 구할 수 있습니다. 단순히 prediction을 수행하는 거라면 regression에서는 Hessian은 필요없고 Jacobian만 구하면 되지만, 애초에 subnetwork를 사용해서 이렇게까지 한 이유가 전반적인 Uncertainty를 알기 위함이었으니, 실 목적을 달성하려면 근사된 Hessian으로 계산된 Covariance Matrix가 수반되어야겠죠. (이 과정에서 Linearized Laplace라는 방법을 사용하는데, 개념 자체는 그냥 MAP로 찾은 posterior에서 테일러근사를 사용하여 근사한다는 개념입니다.)

Experiments:

UCI dataset, rotate MNIST, corruption을 준 CIFAR 데이터셋 등에 대해서 uncertainty estimation을 평가합니다.

유명한 Bayesian 방법론들을 다 모아놨네요. 저번에 살펴봤던 SWAG도 들어가 있습니다ㅎㅎ 그런데 생각보다 성능 차이가 그렇게까지 나는지는 잘 모르겠습니다. reporting한 result에서도 퍼포먼스 차이가 그렇게까지 크지 않은 걸 보니 다른 데이터들에 대해서도 써먹을 만할 정도의 퍼포먼스를 보여줄지는 잘 모르겠습니다.

Conclusion:

low dimesion으로 근사시켜 uncertainty를 측정하는 것이기 때문에 다른 Bayesian 방법론들보다 '더 위험한' (training set과 test set의 분포 차이가 큰 경우) 에서는 고려해볼 만한 방법론이라는 생각이 들지만, 성능이 좋게 나올지는 모르겠습니다. 복잡하게 생각하지 말고 앙상블이나 dropout을 쓰는 편이 훨씬 정신건강에 이로울 것 같다는 생각이 들기도 합니다ㅋㅋ..

한편으로는 Uncertainty estimation 쪽에서 성능을 reporting할 때의 국룰 metric 같은 게 만들어지면 어떨까 하는 생각이 듭니다. 위에서 제시된 Bayesian 방법론들을 보면 마지막에 experiment results 할 때 사용하는 experiment들이 살짝 제각각인 감이 있습니다. 저번의 SWAG과 오늘의 Subnetwork inference를 봐도 reporting하는 초점이 조금 다릅니다. 

SWAG에서는 앙상블 개념까지 포함했기 때문에 성능에 관한 reporting도 있었고, uncertainty 차원에서는 calibration에 관한 내용이 들어가 있었는데 오늘의 알고리즘은 calibration이 아닌 corruption이 들어간 dataset에 대한 acc & loss를 따진 것을 보면 아직 이 분야에서의 metric이 정립되지 않았다는 느낌이 듭니다. 물론 돌려보고 별로 안 좋게 나와서 reporting을 안했을 수도 있지만요. 

오늘은 이렇게 마무리하겠습니다. 읽어주셔서 감사합니다!