[Backbone] DreamTeacher: Pretraining Image Backbones with Deep Generative Models

 

 

Overview 

오늘 리뷰할 논문은 NVIDIA, University of Toronto, Vector Institute, MIT에서 낸 "DreamTeacher: Pretraining Image Backbones with Deep Generative Models" 입니다. ICCV 2023에 채택되었네요. 이 논문에서는 downstream image backbone을 사전 학습 하기 위해 generative network를 활용하는 DreamTeacher 기법을 소개합니다. 즉, 생성 모델의 knowledge를 특정 인식 작업을 위해 잘 설계된 standard image backbone으로 추출하는거죠. 어떻게 보면 역으로 사용했다고 보면 되네요. 

 

 

먼저 label-efficient semantic segmentation 결과를 보시면 아래와 같습니다. 최소 20개의 레이블이 지정된 이미지를 이용하여 ConvNX-B 라는 백본으로 mix-distillation을 사용하는 ADM generator를 사용했다고 합니다. 자세히 보시면 얼굴 영역을 segmentation 해놓기도 하고, GT와 별다른 차이 없이 잘 분할하는 모습을 볼 수 있습니다. 멋찌네요. 

 

본 논문에서는 knowledge distillation을 2가지 타입으로 나눕니다. 

 

1. ImageNet과 같은 큰 레이블이 지정된 데이터세트에서 이러한 백본을 사전 학습 하는 대신 학습된 generative feature을 image backbone으로 distilling 하는 방법

 

2. task head가 있는 generative network에서 얻은 레이블을 target backbone으로 distilling 하는 방법 

 

이런식으로 학습하게 되면 기존에 self-supervised representation learning 방식으로 학습하던 것 보다 훨씬 결과가 더 좋다고 합니다. 

 

 

DreamTeacher Framework

본 논문에서는 DreamTeacher를 2가지 시나리오인 unsupervised representation learning 에서는 사전 학습 중에 레이블을 사용할 수 없는 경우와 일부 레이블을 사용할 수 있는 semi-supervised learning으로 설명합니다.

 

학습된 생성 모델 $G$를 활용하고 학습된 feature를 target image backbone $f$로 추출합니다. $f$를 학습하기 위한 레시피는 $G$와 $f$의 시나리오와 선택 모두에서 동일하게 유지된다고 하네요. 먼저 이미지 $x_i$에 대해 feature dataset $ D = \left\{ x_i, f_{i}^{g} \right\}_i^N $를 만들고, 그에 해당하는 feature $f_{i}^{g}$를 생성 모델로부터 추출합니다. 그 다음 feature $f$를 $f(x_i)$의 intermediate feature로 추출하여 데이터세트 $D$를 사용하여 $f$를 학습합니다. 여기서 conv backbone $f$에 초첨을 맞추고, 향후 transformer에 대한 작업을 남겨둡니다. 

 

여기서 b는 Unsupervised Representation Learning을 나타내고, c는 Label-Guided Representation Learning을 나타냅니다.

Unsupervised Representation Learning

feature dataset $D$가 주어졌을 때 unsupervised representation learning을 위해 이미지 $x_i$에서 해당 생성 특징을 regression하기 위해 백본 $f$의 다른 hierarchical level에서 feature regressor를 연결시킵니다. 먼저 feature dataset 생성은 생성 모델 $G$에서 이미지를 샘플링하고 생성 프로세스에서 intermediate feature를 기록할 수 있다고 하네요. (보통 생성 모델에서 원하는 영상이 나오게 할려고 중간 feature에 대해 이미지를 샘플링 해보면서 체크포인트를 설정하곤 하죠) 이러한 방법은 무한한 크기의 데이터세트를 합성시킬 수 있지만, 생성 모델이 분포의 일부를 충분히 학습하지 못할 수도 있기 때문에 mode dropping 현상이 발생할 수 있다고 하네요. 위 그림 2에서 b 부분에 해당합니다. 

 

여기서 mode dropping은 mode collapsing 이라고도 하는데요, 아래 그림에서 MNIST를 예로 봤을때 숫자가 골고루 학습되지 않고, 아래 그림에서 이상한 숫자를 학습한 것과 같은 영향을 말합니다. 

 

https://jaejunyoo.blogspot.com/2017/02/unrolled-generative-adversarial-network-1.html

 

아무튼 이러한 데이터세트를 synthesized dataset이라고 하고, encode 과정을 통해 레이블이 지정되었거나 지정되지 않은 실제 이미지를 모델 $G$의 latent space로 encode 할 수 있습니다. 이러한 데이터를 encoded dataset라고 합니다.

 

먼저 synthesized dataset D는 생성 모델 $G(z)$로부터 샘플링되었고, 여기서 $z$는 생성모델 G의 prior distribution으로부터 샘플링된 값입니다. $G(z)$의 intermediate feature을 $\left\{ f_{l}^{g} \right\}_{l=1}^{L}$로 저장합니다. 여기서 $l$은 총 L level의 hierarchy level을 나타냅니다. 본 논문에서 주요 목표는 이미지 품질에 관계 없이 이미지를 feature에 매핑시킬 수 있도록 백본을 학습하는 것이기 때문에 bad sample에 대해서는 관심이 없습니다. 

 

encoded dataset은 latent variable $\tilde{z}$을 얻기 위해 encoding 과정을 사용하여 생성 모델의 latent space로 real image를 encoding하여 생성합니다. 그 다음 유사하게 생성과정을 실행하고 $G(\tilde{z})$에서 hierarchical intermediate feature을 저장하여 데이터세트를 얻습니다. 

 

Feature Regressor

generative representation $f^g$를 백본으로 추출하기 위해 이미지 백본의 feature를 generative feature로 매핑하고 정렬하는 feature regressor module을 설계합니다. FPN의 설계에서 영감을 얻어 multi-level feature을 가져오고, lateral(측면) skip connection이 있는 top-down 아키텍처를 사용하여 백본 기능을 fuse 하고, multi-scale feature를 생성합니다. 또한 PSPNet의 PPM(Pyramid Pooling Module)을 FPN branch 전 이미지 백본의 마지막 레이어에 적용하여 feature mixing을 향상 시킵니다. 

 

 

Feature Distillation

encoder $f$로 부터 각 intermidiate feature들을 $f_{2}^{e}, f_{3}^{e}, f_{4}^{e}, f_{5}^{e}$ 라고 정의하고, 각 대응하는 feature regressor output을 $f_{2}^{r}$, $f_{3}^{r}$,$f_{4}^{r}$,$f_{5}^{r}$ 이라고 정의합니다.

여기서 $f_{l}^{r}$과 $f_{l}^{g}$의 채널 수가 다를 경우에는 1x1 conv 를 사용하여 채널 수를 일치 시킵니다. 로스 함수는 다음과 같이 정의되며, intermediate feature activation을 모방하여 teacher로부터 student를 학습하는 FitNet에서 영감을 받았다고 하네요. 

 

 

여기서 W는 LayerNorm에서 구현된 non-learnable whitening operator이며, 이는 여러 계층에서 서로 다른 feature의 크기를 정규화 하고, layer number 2, 3, 4, 5는 input resolution에 대해 stride가 $2^l$인 feature를 나타냅니다. 또한 activation 기반 AT(Attention Transfer)를 이용합니다. 이는 spatial feature 마다 attention map을 증류하게 됩니다. 고차원 feature를 직접 regression 하는 것 보다 수렴 속도를 향상시킨다고 하네요. 

 

 

위 수식에서 Q는 벡터화 된 형태의 각 regressor 및 생성 모델 feature의 레이어에서 i, j 번째 쌍을 의미합니다. 실험에서 p=2 값을 이용하였다고 하네요. 최종 feature regression 로스의 형태는 다음과 같습니다. 여기서 lamda는 AT loss에 가중치를 주게 되며, 10.0을 사용했다고 합니다. 

 

 

 

Label-Guided Representation Learning

downstream task label의 일부를 사전 학습에 사용할 수 있는 semi-supervised 방법에서, 본 논문에서는 DatasetGAN 방법에 따라 frozen generative network G 위에서 feature interpreter라고 불리는 task-dependent branch를 학습합니다. 위 그림 2에서 c 부분에 해당합니다. 

 

Feature Interpreter

Big-DatasetGAN과 유사한 설계를 활용하여 더 나은 메모리 효율성과 예측 정확도로 설계를 개선합니다. 특히 interpreter는 생성기에서 multi lavel feature를 입력으로 사용하게 되며 이는 Feature Fusion Layer에 입력되어 feature dimension을 낮추고 다음 level feature와 합쳐 최종적으로 픽셀 당 logit을 출력하게 됩니다. 그리고 본 논문에서는 BigDatasetGAN의 interpreter 설계를 따르게 되고, conv fused block을 depth0wise separable conv, GroupNorm, Swish activation으로만 대체한다고 하네요. 또한 segmentation label을 사용하여 interpreter branch를 학습하고, 학습을 위해 cross-entropy와 Dice 로스를 사용한다고 합니다. 

여기서 $I_\Theta $는 feature interpreter의 weight를 뜻하고, H 함수는 pixel-wise cross-entropy를 의미하고, D는 Dice loss를 의미합니다. lamda 값은 dice의 hyperparameter라고 하네요! 

 

Label Distillation

본 논문에서는 label distillation을 위해 아래와 같은 loss를 사용한다고 합니다. 여기서 P들은 각각 feature interpreter로부터의 logit과 backbone $f$의 logit-head를 나타냅니다. H는 cross-entropy loss이고, τ 값은 output distribution의 sharpness를 조절한다고 하네요. 4.0 값을 씁니다. 

 

label distillation의 로스는 다음과 같습니다. 여기서 사용하는 lamda 값은 두 로스 사이의 값을 조절하기 위한 hyperparameter 이고, 1.0 값을 사용한다고 합니다. 

 

 

이렇게 학습한 결과를 ADM으로 확인하면 다음과 같습니다. time step T에 따라 결괄르 확인할 수 있는데, T가 증가할수록 값이 더 부드러워지고, low resolution block 에서는 human이나 car 와 같은 객체의 feature에 집중하는 경향을 보입니다. higher resolution block에서는 wheel과 headlight 같이 작은 부품에 집중하는 경향을 보입니다. 

 

 

Experiments

본 논문에서는 백본에 대한 KD 사전 학습을 위해 BDD100K[76], ImageNet-1k(IN1k-1M), LSUN[77], FFHQ[36] 데이터 세트를 사용하였다고 하고, 각 100k, 1.28 million, 10 million, 100k 개라고 합니다. 

 

ImageNet Evaluation and Transfer 

Imagenet Pretraining 

먼저 ImageNet pre-train을 위해 DreamTeacher의 효과를 검증합니다. SOTA 방식인 SparK 방법을 따르며 ConvNext-B와 ResNet50으로 두 개의 컨볼루션 아키텍처를 평가합니다. 일반적인 관행에 따라 ImageNet-1k에서 감독되지 않은 이미지 백본을 사전에 학습합니다. Transformer 기반의 self-supervised 방법과 비교하기 위해 SparK의 방법론을 따르고 최신 CNN을 사전 학습합니다. ConvNeXt는 비슷한 수의 매개변수를 사용하고, 공정한 비교를 위해 ResNet-50을 사전 학습하고 평가합니다. 

 

 

 

 

앞으로 이 논문의 내용 대로 생성 모델 기반 사전 학습 방법이 향후 vision task에서 많이 사용될 수 있을 것 같네요!

결과를 더 살펴보면 다음과 같습니다. 

 

 

 

 

 

 

 

 

 

Paper : https://arxiv.org/abs/2307.07487

DreamTeacher: Pretraining Image Backbones with Deep Generative Models

Project page : https://research.nvidia.com/labs/toronto-ai/DreamTeacher/

DreamTeacher