본문 바로가기

Study/Paper Review

[CV] DIVA: Domain Invariant Variational Autoencoders

Maximilian Ilse, Jakub M. Tomczak, Christos Louizos, Max Welling

PMLR 2020

논문 링크: https://arxiv.org/abs/1905.10427


Introduction

DNNs는 주어진 학습 데이터의 bias를 학습하게 되고, 이에 test 단계에서 다른 분포의 데이터가 주어졌을 때 성능 저하가 일어난다. Domain generalization의 목표는 test에서 이전에 보지 못했던 domains와 유사한 분포의 representations을 학습함으로써 test에서의 성능 저하를 극복하는 것이다. 본 논문에서는 독립적인 세 가지 latent space를 도입하여 VAE framework를 확장시킨 Domain Invarianty Variational Autoencoders(DIVA)를 제안한다.

Notation

p(x, y)는 joint distribution으로서 domain을 의미한다. S는 N개의 domain에서 추출한 샘플들로 이루어진 training set이고 input과 label 쌍으로 이루어져 있다. Domain generalization은 non-i.i.d. condition을 가정하며 입력 데이터로는 label과 unlabeled 모두 고려하고 있다.

Method

DIVA는 3개의 encoder와 단일 decoder로 구성되어 있다. input x가 주어지면 각 encoder는 domain, class, 나머지 residual variations에 대해 latent variable z를 학습한다. 도출된 latent variables는 decoder의 입력으로 주어져 input x를 복원하게 된다. 특히 DIVA는 단일 decoder 외에, 학습된 latent variable로부터 각각 domain과 class를 복원할 수 있도록 auxiliary 분포를 도입하였는데 이들이 마치 domain과 class의 decoder와 같이 동작한다.

Domain-invariant latent subspace를 위해서는, 독립적인 3개의 latent variabl이 필요하다. Domain-specific zd, class-specific zy, 그 외의 input x에 대한 특징들을 학습할 zx가 그에 해당한다. zx는 Gaussian prior를 가정하고, zd~Pd와 zy~Py는 conditional prior를 따른다. Domain이 i.i.d.가 아니라고 가정했으니, domain d가 조건부로 주어질 때 마다 그에 따른 latent variable zd의 분포도 달라지게 되는 것을 반영했기 때문에 사전분포가 조건부 분포의 형태를 따르게 되는 것이다. Label 또한 개별적으로 개수가 다르니 분포가 다를 것이다. Encoder는 input x가 주어졌을 때 각 latent variable이 따르는 확률분포이고, decoder는 각 latent variable이 주어졌을 때 x에 대한 확률분포이다.

DIVA의 loss는 위와 같다. 각 항들에 대해 살펴보면, 붉은 첫번째 항은 encoder의 출력이 주어졌을 때 입력을 복원하는 확률을 높이는 reconstruction term이다. 다음에 이어지는 KL divergence term은 각각 들어오는 domain, class에 대한 latent variable의 분포와 x로부터 학습되어지는 인코더의 분포가 가까워지게 하고, x로부터 학습되는 인코더의 분포와 가정했던 gaussian prior의 거리를 축소한다.

추가적으로 zd, zy가 x로부터 구분을 이루며 학습할 수 있도록 auxiliary term을 추가하였다. VAE의 decoder처럼 zd가 주어졌을 때 domain을 잘 복원하고, zy가 주어졌을 때 label을 잘 복원하도록 확률을 높이는 term이다. domain과 class에 대해 decoder를 따로 구성하진 않았지만 보조적인 loss를 추가함으로써 Qwd,qwy도 decoder와 같은 역할을 수행하게 된다.

Semi-supervised condition에서 DIVA의 loss이다. Reconstruction term과 domain 및 input x의 residual variations에 대한 KL divergence term은 이전과 동일하다. 앞선 식에서 label y가 사용된 term들은 highlight된 부분과 같이 변형되었다. 먼저 첫 번째 식은 class에 대한 latent variable의 분포와 x로부터 학습되는 인코더의 분포 사이 거리를 축소하는 term이다. 이는 기존의 KL divergence항을 풀어서 쓴 형태인데, y가 주어지지 않으므로 기존의 식을 유지할 수 없어 x에서 zy를 샘플링하고 그 zy로 y를 샘플링하여 KL term을 대체하였다. 두 번째 식은 auxiliary에서 y를 복원하는 항이, 복원할 y가 없으니 샘플링된 y들의 분포와 거리를 줄이는 것으로 대체되었다.

Experiments

상단 행의 embeddings은 domain에 따라 채색되었고, 하단 행은 class에 따라 채색되었다. 첫 번째 열과 세 번째 열에서 각각 domain과 class에 대해 잘 군집이 잡혀있음을 알 수 있고 domain은 class와 무관하게 class는 domain에 무관하게 군집된 것을 볼 수 있다. 두 번째 열의 상단 plot은 회전 각도(domain) 따라 색상 지정되어있는데 zx[0]는 MNIST 숫자의 회전 각도와 상관관계가 있는 것으로 관찰되었고, 이는 zx가 domain 정보의 일부(회전 각도)를 포착하고 있음을 나타내는 것이다. 하단 plot은 class(숫자)에 따라 색상이 지정돼있다. zx[0]은 숫자의 선 두께와 상관관계가, zx[1]은 숫자의 너비와 상관관계가 있다. 특히 class 1에 해당하는 embedding이 plot 좌측 하단에 군집을 이루고 있음을 관찰할 수 있다고 저자들은 말하고 있다. Class나 domain과 직접적 관련 없는 residual variations을 포착하도록 설계된 zx가 실제로는 회전각도, 선 두께, 숫자 너비 등 이미지의 스타일적 특성을 포착하고 있다. 이는 클래스나 도메인과 완전히 독립적이지도 종속적이지도 않다고 할 수 있다.

DIVA의 분리 능력에 대해 더 살펴보기 위한 추가 실험이다. 학습된 y의 조건부 사전분포로부터 샘플된 z'y와 잠재 변수 zd,zx를 이용해 학습된 decoder로 새로운 이미지를 생성하였을 때, zd에서 encoding된 회전각과 zx에서 encoding된 선 굵기가 잘 보존됨과 동시에 class 변화가 이루어졌음을 확인할 수 있었다.

학습된 조건부 사전분포로부터 샘플링된 z'd와 잠재변수 zx, zy를 이용해 이미지를 생성하였을 때에도, class와 zx에서 encoding된 선 굵기 등이 잘 보존됨과 동시에 회전각의 변화가 이루어졌음을 확인할 수 있었다.

기존의 방법론들과의 test accuracy 및 1,3,5,9배수만큼 unlabeled data를 추가한 데이터셋에 대해 실험한 결과이다. Unlabeled data의 비율이 높아질수록 test 정확도의 평균값이 소폭 감소, 분산이 소폭 증가하는 것을 볼 수 있다. 이에 대해 저자들은 unlabeled를 많이 추가할수록 loss term의 불균형으로 성능 저하를 이끌 수 있다고 설명한다.

Unseen domain이 입력으로 주어지는 상황을 가정하여, unlabeled data로만 이루어진 domain을 train set에 추가한 상황이다. Semi-supervised가 성능이 조금 더 높은 것을 보아 unlabeled data가 일반화 성능 향상에 도움을 주었음을 알 수 있었으며, 60도 회전각에서 공통적으로 더 낮은 성능을 보이는데 30도는 주어진 labeled 도메인들로부터 잘 학습되었지만 60도는 변화의 끝에 위치해있기에 해당 domain에서 unlabeled를 활용한 모델이 일반화 성능 향상에 도움을 주고 있음을 확인할 수 있다. 60도에서 더 큰 성능 향상이 일어난 것은 unlabeled를 통해 더 큰 회전각에 대한 정보를 학습하였기 때문인 것으로 판단된다.

Lemon box를 통해 재구성된 이미지에 구멍이 나타나고 있음을 알 수 있다. 이는 모델이 zx만으로는 이미지의 일부 정보를 재현할 수 없다는 것을 의미하고 Decoder의 관점에서는 zd와 zy의 확률 질량이 0임을 가리킨다. 이는 이 변수들이 0이 될 확률이 매우 낮거나 없다는 것으로, domain 정보(zd)와 class 정보(zy)는 항상 어떤 특정한 값을 가지며, "정보가 없음"을 의미하는 0은 거의 발생하지 않는 것을 확인할 수 있다. 즉, zx만으로 재구성할 때 구멍이 생기는 이유는 원래 이미지의 일부 특성이 zd나 zy에 encoding되어 있기 때문이며, 이는 시각화에서 확인한 zx의 domain과 class 연관성과 일치한다.

이상적으로는 모델이 모든 domain에서 일관된 성능을 보이는 것이 좋고 이는 모델이 도메인 불변 특성을 학습했다는 증거가 될 수 있다. 그러나 green box처럼 정상 세포도 pink로 염색된 경우, 이는 실제 감염된 세포와 구별하기 어려워진다. 이는 false positive (감염되지 않았는데 감염된 것으로 오판) 또는 false negative (감염되었는데 정상으로 오판) 의 가능성을 높인다. 이는 세포 전체가 pink로 염색되었거나, 정상 세포 내부에 pink 영역이 생겼을 가능성을 모두 포함할 수 있다.

모델이 제대로 학습하지 못한 2개의 도메인 중 하나를 unlabeled data로, 나머지 하나를 test domain으로 설정하고 실험을 진행한 결과이다. 유사한 domain의 unlabelded data만으로도 성능 향상이 가능한지 확인함으로써, 모델의 데이터 효율성을 테스트한 것으로 보인다. Dataset에서 어려운 domain을 학습하고 test 시 타 방법론과 비슷하게 높은 성능을 보였다. Unlabeled data가 추가된 실험에선 기존 방법론과 비교해 unlabeled data를 사용가능한 점과 비교적 향상된 성능을 확인할 수 있다. Label이 없이도 유사 domain에 대한 학습이 잘 이루어졌음을 확인할 수 있다. Test domain과 비슷한 domain을 학습하지 못했을 때는 비슷하게 저조한 결과를 보였다.

각각의 latent variable이 잘 학습하고 있는지를 확인한 실험이다. zd 혹은 zx에서 학습한 y classifier가 항상 다수 클래스를 예측하는 모델보다 성능이 낮다. 앞선 실험과 종합하여 Domain specific과 residual variations specific을 y는 y specific을 잘 분리해서 학습했음을 알 수 있다. zy에 성공적으로 class 관련 정보가 encoding되었음을 알 수 있다.

Conclusion

DIVA는 세 가지 독립적인 latent subspace들을 학습하여 분리된 표현을 학습할 수 있다. 개별적으로 분리된 latent subspaces가 각각에 특화된 표현들을 학습하여, input을 성공적으로 복원함과 동시에 unseen domains의 일반화에 있어 기존과 유사하거나 능가하는 성능이 나타난 것을 확인하였다. 또한, additional unlabeled domains가 외삽 시 모델 성능을 향상시키는데 도움이 되었음을 알 수 있었다.


Loss term 전개를 설명하기 위해 공부했던 경험은 분명 어려웠지만,

이해하는 과정에서 성취감을 느꼈던 기억이 난다. (´▽`)

반응형