연금술사

요약

SimCLR에서는 Instance discrimiation이 학습목적으로 원본이미지를 서로 다른 증강방법을 이용해도, 서로 같은 임베딩이 되게끔 유도한다. CSI은 SimCLR방법과 매우 유사한데, 원본이미지를 증강한 경우만 OOD(Out Of Distribution)으로 학습하는 방법이다.

• 이미지 증강방법들의 집합 $S:=\{S_{0}, S_{i}, ..., S_{K-1}\}$
• 동일 이미지 반환: $I=S_{0}$

여러 이미지 증강방법S로부터 하나를 뽑아, 이미지 모든 이미지를 증강하고(원본반환 포함)이를 SimCLR을 돌림. 이 과정을 여러 증강방법에 대해서 반복함. 

추가적인 학습 테스크로, 증강된 이미지가, 어떤 증강방법으로 이용되었는지를 분류하는 방법도 진행

최종학습 Objective은 증강에 따른 이미지를 OOD Image로 SImCLR하는 목적과 증강방법론 분류모델을 혼합하여 예측하게 됨

Preliminary

• Positive sample: Origin이 동일한 이미지(또는 데이터포인트)
• Negative sample: Origin이 다른 이미지(또는 데이터포인트)
• Dictionary look-up task: 여기서는 contrastive learning을 의미합니다. Query image에서 하나의 positive이미와 나머지 (N-1)개의 negative이미지와 유사도를 계산하는 과정이기에 k-value와 같은 dictionary구조라고 일컫습니다.
• InfoNCE: InfoNCE은 미니배치 K+1개에서 1개의 positive sample을 찾는 softmax 함수 + temperature scale추가

Methods: 딕셔너리 사이즈를 키워 negative samples을 포함할 수 있으면 좋은 학습방법이 될 것

Moco은 딕셔너리 사이즈를 키워, Negative key을 많이 갖고 있게되면 더 학습이 잘된다는 가설로 시작됩니다. 위의 딕셔너리가 가변적(dynamic)합니다. 하나의 query에 대해서 key도 augmentation을 달리하고, 인코더도 점차 학습되기 때문입니다. MoCo은 Negative sample을 임베딩하기에 Key encoder가 비교적 천천히 학습되고(=급격히 변하지 않고) 일관적으로 학습되는 것을 목표로 합니다.

1. MoCo은 샘플을 Queue방식으로 처리합니다. 즉, 새로운 샘플을 Queue에 넣으면, 먼저 들어가있던 오래된 샘플은 제외되는 방식입니다. 샘플을 넣는 과정을 "Enqueue" 제외하는 과정을 "dequeue"라 합니다. 따라서, 비교적 최근의 인코딩된 Key Image들은 enqueue후에 몇 회 재사용될 수 있습니다.

2. Momentum update: key encoder을 직접학습하는 대신에, query encoder을 학습하고 학습의 일부만 반영(=momentum)합니다. 이유는 key encoder에는 많은 negative samples에 대해서 backpropagation하기 어렵기 때문입니다즉, computational cost문제 때문입니다. 예를 들어, 이번 배치가 16개였고, 딕셔너리사이즈가 1024개면, 1,024개의 샘플에 대해서 backpropgation을 해야하기 때문에 계산량이 많아 현실적으로 어렵다는 것입니다. enqueue하는 방식으로는 back-propagation이가능하다는 글을 같이 공유합니다(URL, URL2) (보통, "intractable"이라 하면 미분이 불가능한, 다루기 어려운 이런 뜻인데 미분과 연관되어 쓰여 혼란이 있습니다)

또한 인코더가 급격하게 변하는 경우, 딕셔너리에 저장되어있던 기존의 인코딩의 값이 현재와 많이 상이하기에 천천히변화시켜야한다는 것입니다.

그래서 아래와 같이 momentum update을 대안으로 선택합니다. 아래의 수식(2)을 보면, query encoder의 파라미터의 내용이용하여 key encoder을 업데이트하는 방식으로 학습합니다.

$\theta_{k} \leftarrow m \theta_{k} + ( 1- m )\theta_{q}$ (2),

• $\theta_{q}$: query encoder의 파라미터
• $ \theta_{k} $: key encoder의 파라미터
• $m$: 모멘텀계수. [0, 1]사이의 값

따라서, query encoder만 back-propagation을 진행합니다. 그리고, 급진적인 key encoder의 파라미터 변화를 막기 위해서, m값을 0.999로 학습합니다. 이는 m=0.9보다 성능이 더 좋았다고 합니다.

기타1: contrastive loss 학습방식

• end-to-end update: query encoder와 key encoder가 동일한 경우를 의미합니다. 그리고, 딕셔너리는 항상 미니배치 사이즈와 같습니다. 따라서, 딕셔너리 사이즈를 엄청 키우진 못합니다.
• memory bank: 메모리뱅크는 데이터셋에 모든 샘플을 일단 인코딩하여 넣어둡니다. 그리고, 매니배치에서 일부를 샘플링하여 contrastive learning에 사용합니다. 장점은 "큰 사이즈의 딕셔너리"라는 것입니다. 단점은 메모리뱅크의 업데이트이후 key encoder가 여러번업데이트되기 떄문에, 둘의 일관성(인코딩 차원에서의)이 다소 떨어지는다는 것입니다.

Linear probe 과정에서 learning rate을 30으로 한 이유

논문에서는 "optimal learning rate을 30, weight decay을 0"으로 진행합니다. 일반적으로 다른 hyper-parameter인데, 이 이유는 feature distirbtion ImageNet supervised training과 Instragram-1B와 크게 다름을 의미합니다.

Motivation

1. 최근 연구들은 언어모델을 차용하여 비전문제를 풀려는 아키텍처가 많이 발표되고 있으며, 대표적으로 ViT(Vision Transformer)가 있습니다.
2. (2021년 당시) Vision Transformer을 기반으로한 모델들은 ViT을 인코더 백본으로 특징을 뽑고, CNN을 여러겹 적용하여 decoders로 하여 꽤 좋은 성능을 보여주었습니다.
3. 위의 아키텍처는 아래의 2문제가 있어, Segformer가 해결하고자 했던 것 입니다.
   1. Single scale low-resolution: 하나의 스케일의 feature만 사용된다는 것입니다. 인코더에서 이미지를 특징화하고, 디코더에서 하나의 입력스케일만 입력받습니다.
   2. Computational cost: 이미지가 클수록 계산량이 많다는 것입니다. 특히, 트랜스포머는 token 수 N의 N(O^{2})의 연산을 갖기에, 이미지 사이즈가 크면 연산량이 너무 많이 듭니다.

Segformer 장점

• Positional encoding(PE) 사용하지 않습니다. PE을 안쓰고, Zero padding과 CNN으로 위치정보를 포함시켰기 때문입니다. 덕분에, 여러 크기의 입력이미지를 수용할 수 있습니다. 
• 계층적 구조: Transformer에 들어가는 입력을 점진적으로 작게 나눠, 앞단에는 큰(거친, coarse)특징들 (예, 위치정보)를 뽑고, 후미의 블록에서는 상세한 특징(Fine)을 뽑도록 아키텍처가 구성되었습니다.
• 경량 All-MLP 디코더: 디코더에 CNN하나없이 MLP로만 구성되어, 더 큰 수용장(Effective Receptive Field)을 가집니다. 따라서, 이미지 분절시에 멀리 떨어진 픽셀정보를 고려할 수 있습니다.

아키텍처 구성

Segformer은 크게 2가지의 구성입니다. 거의 모든 Segmentation 아키텍처와 동일하게 인코더와 디코더 구조를 갖습니다.

1) 인코더(Encoder): 계층적 트랜스포머(Hierachial Transformer). 트렌스포머의 패치크기를 다양하여 입력받습니다.

2) 디코더(All-MLP decoder) : 1)의 다양한 level의 features을 입력으로하는 디코더로, 최종적으로 $\frac{H}{4} \times \frac{W}{4} \times N_{cls}$의 이미지를 반환합니다. 이미지가 4배 작아져 다운샘플링 됩니다.

상세한 구성은 아래와 같습니다.

1. 입력 이미지: $H \times W \times 3$

2.계층패치화: 이미지를 $\frac{H}{2^{i+1}} \times \frac{W}{2^{i+1}} \times C_{i}$으로 패치화합니다. 예를 들어, 첫번째 트렌스포머 블럭에서는 패치(토큰)하나당 256/4 x 256/4 x C, 이후에는  256/8 x 256/8 x C2 이런식입니다.

3.패치 결합(Overlapped Patch Merge):

• 결합인데, 사실은 패치를 결합된 상태에서 분리한다는 것이 더 이해가 쉽습니다. $ N \times N \times 3$ 패치를 $ 1 \times 1 \times C $벡터로 표현이 가능합니다. 일단 백터화해놓으면, 이후의 인코더블럭들에서는 $ 1 \times 1 \times C_{i} $ 이든, $ 1 \times 1 \times C_{i+1} $이든 표현이 가능합니다. 일단 벡터화를 하면 아래와 같이 패치를 쪼갤 수 있습니다.
• 예를 들어,  첫 번째 입력이미지를 4x4로 쪼갠경우, 한 패치당 $\frac{224}{4} \times \frac{224}{4} \times C_{1}$ 크기를 갖습니다. 이 입력크기를 F1이라고 합니다. 이 F1을 입력크기를 다시 더 작은 단위의 크기로 쪼개서($ \frac{224}{8} \times \frac{224}{8} \times C_{2} $) 트랜스포머의 입력에 넣습니다. 트렌스포머 블럭1번째에서 1개 패치가 2번째트렌스포머 블럭의 4개가 되는것입니다.
• 이런 패치과정을 겹치지않게(non overlapping)하게 하면 패치간의 연속성이 떨어져서, 겹쳐가면서 진행합니다.  (사실 정확히 4개가 아닐 수 있습니다)

4. Efficient Self-Attention

어텐션은 패치 수(토큰 수)에 영향을 받습니다. 정확히는 패치수(N)의 제곱입니다. Segformer을 이 N을 R비율 맞큼 줄입니다.

N=W x H 라고 할 때,

1. $\text{ Reshape }(\frac{N}{R}, C \cdot R)$: 연산할 이미지의 수를 R배만큼 줄이고, 특징값을 R배만큼 늘립니다.
2. $Linear(C\cdot R, C)$: R배만큼 늘어난 이미지 특징값을 C로 다시 선형변환해서, 압축합니다.

이러면 원래 $O(N^{2})$ 이던 계산량이 $O(\frac{N^{2}}{R})$으로 변경됩니다.

5. Mix-FFN: CNN과 Feed-forward network(FFN)을 단계적으로 진행하는 모듈입니다.이미지 분할(Segmentation)에서는 사실 포지셔널 임베딩 필요없다고 저자들은 주장하는데요. 제로패딩만 어느정도해도 위치정보가 들어간다는 연구가 있기 때문입니다. 

$x_{out} = MLP(GELU(Conv_{3 \times 3}(MLP(x_{in}))) + x_{in}$

• $x_{in}$: self-attention 에서 나온 피쳐입니다.

6. All-MLP Decoder: CNN 보다 더 큰 수용장(Effective recept field, ERF)을 갖기위해서, 일부러 CNN을 사용하지 않고 MLP layer을 구차했습니다. 문제는 Transformer block에서 각기 다른 token(이미지)사이즈를 입력받아야하는데 아래와 같이 처리했습니다.

1. 채널 수 맞춤: $\hat {F_{i}} = Linear(C_{i}, C)(F_{i}), \forall i $. 모든 i(Transformer block의 output)에 대해서 채널을 C로 맞춥니다. 이렇게하려면 MLP가 4개가 있어야겠죠.
2. 업샘플링(upsampling): 모든 피처맵(F_{i})을 $\frac{H}{4} \times \frac{W}{4} $ 크기로 업샘플링을 합니다. 업샘플링은 Bilinear interpolation또는 Nearest interpolation로 해볼 수 있습니다.
3. 결합(Concatenate) 후 Linear: ${F = Linear(4C, C)(Concat( \hat {F_{i}})), \forall i $ 2.에서 나온 모든 피쳐맵을 한축으로 concat합니다. 그 이후 Linear 채널을 교정합니다.  
4. 분류: $M = Linear(C, N_{cls})(F)$ linear layer을 하나 넣어 클레스 수에 맞게 분류합니다.

아래 그림으로 보면 더 이해가 쉽습니다.

결과 1: 수용장이 더 커진게 맞나?

네, 수용장이 커졌습니다.

아래는 Cityscapes데이터셋에서의 100개이상의 이미지의 예측에서의 수용장을 시각화 한 것입니다. 실제로는 output값을 역전파해서 어떤 픽셀이 gradient가 컸는지 계산했을 것입니다. (https://ahmdtaha.medium.com/understanding-the-effective-receptive-field-in-deep-convolutional-neural-networks-b2642297927e)

DeepLabv3+ 모델과 SegFormer을 비교해봤을 떄, 더 큰 수용장을 보입니다. 특히, stage 4에서는 DeepLabV3+가 중심부에만 몰려있습니다.

모델사이즈에 따른 성능은?

인코더와 디코더 사이즈를 키울수록 mIoU지표는 확실히 좋아집니다.

PE 대신에 Mix-FFN(CNN+MLP)을 쓴게 더 좋은가?

포지셔널인코딩은 이미지 사이즈가 커지면, 포지셔널 인코딩을 Interpolation해서 사용해야하는데, 이 Interpolation때문에 성능열하가 존재한다고 이미 알려져있고, 그 결과도 재현되었습니다. PE로는 성능차이가 크게 떨어반해, Mix-FFN으로는 성능 열하가 적게 보여집니다.

Discussion

Segformer는 positional encoding 없이도 다양한 입력 이미지 크기를 처리할 수 있는 유연성을 가지며, 계층적 구조와 효율적인 self-attention 메커니즘을 통해 계산 효율성을 극대화합니다. 또한, 경량의 All-MLP 디코더를 채택하여 더 큰 수용장을 확보함으로써, 이미지 분할과 같은 작업에서 더욱 정교한 결과를 보였습니다.

Motivation

CE(Cross-entropy)은 지도학습의 분류에 주로 사용됩니다. 하지만, 많은 단점들이 존재하는데, 예를 들어 noisy label이나 poor margin 같은게 있어 일반화 성능이 떨어집니다. CE의 대안으로 나온 여러가지 손실함수가 발명됬지만, 여러 챌린지에서 보면 여전히 CE을 쓰듯이, 실무에서는 큰 도움은 안됩니다.

최근에 대조학습(Constraive learning)으로, 라벨이 없이도 자기지도학습에서 사용됩니다. 미니 배치 내에서, 앵커와 같은 origin data은 가깝게, 앵커와 먼 데이터는 멀게 학습하는 방법으로 학습합니다.

이 논문은 자기지도학습에서의 라벨을 이용하여 contrastive learning을 이용해서 학습하는 방법을 제안했습니다. 같은 클레스의 임베딩을 정규화한다음에 가깝게 위치시키고, 다른 클레스보다 멀게 학습시킵니다. 앵커하나에 여러 positive pair을 만들 수 있다는게 장점이고, 많은 negative pair도 만들 수 있습니다.

Method

1. $\text{Aug()}$: 한 이미지 x로부터 이미지를 증강하는 모듈입니다. $\tilde{x}=Aug(x)$와 같이 $\tilde{x} $가 2개씩 생성됩니다. 각 $\tilde{x}$은 x로부터 새로운 view(증강법에 따른 변형)을 의미합니다. crop일수도, random rotation일 수도있습니다.
2. $\text{Enc()}$: 인코더 파트입니다. 벡터표현을 위해서 인코더를 하나 만듭니다. 이 인코더는 1에서 만들어진 2개의 증강된이미지를 포워딩하여 벡터를 만듭니다. 벡터의 차원은 2,048입니다. $r=Enc(x)\in R^{D_{E}}$
3. $Proj()$: Projection network입니다. 벡터표현이후에 128차원의 벡터로 프로젝션합니다. 이 벡터는 단위의 길이를 가질 수 있도록(unit hypersphere), 정규화를 합니다. Projection layer은 SSL에서는 사용하지 않습니다.

1. Contrastive loss function의 정의: 

1. $\{x_{k}, y_{k}\}_{k=1,...N}$: 지도학습에서는 N개의 임의의 샘플이 있다고 가정합니다. N은 배치사이즈입니다.
2. $\{\tilde{x_{l}}, \tilde{y_{k}|\}_{l=1,...,2N}$: 1.에서의 이미지를 각각 2개의 이미지 증강을한 결과를 의미합니다. 여기서 라벨은 $y_{k}$은 증강후에 $ \tilde{y_{2k-1}, \tilde{y_{2k} $로 되니, 다 같은 값입니다. 2N은 증강후의 배치입니다. 본문내에서는 "multi-viewed batch"라 합니다.

2. Self-supervised Contrastive Loss의 정의

Contrastive learning과 Supervised constrastive learning을 비교하면 아래와 같습니다. 라벨이 있으니, 같은 클레스끼리는 가깝게 학습시키게끔 변형되었습니다.

• $P(i)$: 은 i랑 동일한 이미지는 아닌, positive 의 집합니다.
• $A(i)$: 은 i(Anchor)가 아닌 나머지의 집합입니다.
• $z_{i}$:은 Projection 이후의 임베딩값입니다.
• $\tau$: temperature scaling

실제 논문에서 제안된건 아래와 같이 |P(i)|의 개수를 어디서 연산하냐에 따라서, 아래와 같이 2가지로 나뉩니다. 두 함수는 완전 동치는 아니라 각각 실험에서의 최적화된 방법을 사용하면 됩니다.

Results: 엄청좋아지나..? 그정도까진...

1. Top 1 classification 성능은 그냥 CE을 쓰나 SimCLR을 쓰나 엄청난 차이를 보이지 않습니다. 하지만 이 영역대에서 x%p올리기 쉽지않은데, SupCon으로 한번 최적화해보는건 좋은 선택같습니다.

2. 그래도 강건한 모델:

mCE라는 지표로 N은 이미지의 왜곡방법의 수를 의미하며, 분자 분모는 오차율입니다. mCE은 낮을수록 이미지 변환에도 강건하고 오차율이 없다는 것을 의미하는데요. 

$mCE = \frac{1}{N} \sum_{c=1}^{N} E_{c \text{baseline}} E_{c}$

여기서 corruption 이 정확히 어떤 이미지 변환인지는 표기는 안되어있으나, 이미지 corruption을 강하게 하더라도 종전모델보다 강건한 이미지를 나타냅니다.

얼마나 학습시켜야하나?

1. ResNet을 CE로 학습시키는데 1400정도 에폭을 활용: "Since SupCon uses 2 views per sample, its batch sizes are effectively twice the cross-entropy equiv-
alent. We therefore also experimented with the cross-entropy ResNet-50 baselines using a batch size
of 12,288. These only achieved 77.5% top-1 accuracy. We additionally experimented with increas-
ing the number of training epochs for cross-entropy all the way to 1400, but this actually decreased
accuracy (77.0%)"

2. 200에폭도 충분함

"The SupCon loss was trained for 700 epochs during pretraining for ResNet-200 and 350 epochs for
smaller models. Fig. 4(c) shows accuracy as a function of SupCon training epochs for a ResNet50,
demonstrating that even 200 epochs is likely sufficient for most purposes."

Image to image translation은 두 이미지 사이의 연관성을 학습하는 방법입니다. 보통은 1) 두 이미지의 짝을 지은 데이터가 가 구하기 어려워서 image to image을 하기 어렵고, 2) 한 이미지를 꼭 반드시 하나의 이미지에만 짝을 지을 필요가 없어 데이터 구성이 매우 어렵습니다.

위의 1), 2)의 예시인데, 이렇게 이미지를 짝을 지어야하는 경우에 이 데이터를 짝짓기도 어렵고, 짝을 짓더라도 다양한 경우가 많아서 golden standard로 짝을 지었다고 보장하기 어렵습니다. 

방법론

DRIT++은 2개의 임베딩을 나눠 만들어내는데, 1) domain invariant content space, 2) domain-specific attribute space 을 나눠 만들어냅니다.

구성은 크게 3파트입니다.

1. 컨텐츠 인코더(content encoder): domain invariant content space을 만들어내는 인코더가 하나 필요합니다.
2. 속성 인코더(Attribute encoder): 속성 인코더도 하나 필요합니다. 속성 인코더가 만드는 임베딩 벡터는 컨텐트 임베딩 벡터와 공유하지 않습니다. 단, X 도메인, Y도메인에 쓰이는 인코더는 동일 합니다(shared).
3. 제너레이터(Generator): 컨텐츠 인코더가 만든 벡터와, 속성인코더가 만든 벡터를 입력으로하여 이미지를 생성합니다.
4. 판별부(Domain discrimator): 합성된 이미지와 진짜 이미지를 판별하는 모듈입니다.

테스트 타임에는 attribute vectors을 prior Gaussian distribution을 따른다고 가정하고  N(0,1)의 attribution vector을 주입합니다.

요약

iBOT(masked image modeling)과 DINO(knowledge distillation with no label)방법을 혼합한 방법

Transactions on Machine Learning Research (01/2024)

Data processing: LVD-142M dataset의 생성과정

많은 양의 데이터 학습을 위해, 데이터 처리방식을 1) 라벨, 비라벨 이미지 수집, 2) 중복이미지 제거, 3) 이미지 검색

1. 데이터 수집: 데이터 수집은 라벨링된 데이터와 라벨링 안된 데이터 크게 2가지의 데이터셋을 구축했습니다. 라벨링 된 데이터는 classifcation, segmentation, depth estimation 등의 다양한 문제에 사용되는 공공데이터를 수집했습니다. 다른 한편으로는 웹크롤링을 통해서 <img>테그가 있는 경우를 모두 가져오고, PCA hash dedupliation 등을 이용해서 이미지의 중복을 제거했습니다.
2. 중복제거(de-duplication): 비슷한이미지만 모여있으면, 이미지의 다양성을 학습하기 어려우니, 유사한 이미지들을 제거하기위해서 중복제거를 image-copy-detection알고리즘을 이용해서 처리했습니다.
3. Self-supervised image retrieval: 라벨링된 데이터셋과 라벨링안된 데이터셋에서도 매우 비슷한 이미지가 있을 수 있기때문에,  이를 제거해야합니다. 첫 번째로, 이미지를 사전학습된 ViT-H/16로 "비정제 /정제 "임베딩 시키고 각 이미지를 코사인유사도로 유사성을 평가합니다. 그 후, K-means clustering을 이용해서 "비정제" 데이터를 군집화합니다. 하나의 정제 이미지의 벡터를 이용해서, 가장 가까운 N(4)의 이미지를 검색합니다. 1개의 정제데이터에 대해서, 비정제 데이터 4개의 유사한 Cosine similarity을 확인하고, 유사도가 낮으면(=이미지가 충분하지 않다면), 쿼리 이미지가 속한 클러스터에서 추가로 M개를 샘플링해서 보완합니다.

실제 구현은 vector DB로 쓰이는 Faiss 을이용해서 진행했습니다. 약 20개의 노드에서 위 과정을 분산처리하여 LVD-142M dataset을 생성합니다.

Discriminative Self-supervised Pre-training

이 연구에서 제안하는 방법론은 DINO와 iBOT loss을 혼합하는 방법론입니다. 이를 혼합하기 위해서, Student모델과 Teacher model은 각각 2개의 head을 가지고 있게됩니다.

1. Image-level objective: DINO v1에서 사용했던 것과 같습니다. Student, Teacher 모델이 있고, Student model이 반환하는 vector 와 teacher model이 반환하는 vector가 유사해지는 것을 목적으로합니다. (+centering + moving average에 대한 설명). 그 후, 파라미터는 EMA(Exponential moving average)로 student model에서 teacher 모델로 반영합니다.

$L_{DINO}=-\sum p_{t}log p_{s}$

2.Patch-level objective: iBOT을 이용한 자기지도 학습입니다. iBOT은 MIM(Masked image modeling0)으로 이미지를 더작은 패치로 나눈다음에 훈련하는 방법을 의미합니다. 이렇게 만든 저 작은 패치를 랜덤으로 마스킹합니다. (Fig 2. iBOT) 마스킹된 경우 마스킹 토큰만 줍니다. 

1. $\textbf{x}=\{x_{i}\}^{N}_{i=1}$: 이미지를 작은 토큰으로 생각해서, 패치단위로 N개로 잘게 쪼갭니다.
2. $\textbf {m}^{N}$: 마스킹인덱스입니다. 0 또는 1입니다. N개의 패치마다 마스킹할건지 여부를 결정합니다. $m_{i}=1$이면 마스킹된 상태입니다.
3. $\tilde{x}\triangleq \{\textbf {x}_{i} | m_{i}=1\}$: 마스크된 이미지를 의미합니다. m=1인 경우, 이미지들이 마스크 토큰으로 변경됩니다.
4. $\hat{x}\triangleq  \{\tilde{x}|(1-m_{i})x_{i} +m_{i}\textbf {e}_{[MASK]}\}$: N개의 패치가 마스크 토큰 또는 보존

학습의 목적은 마스크된 토큰($\tilde{x}$)을 복원하는 학습을 합니다. 아래의 그림 Figure 2에서, VIT(student)에는 마스크된 토큰을, teacher 모델에서는 전체의 패치를 다 본 후, masking 된 패치만 가져와 두 분포가 동일하도록 하는 것을 원합니다.

$L_{MIM} = -\sum _{i=1}^{N}m_{i}\cdot P_{\theta}^{patch}(u_{i})^{T} log P_{{\theta}}^{patch}(\hat{u_{i}})$

위 식에서 $m_{i}$가 곱해져서 반드시 마스킹된 경우만, 임베딩값을 복원하는 문제로 변경됩니다. 여기서도 DINO와 마찬가지로 softmax, centering이 쓰입니다. 논문에서는 아래와 같이 간략하게 표현했습니다. 

$L_{iBOT}=\sum_{i} p_{ti}logp_{si}$ (i:masked patch index, t: teacher model, s: student model)

3. 두 해드의 파라미터는 공유: Zhou et al(2022)에 논문에선서는 DINO와 iBOT의 해드 파라미터를 공유하는게 더 나은 성능을 보였으나, 본인들이 실험해보니 따로두는게 더 나았다고합니다.

4. Sinkhorn-Knopp centering: DINO에서는 teacher모델의 표현을 centering 후의 temperature scaling을 진행합니다. DINO v2에서는 SK(Sinkhorn-Knopp) batch normalization을 사용하는게 더 낫다고하여 Sinkhorn-knopp algorithm을 진행합니다.

5. KoLeo regularizer: 배치내에서 피쳐가 균등하게 만들기위해서 KoLeo regularizer을 사용합니다. KoLeo 정규화전에 각 벡터에 L2 norm을 진행합니다. 

$L_{koleo}= -\frac{1}{n}\sum_{i=1}^{n}log(d_{n,i})$,

$d_{n,i}=min_{j=i}||x_{i}-x{j}||$은 벡터i와 배치 내에 어떤 벡터든간의 최소거리를 의미합니다. 벡터를 L2 norm을 진행하면 거리가 1로 정규화되고, 각만 달라지게됩니다. 이 벡터들간의 거리가 좁을 때, 점점더 좁아지면 -log값이 매우커지니 서로 띄어놓게하기위해(uniform span) 이 정규화를 진행합니다. 

6. Adapting the resolution: 작은 이미지를 패치로 나누면 더 작게되어 pixel level task(semgnetaion, detection)등에서 성능열하가 발생합니다. 그렇기에 학습후반에만 일부 큰 이미지518x518로 학습합니다.

DINO v2의 PCA시각화

비지도학습을 이용해서, 임베딩값이 어떻게 픽셀상에 표현되는지 알기위해, 각 이미지를 RGB로 3개의 주성분을 가진 경우로 표현합니다. 임베딩을 주성분(각 성분은 RGB로 취급)하여 시각화합니다. 이 단계는 4단계로 구분됩니다.

• 첫 번째 PCA 수행:
  • 모델을 사용하여 이미지 패치(조각)들을 추출하고 이를 통해 특성(features)을 얻습니다.
  • 이 특성들에 대해 첫 번째 PCA를 수행하여 주성분들을 구합니다.
  • 첫 번째 주성분(첫 번째 컴포넌트)을 기준으로, 값을 임계값(threshold) 이상인 패치들만 남깁니다. 이렇게 하면 이미지의 주요 객체와 배경이 분리됩니다. 이 과정에서 첫 번째 주성분 방향이 가장 큰 분산을 설명하기 때문에, 주성분 값을 기준으로 분리하면 주로 배경과 전경이 나눠집니다.
• 배경과 전경 분리:
  • 첫 번째 주성분 값을 기준으로 배경과 전경이 분리된 후, 배경 패치를 제외한 나머지 패치들만 남깁니다. 즉, 전경 패치들만 남기게 됩니다.
• 두 번째 PCA 수행:
  • 전경 패치들에 대해 두 번째 PCA를 수행합니다. 여기서 주성분 3개를 추출합니다.
  • 이 주성분들은 이미지의 주요 객체의 다양한 부분들을 설명하는 데 사용됩니다.
• 시각화:
  • 두 번째 PCA로 얻은 3개의 주성분을 각각 R, G, B 채널에 할당하여 시각화합니다. 이렇게 하면, 이미지의 주요 객체의 다른 부분들이 서로 다른 색상으로 표현됩니다.

아래의 이미지를 샘플로 진행해보겠습니다.

from torchvision.io.image import read_image
from torchvision.transforms import ToPILImage
from matplotlib import pyplot as plt
image = read_image("....png")

plt.imshow(ToPILImage()(image))

dino 모델의 입력스케일이 [0, 1]의 float이기 때문에 /255로 나누고 `forward_features`메서드로 포워딩해줍니다.  features은 dictionary 형태로 나오고, cls token을 포함한 여러가지 output이 나옵니다. 이 중에서, 패치단위로 나눈 토큰(token)의 벡터를 구하기 위해서 "x_norm_patchtoken"의 값을 가져옵니다.

embedding_patch(E_patch)의 shape은 (1, 256, 178)인데, 1은 하나의 이미지, 256은 token의 수(패치수), 768은 임베딩 차원을 의미합니다.

import torch
from einops import rearrange
from torchvision.transforms import Normalize
from torchvision.transforms.functional import resize

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)

norm = Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD)

I_norm = norm(image / 255)
dinov2 = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
features = dinov2.forward_features(I_norm.unsqueeze(0))
E_patch = features["x_norm_patchtokens"]
print(E_patch.shape) # (1, 256, 768)


<class 'dict'> dict_keys(['x_norm_clstoken', 'x_norm_regtokens', 'x_norm_patchtokens', 'x_prenorm', 'masks'])
torch.Size([1, 256, 768])

PCA을 돌리기위해서 3차원의 1의 배치를 squeeze하여 2차원으로 변경해줍니다. (1, 256, 768) -> (256,768). 그리고, PCA로 주성분을 구하고, 그 주성분중에 가장 큰 주성분 벡터(V[:, 0])을 뽑아서, 임베딩 차원 -> 주성분차원으로 정사영(projection)시켜줍니다.

E_patch_norm = E_patch.squeeze(dim=0)
print(E_patch_norm.shape)

_, _, V = torch.pca_lowrank(E_patch_norm)  # V은 주성분. 각 벡터는 직교(orthogonal)

E_pca_1 = torch.matmul(E_patch_norm, V[:, 0])

주성분으로 프로젝션된 각 패치들 중에는 배경도 있고, 전경도 있을것입니다. 이 주성분은 가장 분산이 큰 설명하는 벡터이기에 배경/전경을 설명할 수 있습니다. 따라서, 주성분 1번(인덱스0번)에 내적하여 projection시킵니다. 

프로젝션된 벡터를 정규화하여, threshold을 걸어 배경과 전경을 분리합니다. 각 256 패치들이 전경에 속하는지, 배경에속하는지에 대한 마스킹을 갖계됩니다(boolean). 여기까지가 첫 번째 PCA이후, threshold을 걸어 배경과 전경을 분리하는 것입니다..

def minmax_norm(x):
    """Min-max normalization"""
    return (x - x.min(0).values) / (x.max(0).values - x.min(0).values)
    
E_pca_1_norm = minmax_norm(E_pca_1)
print(E_pca_1_norm.shape)
M_fg = E_pca_1_norm.squeeze() > 0.5 # (256, )
M_bg = E_pca_1_norm.squeeze() <= 0.5 # (256, )

여기서, 전경의 픽셀만 다시 분리하여 PCA을 새롭게 진행합니다. 이 두 번째 PCA의 목적은 주성분을 3개를 뽑아서RGB에 대응시켜 모델이 어느 패치(포지션)을 보고 임베딩을 했는지, 의미있게 임베딩했는지 확인해보자는 것입니다. 마찬가지로, 전경만 뽑아 PCA을 합니다.

PCA의 결과는 (N, 3)입니다. N은 전경의 패치수, 3은 프로젝션된 값입니다.

# 전경 패치만 뽑아서, PCA을 돌립니다.
_, _, V = torch.pca_lowrank(E_patch_norm[M_fg])

# PCA에서 가장 분산이 큰 주성분 3개를 뽑아 각 벡터에 projection합니다.
E_pca_3_fg = torch.matmul(E_patch_norm[M_fg], V[:, :3])
E_pca_3_fg = minmax_norm(E_pca_3_fg)  # (N, 3)

여기서 (N, 3)을 다시 이미지로 그려야합니다. (256,)의 패치에 전경부분만 (N, 3)의 (, 3)의 값들을 넣어줘야합니다. 그리고 각 패치에 맞도록 다시 rearange해줍니다. 이 과정은 (256, 3)을 패치 형태에 맞는 (16, 16, 3)으로 바꾸는 과정입니다.

B, L, _ = E_patch.shape  
Z = B * L
print(Z)
I_draw = torch.zeros(Z,3)
I_draw[M_fg] = E_pca_3_fg
I_draw = rearrange(I_draw, "(B L) C -> B L C", B=B)
I_draw = rearrange(I_draw, "B (h w) C -> B h w C", h=224//14, w=224//14)
print(I_draw.shape)

R, G, B형태에 맞게 시각화하고 resize해줍니다. 아래의 우측그림과 같이 우 하단영역의 염증부분에 더 많이 강조가 됨을 알 수 있습니다. 반면 좌측상단은 stroma부분으로 이 부분을 주의깊게 설명하는건 아니라는 것으로 해석해볼 수 있습니다.

# Unpacking PCA images
image_1_pca = I_draw[0]

# To chanel first format torchvision format
image_1_pca = rearrange(image_1_pca, "H W C -> C H W")

H, W = 224, 224
# Resizing it to ease visualization 
image_1_pca = resize(image_1_pca, (H,W))

from matplotlib import pyplot as plt
fig, axes = plt.subplots(1, 2)
axes[0].imshow(ToPILImage()(I))
axes[1].imshow(ToPILImage()(image_1_pca.detach()))