[Paper Review] Learning Loss for Active Learning 논문 리뷰 및 구현

Paper-Review

[Paper Review] Learning Loss for Active Learning 논문 리뷰 및 구현

Dobby-HJ 2023. 1. 5. 19:58

0 BackGround Knowledge

1. Abstract

Deep Convolution Network는 더 많은 Annotated Data가 있으면 있을수록 성능이 증가합니다. 하지만 여기서 문제는 Annotated Data를 만들기 위한 한정된 Budget입니다. 본 논문에서는 이를 타개하기 위한 하나의 Solution으로 Loss Prediction하는 Module을 추가적으로 장착해 Active Learning을 하도록 제안합니다. 기존의 Active Learning 방법론에 비해 Task-agnostic하며, Architecture에 영향을 받지 않는 굉장히 General하고 원래 Network에 비해 굉장히 작은 Module만 추가적으로 부착하면 되므로 Efficient하다고 합니다.

2. Introduction

위의 사진처럼 Annotated Data를 주면 줄수록 현재 Deep Convolution Network는 성능이 Saturate되지 않고, 점점 증가하고 있습니다. 그래서 Large-scale Dataset을 만들기 위해 Semi-supervised learning에서부터 unsupervsied까지 시도하고 있지만, fully-supverised learning에 비해서 성능이 떨어집니다. 그리고 특히 "Annotated Data"가 중요한 Medical AI에서는 그 격차가 더 큽니다.

하지만 AnnotationData 또한 Target Task에 따라 Cost가 굉장히 차이 납니다. Task가 어려워지고 Labeler의 전문성이 더 필요해질수록 Cost는 높아집니다. 직관적으로 일반적인 Image Classification가 Medical AI에서의 Classification 비해 저렴한 Annotation 비용이 듭니다.

그래서 저자는 같은 양의 annotated Data가 있을 때, 이를 이용해 효과적으로 학습시키는 방법이 필요했고, Active Learning이 적절한 방법이라 생각했습니다. Active Learning은 학습할 Model이 uncertain한 Data point를 효과적으로 Select해 학습을 하는 방법입니다.

Unlabeled Data Pool에서 Data를 Select하는 방법은 크게 3가지가 있습니다.

Uncertainty-Based Approach
- 잘 못 맞출 것 같은 Data들을 Annotation
Diversity-Based Approach
- Dataset 전체를 모사(근사, represent)하는 Subset을 Annotation
Expected Model Change
- Model의 output을 변화시키거나, parameter를 가장 많이 변화시킬 것 같은 Data를 Annotation

1번의 Uncertainty-Based Approach의 가장 쉬운 구현 방법은 바로 Class의 posterior probabilities를 활용하는 것입니다. 예측된 확률 혹은 사후 확률에 의거한 entropy는 data point의 uncertainty로 정의하는 것입니다. 직관적이지만, 단점은 Classification와 같은 Task에만 적용가능합니다. 더 어려운 recognition Task인 Detection, Segmentation, Pose Estimation등에서는 잘 동작하지 않아, 우리의 목적인 Task-Agnostic하지 못하다는 단점이 있습니다.

다른 Task-Agnostic한 Uncetainty Approach는 다양한 Model을 학습시켜, 각 모델들의 Prediction에 대한 차이를 이용해 unccertainty를 측정하는 것입니다. 같은 Sample에 대한 Prediction의 분산이 큰수록 Uncertainty하다고 정의하는 게 대표적입니다. 하지만 이러한 방법은 최신 Deep Nerual Network를 학습하는데 너무 많은 Cost가 발생한다는 단점이 존재합니다.

최근에는 Deep Network에 Monte Carlo Dropout과 같은 방식을 도입해 1개의 Model을 학습시키는 방법도 있습니다. 해당 방법은 모델의 전체에 Dropout을 적용해 반복학습한 합니다. 이후 Inference시에 일반적으로 Dropout을 제거하는 방식과 달리 학습때와 동일하게 Dropout을 적용해 $N$번 Inference하여 해당 N개의 결과의 분산을 Uncertainty로 정의하는 방법입니다. 이 또한 최신 Deep Neural Network에서는 Dropout이 적용된 거대 모델을 학습하는데 어렵고, 특히 큰 Dataset에서는 많은 Cost가 소요된다고 합니다.(CIFAR-10 기준으로도 며칠씩 걸린다고 합니다)

Distribution Approach(Diversity-Based Approach)는 Task-Agnostic 할 수 있습니다. 왜냐하면 Distribution Approach는 Prediction이 아닌 Feature Space에만 의존하기 때문입니다. 하지만 이러한 방식은 localization Task(Detection, Segmentation)등을 위해 location-invariant feature sapce(위치에 의존적이지 않은 Feature Extraction을 위해)를 추가적으로 engineering 해야 합니다. 이러한 방법은 작은 Model에서는 괜찮지만, 최근의 Deep Nerual Network에서는 Computation Cost가 Impractical 하다고 합니다.

본 논문에서는 Task-Agnostic하고, Architecture에 상관없이 Loss를 Prediction하는 작은 Module을 부착해, 해당 모델을 통과했을 때 나오는 실제 Loss를 예측할 수만 있다면, 우리는 high loss가 나오는 Unlabeled Sample을 선택할 수 있게 됩니다. 이렇게 선택된 Data들은 Model에 더 나은 Information을 주입할 수 있습니다.

Figure 1 - (b) Active Learning with a loss prediction module

2.1 Contributions

최신 Deep network에 바로 적용가능한 loss prediction module을 이용해 쉽지만 efficient한 active learning Method를 제안합니다.
대표적인 3가지 Task(Image Classification, Object Detection, Human Pose Estimation)를 통해 성능 검증을 했습니다.

3. Related Research

Uncertainty Approach

먼저 Uncertainty Approach에서 Classical한 active learning은 posterior probability of a predicted class 혹은 margin between probability of predicted class 혹은 margin between posterior probabilities of a predicted class and the secondly predicted class입니다. 위의 컨셉으로부터 posterior probabilites의 entropy는 definition의 방정식을 일반화할 수 있습니다. 예를 들어 SVM에서는 Decision boundary로부터의 distance가 uncertainty로 정의할 수 있습니다. 다른 방식으로는 여러 개의 독립 모델을 구성한 뒤, 1개의 sample에 대해서 각 모델의 불일치를 측정하여 uncertainty를 정의합니다.

Distribution Approach

Distribution Approach는 Unlabeled Pool을 represent할 수 있는 Data Sample을 선택하는 전략입니다. 직관적으로 전체를 represent할 수 있는 subset을 선택하는 discrete optimization problem으로 formulate 됩니다.(그렇다고 합니다;;) 각각의 Data Sample들이 주변의 데이터 포인트와 얼마나 가까운지 고려하여 Dataset(knowledge)을 잘 represent할 수 있는 Datasample들을 선택하는 전략입니다.(앞서 언급한 대로 이는 굉장히 높은 time complexity를 가지게 되는데, 직관적으로 permutation의 형태로 K개의 smaple을 선택해 K-Means가 정해이고, 휴리스틱 하게 근사를 할 수 있도록 optimizing을 추가적으로 해 time complexity를 줄일 수 있습니다. Trade-off입니다.)

Distribution Approach(Core-Set Method)

Expected Model Change

Expected Model Change 방법은 조금 더 정교하고 Decision-Theoretic한 접근방법입니다. 해당 방법은 모든 Label에 대해서 Gradient length, Future error, output change를 Estimate 합니다.(추가예정입니다.)

각각의 전략들은 모두 각각의 장단점이 있습니다.
Entropy를 이용한 Uncertainty를 측정하는 방법은 Classification에서 밖에 사용할 수 없고, 나머지 Task에서는 Entropy를 정의하는 게 어렵습니다.
Ensemble Approach는 학습시켜야 할 CNN이 너무 커지거나 Data가 많아지면 Computation Cost가 너무 많이 들어갑니다.
Monte Carlo Dropout을 통해 Multiple Forwardpass는 학습(수렴)이 어렵고, Inference시에도 $$N$$번 반복해야 하므로 좋지 않습니다.
Distribution Approach는 Unlabeled pool이 클 경우에는 Optimization 방법론이 굉장히 무거워집니다.

4. Method

4.1 Overview

본 챕터에서는 loss prediction module을 이용해 active learning 시나리오를 명확히 정의합니다.
Target Model인 $$\Theta_{\text {target}}$$과 loss module인 $$\Theta_{\text {loss}}$$을 이용할 예정입니다. loss prediction은 위의 그림처럼 target 모델에 부착되어 있습니다.

Figure 1 - (A) A model with a loss prediction module

Target Model은 input $$x$$를 넣어 $$\hat{y}$$를 얻어내는 task를 수행합니다. $$\hat{y}=\Theta_{\text{target}}(x)$$. 반면 loss predcition module은 $$h$$($$x$$를 이용해 $$\Theta_{\text{target}}$$에서 Extraction된 Feature map들의 set)을 이용해 $$\hat{l}= \Theta_{\text{loss}}(h)$$를 수행해 loss를 prediction합니다.

대부분의 real-world의 learning problem은 large pool of unlabeled data $$\mathcal{U}N$$를 한 번에 얻을 수 있다는 것에 있습니다. ($$N$$은 데이터 point(sample)의 수를 의미합니다.) 처음 Dataset을 sampling 할 때에는 $$K$$개의 Data를 Sampling 합니다.(labeling작업을 합니다) 이를 $$\mathcal{L}^0_K$$라 하겠습니다. 이러면 unlabeld data set의 크기는 $$\mathcal{U}^0{N-K}$$로 줄어들게 됩니다.

이렇게 K개의 Sample이 sampling된다면, 해당 Dataset을 이용해 Target Model $$\Theta_{\text{target}}^0$$를 학습시킵니다. 그리고 loss prediction module 또한 $$\Theta_{\text{loss}}^0$$ 학습시킵니다. 이후 우리는 unlabeld pool로부터 loss prediction module을 이용해 loss를 prediction합니다. 그래서 data-loss pair를 obtain하게 되고, ($$\{(x, \hat{l})\ |\ x \in \mathcal{U}^0_{N-K}\}$$)

이후, 사람이 $$K$$-highest loss를 가진 data들에 대해서 annotation을 달았습니다. 이렇게 다시 K개의 Dataset이 생겼고, 이를 이전의 Dataset과 합쳐서 이번엔 $$\mathcal{L}_{2K}^1$$의 Labeled Dataset을 만들어냅니다. 이러한 cycle을 반복해 원하는 성능을 얻게 되거나, annotation에 들어갈 budget이 없을 때까지 반복합니다.

4.2 Loss Prediction Module

loss prediction module이 target model에서 정의된 loss를 imitate(맞추는 것)이 task-agnostic active learning에서의 핵심입니다.

loss prediction module을 이용해 loss를 prediction을 하는 해 active learning을 하는 것에 초점을 뒀지만, 이미 Deep Nerual Network로 인해 높은 Computational Cost가 적용되는 상황이므로 해당 Module이 Computation Cost가 Minimize 되었으면 좋겠습니다. 그래서

Loss Prediction Module은 Target Model보다는 훨씬 작으면 좋겠습니다.
Target Model과 함께 학습되면 좋을 것 같습니다.(2개의 분리된 stage 형태가 아니었으면 좋을 것 같습니다.)

그래서 다음과 같은 형태의 Module이 만들어졌습니다.

Figure 2. The architecture of the loss prediction module

일반적으로 이러한 loss prediction의 loss function으로 MSE가 떠오릅니다. 고로 다음과 같이 $$(\text{MSE}) \ L_{\text{loss}}(\hat{l},l) = (\hat{l} - l)^2 $$ 과 같은 형태가 될 것입니다. 문제는 모델이 information을 얻으면서 unlabeled Data의 Uncertainty가 점점 줄어들고 이는 곧 전체적인 $l$값이 줄어드는 것을 알 수 있습니다. 그래서 위와 같은 MSE loss function을 naive하게 적용할 경우, loss prediction에 집중하는 것이 아닌, 그저 $$\hat{l}$$의 절대적인 수치를 줄이는 것에 집중하는 것을 포착했습니다. 그래서 기존의 active learning 방법들보다 더 낮은 performance를 보여줍니다.

우리는 이를 해결하기 위해 Ranking Loss를 도입하는 것입니다. Ranking Loss는 기존과 달리 2개의 Data를 묶어 Pair를 맞춰준 후, 해당 Paired Data쌍의 순서가 들어온 Data의 Loss의 크기순이 아니면 loss를 발생시키는 방법입니다.

$$ L_{\text{loss}}(\hat{l^p},l^p) = \max(0, -1(l_i, l_j)\cdot (\hat{l_i} - \hat{l_j} + \xi)\\ \begin{align} \text{s.t.} 1(l_i, l_j) = \begin{cases} +1 , &\text{if} \quad l_i > l_j\\ -1, &\text{otherwise} \end{cases} \end{align} $$

여기서 $\xi$는 pre-defineed positive margin이고, 최종적으로 $$\hat{l_i}$$가 $$\hat{l_j} + \xi$$보다 큰 경우에는 loss를 주어 $$ \hat{l_i}$$는 increase하게, $$\hat{l_j}$$는 decrease하도록 합니다.

즉 최종적인 $$\mathcal{B}^s$$인 mini-batch가 active learning stage $$s$$에 주어졌을 때, 다음과 같이 Target Model과 Loss prediction module을 joint학습을 위한 loss function은 다음과 같습니다.

$$ \cfrac{1}{B}\sum_{(x,y)\in\mathcal{B}^s} L_{\text{target}}(\hat{y},y) + \lambda\cfrac{2}{B}\cdot \sum_{(x^p, y^p)\in \mathcal{B}^s}L_{\text{loss}}(\hat{l^p}, l^p)\\ \begin{split} \hat{y } & = \Theta_{\text{target}}(x)\\ \text{s.t } \hat{l^p}
&= \Theta_{\text{loss}}(h^p) \qquad \\ l^p &= L_{\text{target}}(\hat{h^p}, y^p) \end{split} $$

결론적으로 분리된 학습 procedure 없이 약간의 추가적인 학습 parameter만으로 loss prediction module은 data에 대해 most informative data smaple을 고를 수 있습니다. (이후에 구현파트에서 언급하겠지만, 위와 같이 단순히 Loss prediction Module로 부터 예측된 Top-K개의 Unlabled Data Sample들을 보게 되면 전부 같은 class의 비슷한 Data들입니다. 직관적으로 이해한 바로는 s stage 상에서 모르는 데이터를 sampling할 때는 잘 모르는 데이터들이지만, 해당 데이터가 몇 개만 들어와도 사실 이제 "잘 아는" 데이터가 되므로, 이러한 diversity를 확보하기 위해 전체 Unlabeled Data에서 Top-K를 얻는 것이 아닌 전체에서 일부를 Random Sampling한 뒤, 해당 Pool에서 Top-K를 얻어옵니다.)

(※ Margin Rank Loss의 목적은 두 Data Sample 같은 Loss가 Margin 만큼 차이 나도록 학습하는 것, 이로 인해 Ranking을 매길 수 있게 됨.)

5. Evaluation

(Task별 세부사항은 추후 추가 예정)

)

위 두 사진 속에서 1번은 본 논문의 방식대로 했을 때이고, 아래는 entropy-based method를 이용해 데이터를 시각화한 것이다.

그래서 실제 Loss를 x축으로 했을 때 그림에서 보면 본 논문의 방식대로 적용했을 때 실제 loss값과 더 높은 correlation을 가지는 것을 증거로 본 방식이 Informative Data point를 잘 Sampling하는 방식이라 설명합니다.

6. Limitations and Future Work

본 논문에서는 굉장히 좋게 Deep Neural Network에 Task-Agnostic한 active learning을 적용할 수 있는 방법을 소개했습니다. 여러 Visual recognition Task를 통해 검증되었지만, Data의 Diversity나 Density를 고려되지 않았습니다. 또한 Classification보다 더 어려운 Pos estimation이나 Detection의 경우 상대적으로 결과가 낮게 나왔습니다. 또한 loss predition moudle의 정확도를 높이기 위해 더 나은 아키텍쳐와 objective function에 대해 더 연구가능합니다.