논문 리뷰: BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning

 

"BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning" (ICLR 2020) - Wen et al.

여러 모델을 사용하여 Ensemble을 하면 보통 single model을 사용하는 것보다 좋은 성능을 가져온다. 하지만 모델을 여러 개 training 해야 하므로 모델 개수만큼 computational cost가 발생하게 된다. 이번 논문은 이런 단점을 극복하기 위해 나온 논문이다.

 

---

Introduction

 

여러 사람이 함께 문제를 풀면 더 좋은 결과를 가져오는 경우가 많다. 이런 직관적 이유에서인지 Ensemble 방법은 오래전부터 사용해왔다. 실제 performance가 좋았고 많은 연구가 되어왔다. 단점도 존재하였는데 training시 computational cost가 많이 든다는 것, memory를 많이 차지할 수 있다는 점이다. 이런 문제 때문에 stochastic weight averaging (SWA) 같이 gradient descent를 하는 도중 여러 weight를 sampling을 하는 방법이나 MCMC를 사용하여 model을 sampling 하는 논문들이 나오고 있다. 아쉽게도 실제 Ensemble 보다는 보통 결과가 좋지 못하다. 그 이유로는 하나의 local minima 내에서 sampling을 하기 때문이라고 생각되고 있다. 서로 다른 initialization에서 training을 하여 다른 minima로 업데이트된 모델들이 좋은 Ensemble 결과를 가진다는 논문이 생각을 뒷받침 해주고 있다. 이번에 소개하는 논문은 빠른 training과 memory efficient 한 장점을 가지면서도 여러 initialization에서 업데이트된 파라미터들로 좋은 성능을 낸다는 강점이 있다.

 

---

Method

 

방법은 아래 Figure 2에 나온 것이 전부인데 상당히 simple하다.

W는 neural network backbone에 해당하는 weight matrix이며 $r_{i}, s_{i}$가 새로운 ensemble weight를 만들어내는 parameter이다. $\bar {W}_{i} = W \circ F_{i}$ $ where F_{i} = r_{i} s_{i}^{T}$ 로 계산되어 $\bar {W}_{i}$가 ensemble model의 weight가 된다. 학습 시 mini-batch를 i 개수만큼 쪼개서 sub-batch를 만들고 각각의 batch로 $r_{i}, s_{i}$를 학습시킨다. 한편 backbone이 되는 W는 공통으로 학습시킨다. 이렇게 했을 때 서로 다른 initialize로부터 학습된 $r_{i}, s_{i}$ parameter로 i개만큼의 ensemble 모델을 만들 수 있으며 parameter 크기가 작으므로 memory efficient 하다는 장점이 있다. 추가로 backbone W는 그대로 두고 새로운 task를 학습시킬 때 $r_{i}, s_{i}$ parameter만 학습시킬 수도 있으므로 lifelong learning에도 효과적으로 쓰일 수 있다.

 

---

Experiments

Figure 3에서 lifelong learning 실험 결과를 잘 보여주고 있다. performance가 좋을 뿐 아니라 time & memory cost가 굉장히 작음을 알 수 있다.

 

Classification에도 Single model이나 MC-dropout을 사용한 ensemble 방법보다 좋은 성능을 내고 있다. 아쉽게 이 논문에서도 Naive Ensemble 방법을 뛰어넘지는 못했지만 훨씬 efficient하다는 점을 생각할 때 큰 발전이라고 생각한다.