이번에 리뷰할 논문은 ICLR 2023에 공개된 SURGICAL FINE-TUNING IMPROVES ADAPTATION TO DISTRIBUTION SHIFTS이라는 논문입니다. 교신 저자로 계신 분이 Chelsea Finn 교수님이셔서 Meta Learning을 연구하는 저에게는 매우 흥미가 가는 논문입니다. 참고로 Chelsea Finn 교수님은 MAML(Model Agnostic Meta Learning)이라는 아주 유명한 논문의 저자이십니다! 그럼 간단 논문 리뷰 시작해 보겠습니다.
리뷰한 내용의 부족한 점, 잘못된 점에 대한 피드백은 언제든 환영합니다!
Abstract
Distribution shift에서 pretrained model을 fine-tuning 하는 대표적인 방법으로는 마지막 몇 개의 layer (the last few layers)만을 새로운 task에 적응하는 것이 대표적입니다. 본 논문에서는 선택적으로 layer를 fine-tuning하는 surgical fine-tuning 방법을 제안합니다. 이론적으로 two-layer neural network에서 모든 layer를 fine-tuning 하는 것보다 첫 번째 layer만을 fine-tuning 하는 것이 더 우월한 성능을 보여줌을 증명합니다.
Method
Figure 1은 distribution shift가 있을 때, Full fine-tuning과 Surgical fine-tuning을 비교한 그림입니다. 결과와 같이 First block/Middle/Last layer만을 fine-tuning 하는 방식이 모두 더 좋은 성능을 보여주고 있습니다. 특히 CelebA에서는 4%나 더 높은 뛰어난 성능을 보여줍니다.
본 논문에서는 ResNet을 기반으로 실험을 진행합니다. fine tuning하는 layer는 Block단위로 Block1, Block2, Block3, Block4로 나누며 마지막 fully connected layer에 대해서는 "Last Layer"라고 지칭합니다.
Figure 2에서는 distribution shift에 대해서 특정 layer의 fine-tuning 성능을 보여줍니다. 각 데이터셋에서의 특성에 따라 fine-tuning을 위한 Block을 선택하는 것이 중요한 문제임을 알 수 있습니다. 해당 Block만 잘 찾을 수 있다면 훌륭한 성능을 보이지만 그게 아니라면 매우 심각한 성능 감소도 같이 보여주고 있습니다.
다음으로 논문의 저자들은 이 문제를 작은 문제로써 two-layer networks의 관점에서 접근합니다. 그들은 이 모델에서 first layer와 last layer를 fine-tuning 함으로써 target data에 대해서 loss가 0에 가까운지, 0보다 큰지를 통해 잘 adaptation 되었다고 판단하며 이를 증명합니다. 이것을 저자들은 Input perturbation이라고 이야기하며 Label perturbation에 대해서도 같은 작업을 수행합니다. 자세한 내용은 생략합니다.
Automatically Selecting which layers to Tune
저자들은 특정 layer에 대한 fine-tuning 을 자동으로 model이 선택하게 만드는 방법 2가지를 제안합니다. 여기에서는 특정 layer들을 freeze 또는 업데이트하는 형태로 접근하며 첫 번째 방법은 Relative Gradient Norm (Auto-RGN) 방법입니다.
매우 간단하게 gradient의 norm을 parameter norm으로 나눠준 것을 RGN값으로 정의합니다. 여기서의 핵심은 gradient norm입니다. gradient norm이 크다는 것은 모델이 현재 입력에 대한 representation의 변화가 크다는 의미로써 target task에 대한 더 많은 정보를 가지고 있다고 판단합니다. 따라서, 본 논문의 저자들은 RGN값에 대해 normalization을 수행한 후 learning rate를 곱해줌으로써 per-tensor learning rate로 전체 모델을 업데이트합니다. 만약 RGN값이 크다는 것은 더 큰 learning rate임을 의미하며 해당하는 layer는 비교적 충분한 업데이트가 가능합니다.
다음으로는 Signal-to-Noise Ratio (Auto-SNR)의 방법을 제안합니다. 신호대잡음비라고도 잘 알려진 이 값은 아래와 같이 정의됩니다.
본 논문에서 제안하는 SNR은 평균과 분산의 기댓값으로 이루어져 있습니다. 이 값은 특정 threshold보다 낮을 경우 해당하는 layer는 freeze 하는 형태로 업데이트합니다. 이와 같은 automatically selecting 방법을 통해 아래와 같은 성능을 보여줍니다.
Auto-RGN의 방법을 사용하면 대부분의 데이터셋에서 Full fine-tuning보다 좋은 성능을 보이며, 다른 fine-tuning 방법 중 하나인 Gradual Unfreeze보다도 우월한 성능을 보여주고 있습니다.
또한 Auto-RGN을 통해 특정 Layer의 weight가 매우 커짐을 알 수 있는데 이는 해당 layer가 target task와의 distribution shift를 조절함에 있어서 많은 기여를 한다고 볼 수 있습니다.
본 논문에서는 특정 Layer만을 자동으로 선택해서 모델을 업데이트하는 Surgical fine-tuning 방법을 제안합니다. 저자들은 Auto-RGN, SNR을 통해 target task에 대한 distribution shift를 줄일 수 있고 이를 다양한 데이터셋에 대한 실험을 통해 보여줍니다. 기존에도 freeze를 통해 fine-tuning 하는 다양한 방법들이 존재하지만 이를 automatically 하게 조절하는 방법을 제안한 점과 특히 마지막 실험을 통해 target task에 민감하게 반응하는 layer를 업데이트하는 게 중요함을 보여준 점이 인상 깊은 논문이었습니다.