[ICML 2023] SAAL: Sharpness-Aware Active Learning

제가 이번에 리뷰하려는 논문은, 새로운 optimizer 기법으로 성능 향상을 도모한 AL 연구입니다. 2021년도 짧게나마 ImageNet 에서 SOTA의 자리에 올랐던, SAM (Sharpness-Aware Minimization) 이라는 방법론이 존재한데요, SAM은 SGD, Adam 과 같은 옵티마이저 중 하나입니다. 이 SAM에 영감을 받아 Active Learning 에 적용해본 논문에 대해 리뷰해보도록 하겠습니다. (참고로 여기서의 SAM은 최근 우리에게 익숙한 Segment Anything Model 인 SAM이 아닙니다!!!)

---

• Title: [ICML 2023] SAAL: Sharpness-Aware Active Learning
• Open Review: Link
• Author Video: Link
• Code: GitHub

---

Introduction

Active Learning의 파이프라인은 모델 학습을 단 한 번만 하는 것이 아니라, 학습할 데이터를 추가하고 모델을 재학습시키는 과정이 반복됩니다. 첫번째 주기에서 (1) 1,000개의 학습데이터 D^L_0 로 모델을 학습했다면 (2) 해당 모델을 사용하여 unlabeled dataset로부터 1,000개의 데이터셋 D^L_1을 선별합니다. 그럼 두번째 주기에서는 (1) 기존에 있던 학습데이터와 D^L_0 선별한 데이터 D^L_1 2,000개의 데이터를 사용하여 첫번째 주기에 이어 모델을 재학습한 뒤 (2) 또 다른 1,000개의 데이터셋을 선별합니다. 즉, Active Learning에서는 (1) 모델 재학습 (2) 데이터 선별 및 추가 를 하나의 주기로 여러 번 반복하며 성능을 평가합니다.

반복적으로 데이터를 추가하다보니, Active Learning의 처음 주기에서 사용하는 데이터는 전체 데이터셋의 2%를 차지할만큼 소량으로 구성되어 있습니다. 소량의 학습 데이터셋이라면 고질적으로 따라오는 문제점들이 존재하죠, 바로 Bias 혹은 over-fitting 문제 입니다. 학습 데이터가 워낙 적으니, 모델은 그 데이터셋에 대해 쉽게 편향된 결과를 보입니다. 그런데 문제는 여기서 끝이 아닙니다. Active Learning은 반복적으로 모델을 업데이트시키는 탓에, 편향된 결과가 모든 주기에 계속 누적되어 전달되기도 합니다. 이로 인해, 한번에 10,000장의 학습데이터를 사용하는 것보다 1,000장씩 10번 학습한 Active Learning 기법의 성능이 더 낮은 결과를 보일 수도 있습니다.

따라서 저자는 Over-fitting 을 극복하기 위한 방식에 대해 고민하였습니다. 이를 해결하기 위해, 본 논문에서는 새로운 액티브 러닝 방법론인 Sharpness-Aware Active Learning (SAAL)을 제안하였습니다. SAAL은 Unlabeled Dataset 중에서 Loss의 날카로움(sharpness)을 최대화하는 인스턴스를 선택하는 새로운 획득 함수(acquisition function)를 도입하였습니다.

여기서 획득 함수란, Unlabeled dataset이 모델 성능에 얼마나 영향을 미치는 지 데이터가 가지는 가치를 정량화하는 함수입니다. 이는 기존의 Sharpness-Aware Minimization (SAM) 방식을 Active Learning 에 적용한 것으로, SAM은 Labeled dataset에서 학습 데이터의 maximally perturbed loss을 최소화하여 flat loss를 달성하는 것을 목표로 합니다.

Related Work

About Sharpness Aware Minimization (SAM)

[ICLR 2021] Sharpness-Aware Minimization for Efficiently Improving Generalization

본 논문을 이해하기 위해서는 SAM 이라는 최적화 알고리즘에 대해 이해해야합니다.

딥러닝 모델의 일반화 성능을 이해하기 예전부터 연구자들은 Loss Landscape의 형태를 연구해왔습니다. 이는 학습 후 모델이 도달하는 최소값의 특성, 즉, 최소값이 평평한(Flat)지, 아니면 날카로운(Sharp)지를 살펴보며, 해당 최소값이 모델의 일반화 성능과 어떤 관련이 있는지를 분석하였습니다.

상단 그림에 Flat과 Sharp minima에 대한 예시를 확인할 수 있습니다.

Flat minima는 일반적으로 파라미터의 작은 변화에 대해 Loss가 크게 변하지 않는 영역을 의미합니다. 이 때, 모델이 새로운 데이터에 대해 더욱 견고하게 작동할 가능성이 높다고 여겨집니다. 이는 학습 데이터에서의 작은 변화나 노이즈에 모델이 더 강건하게 반응하기 때문이라고 합니다.

반면, Sharp minima는 파라미터의 작은 변화에도 Loss가 급격히 증가하는 영역을 의미합니다. 여기서는 모델이 학습 데이터에 과도하게 최적화되어 있을 가능성이 높고, 결과적으로 새로운 데이터나 약간의 변동에 대해 취약할 수 있습니다. 이는 과적합(overfitting)의 현상과 관련이 깊다고 합니다.

다만 많은 연구가 이루어졌음에도 불구하고, 아직까지 어떤 유형의 최소값이 최적의 일반화 성능을 제공하는지에 대해 명확한 결론을 내리지 못하고 있다고 합니다. 일부 연구에서는 평평한 최소값이 더 좋다고 주장하는 반면, 다른 연구에서는 날카로운 최소값이 특정 조건에서 유리할 수 있다고 리포팅합니다. 이는 모델의 구조, 훈련 데이터의 특성, 최적화 기법 등 다양한 요소에 따라 달라질 수 있습니다.

SAM 논문에서는 Flat Minima가 일반화 성능 향상에 더 도움이 된다는 가정을 기반으로 시작됩니다. 이를 직관적으로 이해할 수 있도록 위의 예시를 함께 알아보겠습니다: 학습 데이터(Train set)와 평가 데이터(Test set)가 서로 달라 Loss Landscape이 옆으로 약간 Shift되어 달라진 상황입니다. Flat minima에서는 학습 데이터에서 찾은 최적의 Loss 값과 평가 데이터에서의 Loss값이 크게 차이 나지 않습니다 (파란색 동그라미). 즉, 학습 데이터에서 잘 작동하는 모델이 테스트 데이터에서도 비슷하게 잘 작동할 가능성이 높습니다. 반면 Sharp minima에서는 값이 아주 크게 차이나는 것을 알 수 있습니다(빨간색 동그라미).

그런데 Loss Landscape을 결정하는 요소는 모델의 파라미터 수, 사용할 수 있는 학습 데이터의 양, 그리고 학습 과정에서 사용되는 다양한 기술 등이 매우 다양합니다. 이러한 요소들을 조정하면서 Loss Landscape를 평평하게 만들어 날카로운 최소값(Sharp Minima)이 존재하지 않도록 하는 것이 이상적이겠지만, 실제로는 불가능한 일입니다.

따라서 저자들은 Loss landscape 자체를 변경하는 대신에, 모델이 sharp minima 방향으로 학습되지 않도록 최적화 기법을 수정하는 방법을 제안하였습니다. 이러한 접근 방식을 ‘Sharpness-Aware Minimization’이라고 하며, 줄여서 SAM이라고 부릅니다. SAM은 모델이 보다 평평한 방향으로 학습되도록 유도하여, 결과적으로 모델의 일반화 성능을 향상시킬 수 있습니다.

본 논문에서 정의한 가정에 따라, 최적화 과정에서는 다양한 최소값 중에서 flat minima 값에 도달하도록 하는 것이 중요합니다. SAM은 loss 값뿐만 아니라 loss의 sharpness도 고려하여, 파라미터 𝑤와 약간의 노이즈 𝜖를 더한 𝑤+𝜖 사이의 loss 차이가 작은 방향으로 학습을 진행하게 됩니다. 즉, sharpness하지 않은 방향으로 더 flat 한 minima로 수렴하도록 최적화 함수를 수정하였습니다.

위의 수식을 정리하면, 기존의 loss 함수에 Sharpness과 Weight Decay를 추가하여 새로운 loss 함수를 정의하고, 이를 최소화하는 파라미터 𝑤를 찾으면 됩니다. 이 과정은 경사 하강법(Gradient Descent)을 통해 수행한다고 하였습니다.

• ϵ: 파라미터 𝜃에 추가하는 perturbation
• 𝑆는 Training Dataset
• ∥𝜖∥≤𝜌 : perturbation 𝜖의 크기가 𝜌 이하로 제한됨을 의미
• 𝐿_𝑆(𝜃+𝜖): perturbation이 추가된 파라미터 𝜃+𝜖에서의 loss 함수.
• 파라미터 𝜃 주변에서 loss 값이 최대가 되는 방향을 찾아서 그 값을 계산하는 것.

더 자세하게 설명하자면, SAM에서는 파라미터 𝜃를 조정하여 다음 두 가지를 최소화하는 것을 목표로 합니다:

1. 파라미터 𝜃 주변의 loss 값의 최대값은 . 이를 통해 모델이 flat loss minima에 도달하도록 유도
2. 파라미터 𝜃의 크기 ∥𝜃∥^2_2 ​는 모델의 복잡도를 줄여 과적합을 방지하는 역할. 즉, 정규화 항. γ는 정규화항에 대한 가중치.

이제 문제는 Loss 함수 (2)를 최소화하는 것으로 변경되었습니다. 그러나 𝜖 을 직접 알 수 없기 때문에 GD를 곧바로 적용하는 것은 불가능합니다, 따라서 이를 위해 내부 최대화 문제를 해결하고, 이를 통해 ∇𝑤𝐿𝑆𝐴𝑀(𝑤)에 대한 근사치를 도출한 후, 확률적 GD을 적용하여 최적화를 수행하게 됩니다.

우선 내부 최적화를 해야하는 수식은 바로 위에 작성된 수식과 같습니다. Loss 𝐿_𝑆(𝑤+𝜖)를 0을 중심으로 1차 테일러 전개(First-order Taylor Expansion)하여 근사하게 됩니다. 이를 통해 다음과 같은 결과를 얻을 수 있습니다:

다음으로, Dual Norm Problem을 풀어서 다음과 같은 수식이 전개됩니다.

수식 (3)은 loss 함수의 gradient를 이용하여 최대 교란 ϵ∗를 계산하는 방법을 나타낸 것으로, ϵ∗는 loss를 최대화하는 방향으로의 작은 변화로, 파라미터 𝜃 주변의 loss landscape에서 가장 sharp한 방향을 찾게 됩니다.

이제 ϵ∗를 이용해 내적 극대화를 풀고 나면 아래와 같이 헤시안 항을 제외하면서 기울기를 구해 궁극적인 최소화 문제를 풀 수 있게 됩니다.

즉, 수식 (4)는 최대 교란 𝜖를 고려한 loss 함수의 기울기를 근사하는 방법을, 파라미터 𝜃에서의 loss 함수 𝐿_𝑆​의 기울기를 최대 교란 방향 ϵ∗에서 계산하여, 이를 통해 최대화된 loss를 반영한 기울기를 얻게 될 수 있습니다.

생소한 수식들이 즐비한 탓에 이해하는 데에 너무 많은 시간이 소요되었네요. 컨셉적으로 기존 Loss 함수에 Sharpness도 고려할 . 수있도록 수식을 수정하였다고 이해하시면 좋을 것 같습니다.

3. Method (Sharpness-Aware Active Learning, SAAL)

3.1 Motivation

이제 위에서 설명한 SAM optimizer를 AL에 접목시킬 때, 기존 수식을 AL 환경에 맞춰 변환하는 과정에 대해 설명하겠습니다.

SAM 논문에서 population dataset의 Loss인 \mathscr{D} 는 학습 데이터의 maximally perturbed loss인 \mathscr{X} 에 의해 상한이 결정된다고 하였습니다. AL에서 학습데이터는 라벨이 있는 \mathscr{X_L} 과 라벨이 없는 \mathscr{X_U} 로 나뉩니다. (즉, \mathscr{X} = \mathscr{X_L} ∪ \mathscr{X_U} ).

따라서 상한값은 다음과 같이 표현될 수 있습니다.

즉, population loss 인 L_\mathscr{D}(\theta)에 접근할 수 없으므로, 수식 (7)의 상한으로부터 네트워크를 학습하여 최소화하고자 하였습니다. 수식 (7)은 수식 (6)으로 전개되는데, 여기서 첫번째와 세번째 항은 위에서 설명한 SAM 최적화 기법을 통해 최소화할 수 있습니다. 남은 두번째 항인 unlabeled 에 대한 Loss가 Active Learning에 적용하기 위해서 해결해야하는 핵심 부분이 됩니다.

여기서 저자는 두번째 항 Loss가 높은 데이터를 선택하도록 acquisition function을 설계하였습니다. 즉 최대 perturb loss를 갖는 인스턴스를 선택하여 레이블을 부여하면, 해당 인스턴스는 더 이상 𝑋𝑈에 속하지 않고 𝑋𝐿에 포함됩니다. 따라서 𝑋𝑈​의 최대 Loss 인스턴스를 제거함으로써 𝐿𝑋𝑈​​ (즉, 𝐿_{𝑆𝐴𝐴𝐿}​)를 직접 줄일 수 있게 됩니다.

3.2 Sharpness-Aware Active Learning

따라서 저자가 제안하는 Active Learning 에서의 acquisition function은 수식 (8)과 같습니다:

l 는 cross-entropy loss function 입니다. 그런데 loss 를 ㄱㅖ산할 때에는 당연히 레이블이 필요합니다. 그런데 Unlabeled data에는 label 이 없기 때문에, 저자는 pseudo-label을 라벨로 사용하였습니다. 수식 (8)에서 \hat{y}_u가 수도 레이블을 의미합니다.

따라서 SAAL 알고리즘은 사실 굉장히 간단합니다.

4. Result

4.1 Image Classification

실험 세팅은 다음과 같습니다.

• Fashion-MNIST, SVHN, CIFAR-10, CIFAR-100 데이터셋을 사용하여 실험을 진행하였고, ResNet-18을 백본으로 사용
• 각 획득 단계 후 네트워크를 50epoch 학습, Adam 옵티마이저(lr=0.001)와 SAM 옵티마이저(lr=0.001 또는 0.1)를 사용
• ImageNet 실험에서는 Adam 옵티마이저를 사용해 각 획득 단계 후 500 epoch 동안 학습, 모든 설정을 세 번씩 반복 실험

실험 결과는 테이블 1에서 확인할 수 있습니다. SAAL은 8개의 중 7개에서 베이스라인보다 우수한 성능을 보였습니다. 특히, SAAL은 SAM 옵티마이저보다 Adam 옵티마이저를 사용할 때 더 높은 성능을 발휘하였습니다. 저자는 이 차이가 식 (7)에서 비롯된다고 추정하였습니다. 이 식은 데이터 선택 후 평평한 지역 최소값을 모델링하는 SAAL의 동기를 설명할 수 있다고 하였는데, 접근할 수 없는 목표인 L_D(\theta) 는 으로 상한이 정해집니다.

Adam 옵티마이저를 사용할 때, 이 상한의 첫 번째 term인 은 SAM 옵티마이저를 사용할 때보다 약하게 최적화됩니다. 이는 SAM 옵티마이저가 을 직접 최소화하기 때문이라고 합니다. 따라서 Adam 옵티마이저를 사용할 때는 상한의 두 번째 항목인 의 중요성이 더 커지게 됩니다.

Figure 12는 획득 반복에 따른 테스트 정확도를 보여주며, SAAL이 다른 베이스라인보다 더 빠르게 높은 정확도에 도달함을 나타냅니다 (Figure 12a, 12d, 12g). 또한, SAAL이 고해상도 데이터셋에서도 확장 가능하다는 것을 증명하기 위해 ImageNet에 대해 세 번의 반복 실험을 추가로 수행한 결과, Figure 2는 각 획득 반복에서 SAAL이 다른 베이스라인보다 우수한 성능을 보임을 보여주었습니다.

Qualitative Analysis

Figure 4에서는 labeled dataset 𝑋𝐿​, unlabeled dataset 𝑋𝑈​, 그리고 전체 데이터셋 𝑋𝐿∪𝑋𝑈​에 대해 최대 perturbation loss를 측정합니다. 모든 unlabeled 인스턴스에 대해 해당 perturb를 계산하는 것은 계산 비용이 많이 들기 때문에, 매 cycle마다 𝑋𝑈​에서 랜덤하게 2,000개의 인스턴스를 샘플링하여 평균 결과를 그래프로 나타낸 것입니다.

• Figure 4a: SAAL은 SAM 옵티마이저를 사용할 때, 다른 베이스라인보다 낮은 최대 perturb loss 값을 보입니다. 이는 SAAL이 높은 perturb loss 값을 가진 인스턴스를 선택하고, 이를 레이블된 데이터셋에 추가하여 제거했기 때문이라고 합니다.
• Figure 4b: 𝑋𝐿​의 perturb loss를 보여주며, 이 값이 낮을수록 모델이 평평함을 나타냅니다. SAAL은 다른 베이스라인보다 낮은 값을 보여, 더 평평한 모델을 생성합니다.
• Figure 4c: 전체 데이터셋의 최대 perturb loss를 보여주며, 이는 식 (7)의 상한값에 해당합니다. SAAL은 가장 낮은 상한값을 달성하여, 궁극적으로 더 낮은 전체 loss를 달성할 가능성이 높습니다.

Visualization of Loss Landscape

SAAL은 1) sharpness가 높은 인스턴스를 선택하고 2) 이러한 sharpness를 가진 인스턴스의 loss를 빠르게 감소시키도록 모델을 학습시켜 flat한 landscape을 구축하는 것을 목표로 합니다. 이를 위해 loss 헤시안 행렬의 첫 번째 eigenvalue를 사용하여 landscape 시각화한 결과를 Figure 5에서 보여주며, visualization 결과와 첫 번째 eigenvalue를 통해 SAAL이 더 flat loss landscape을 가지고 있음을 확인할 수 있습니다.

4.2. Ablation Study on Image Classification

Robustness to Class Imbalanced

SAAL의 robustness를 추가로 평가하기 위해, long-tailed CIFAR-10 데이터셋을 만들어 low budget 설정에서 실험을 진행했습니다. Table 3에 따라, SAAL은 불균형한 상황에서도 다른 베이스라인보다 뛰어난 성능을 보였습니다.

Sensitivity Analysis on ρ

SAAL은 perturb 영역의 크기를 나타내는 하이퍼파라미터 ρ를 사용하였는데, 본 실험을 통해 CIFAR-10 데이터셋을 사용하여 𝜌의 민감도 분석을 시도하였습니다. 𝜌의 후보 값으로 0.01, 0.05, 0.10을 설정했습니다.

• Figure 6과 Figure 7의 분석을 통해 𝜌=0.05가 최적의 값임을 확인할 수 있었고,
• 𝜌=0.05는 데이터 인스턴스의 예측 레이블을 높은 확률로 유지하며, 최대 perturb loss에 비해 여유 값을 상대적으로 작게 유지
• 결론적으로, 𝜌=0.05는 𝜖의 적절한 값으로, 예측 레이블을 유지하고, 여유 값을 적절히 유지하며, 테스트 정확도를 높이는 데 가장 효과적임을 확인했습니다.

4.3. Object Detection

SAAL의 효과를 복잡한 작업에서 보여주기 위해 object detection으로 확장한 결과도 보였습니다. detection 모델의 학습 loss는 바운딩 박스 regression 과 classification loss로 구성됩니다. 저자는 PASCAL VOC 2007 및 2012 데이터셋(각각 5,011장과 4,952장 이미지)으로, detector는 Single Shot Multibox Detector(SSD)를 사용하였습ㄴ디ㅏ.

SAAL을 객체 탐지에 적용하기 위해, classification loss를 최대화하도록 파라미터를 교란시키고, 이미지 𝑥의 각 bounding box에서 perturb loss의 합을 획득 점수로 사용하였다고 합ㄴ디ㅏ. 이후, 점수가 가장 높은 이미지를 선택했습니다. 초기 레이블링된 데이터셋은 무작위로 선택한 1,000개의 이미지로 구성했고, 각 획득 반복마다 추가로 1,000개의 인스턴스를 선택하여 10,000개의 최종 인스턴스를 얻었다고 합니다.

Figure 10은 SAAL과 다른 베이스라인의 mAP를 세 번 반복한 결과입니다. SAAL은 초기 iteration에서 높은 성능을 보였고, 마지막에는 0.7541의 최고 mAP를 기록했습니다. 반면, BADGE, Entropy, Random은 각각 0.7493, 0.7518, 0.7403을 기록하엿다고 합니다.

Conclusion

AL에 SAM을 적용한 연구에 대해 알아보았습니다.

방법론에 대한 실험 및 증명은 어떻게 다루고 있는지, 그나마 최근 알고리즘인 SAAL에 대해 자세히 알아보았습니다.