SAM (Segment Anything Model)과 같이 특정 task에 맞춰진 모델에서 광범위한 task를 수행할 수 있는 foundation 모델 개발로 패러다임이 전환됐다. SAM은 이미지 분할 task에서 놀라운 결과를 보였지만, 의료 도메인은 자연 도메인과 격차가 있다보니 의료 도메인에 적용 시 아쉬운 성능을 보였다. 이러한 한계를 극복하기 위해 MedSAM, SAM-Med와 같이 SAM을 의료 영상으로 fine-tuning한 연구들이 나왔지만, 여전히 한계가 존재한다.

1. 위 모델들은 2D 영상 분할을 위해 디자인된 아키텍쳐다.
2. 3D 영상 분할을 위해 SAM을 확장했지만, 임상 워크플로우에서 요구되는 3D 영상 분할에는 여전히 부족하다.
3. 지금까지 나온 모델들에 대한 대규모 검증이 부족하다.

본 연구는 이러한 한계점을 극복하기 위해 다양한 영상 모달리티, 병변, 장기들의 정보를 담고 있는 데이터셋을 구축하고, 3D 영상 분할 foundation 모델인 MedSAM2를 제안한다.

2. Methods

MedSAM2는 SAM2 모델을 그대로 가지고 와서 의료 도메인에 fine-tuning한 형태라고 볼 수 있다. 모델의 전반적인 구조는 이전에 정리해둔 SAM2 내용을 참고하고 (https://velog.io/@jjun_8177/SAM2-Segment-Anything-in-Images-and-Videos-%EC%A7%84%ED%96%89-%EC%A4%91), SAM2에 어떻게 의료 영상을 입력하여 학습하는지에 대해 더 살펴보겠다.

2.1. Image pre-processing

영상 전처리는 nnU-Net, Segment anything in medical images 논문에서 소개한 내용을 따른다.

1. 이미지 강도 조정

• CT : Hounsfield unit ( brain$(80/40)$, abdomen$(400/40)$, bone$(1800/400)$, lung$(1500/-600)$, mediastinum$(400/40)$ )
• MRI / PET : cut-off with lower-bound and upper-bound in 0.5% and 99.5%

2. 이미지 normalization 및 resampling$(512,512)$
3. 'npy' 파일로 저장

• CT, MRI, PET 영상 모두 NifTi 파일로 처리 후 numpy array로 저장
• Data loading을 빠르게 할 수 있고, 서로 다른 데이터 포맷을 통일할 수 있다

전처리 코드는 다음과 같다. 영상 모달리티에 따라 이미지 강도 조정 및 min-max normalization을 적용하는 것을 볼 수 있다. 전처리된 영상, 정답 마스크, 영상의 기하학적 정보를 npz 형태로 최종 저장하여 데이터셋으로 구축된다.

2.2. Network architecture

네트워크 구조는 이전의 SAM2와 동일하다. CT, MRI, PET는 모두 3차원 영상인데, 이 중 axial view 기준 하나의 slice를 image encoder에 태워서 frame embedding을 얻고, 프롬프트인 bounding box를 prompt encoder에 태워 prompt embedding을 얻어 mask decoder에 이 두 embeddings 간의 attention 연산을 수행하여 예측 마스크를 얻게 된다. 현재 slice를 기준으로 이전 slice 및 이후 slice도 같은 방법을 수행하지만, 여기서 현재 slice를 통해 얻은 예측 마스크로 memory attention을 수행하여 conditioned frame embedding을 얻은 후 수행하게 된다. 첫 번째 figure는 MedSAM2에서, 두 번째 figure는 SAM2에서 가져온 figure이다. 네트워크 구조가 동일하다는 것을 보이기 위해 이렇게 준비해봤다.

Fine-tuning은 네트워크의 모든 부분을 수행하게 된다. 네트워크의 모든 부분이 아닌 일부분, 예를 들어 prompt encoder와 mask decoder만 fine-tuning을 하면 더 좋은 결과를 보일 수도 있지 않을까 의문이 들 수 있지만, Segment Anything in Microscopy라는 논문에서 네트워크의 모든 부분을 fine-tuning 했을 때 제일 좋은 성능을 보인다고 입증하여 MedSAM2도 위와 같이 진행한다.

Image encoder가 받는 입력 이미지의 크기가 이전의 SAM2와는 차이가 있다. SAM2는 $1024 \times 1024 \times 3$ 인 이미지를 입력받지만, CT, MRI, PET 영상의 각 slice는 $512 \times 512$여서, MedSAM2는 $512 \times 512 \times 3$인 이미지를 입력받도록 수정됐다. 의료 영상은 흑백 이미지이다 보니 1-channel이고, 자연 영상은 컬러 이미지이다 보니 3-channel로 구성된다. MedSAM2는 SAM2의 image encoder를 그대로 사용하므로 해당 영상을 그대로 image encoder에 태우면 사이즈가 맞지 않으므로 에러가 날 것이다. 그래서 저자는 흑백 이미지를 복사하여 나머지 두 개의 channel에 붙임으로서 3-channel로 구성하는 방식을 사용한다.

Memory attention 모듈은 SAM2와 마찬가지로 4개의 transformer 블록으로 구성되어 있고, bounding box를 프롬프트로 받으므로 SAM 및 SAM2에서와 마찬가지로 bbox의 좌상단, 우하단 좌표를 임베딩하여 mask decoder에 입력하게 된다.

최종적으로 이미지 embedding과 프롬프트 embedding을 mask decoder에 태워 예측 마스크를 획득한다.

2.3. Training protocol

모델은 자연 영상으로 pre-train된 SAM2.1-Tiny 모델을 가져와서 fine-tuning한다. Fine-tuning 된 모델은 이미 feature를 뽑는데 특화되어 있을 것이므로 이렇게 진행한다. Fine-tuning 전략은 image encoder는 learning rate를 낮춰서 feature를 뽑는 능력을 유지할 수 있도록 하고, 나머지 모듈은 learning rate를 높여서 의료 도메인에 특화될 수 있도록 한다. Batch size는 SAM2와 마찬가지로 GPU 당 8로 설정하여, SAM2에서는 동영상에서 8 프레임씩 뽑아 입력했다면 MedSAM2는 8개의 slice를 뽑아 입력하게 된다. Loss는 focal loss와 dice loss를 선형 조합하여 사용하고, 비율은 20:1로 설정하여 focal loss에 조금 더 가중치를 준다.

3. Results

성능 평가는 다양한 모달리티, 병변, 장기들이 포함된 40 개의 segmentation task에서 진행되었다.

성능 비교는 SAM2.1의 tiny / small / base / large 모델, EfficientMedSAM-1 모델과 비교했다. EfficientMedSAM-1은 MedSAM2와 마찬가지로 bounding box를 입력 프롬프트로 받고, bbox segmentation foundation 모델에서 SOTA를 기록한 연구이므로 비교군으로 골랐다고 한다. 3D 이미지 성능 평가 결과는 다음과 같다. 결과는 위의 표를 보면 SAM2.1의 모든 모델에서 DSC 값에 큰 차이가 없다는 것을 볼 수 있고, EfficientMedSAM-1은 특정 task에서는 SAM2.1보다 좋은 결과를 보인다. 전반적으로 MedSAM2가 다른 비교군들보다 좋은 결과를 보인다.

다음은 비디오 성능 평가 결과이다. 두 번째 figure에서 첫 번째 영상은 초음파, 두 번째 영상은 내시경 영상이다. 형태 변화가 까다로운 영상에서 SAM2.1 모델들은 점수가 크게 떨어지고 오차범위가 넓어지는 반면에 MedSAM2는 분할 난이도가 높은 비디오에서도 안정적인 정확도와 분포를 보인다.

- MedSAM2 enables efficient 3D lesion annotation for large 3D CT and MRI datasets

병변의 이질성, noise 및 artifact는 라벨링 작업 시간이 많이 걸리고 노동 집약적인 작업으로 만든다. MedSAM2는 human-in-the-loop 파이프라인을 구축하여 라벨링 작업 시간을 압도적으로 줄이고 대용량 의료 분할 데이터셋을 구축한다.

순서는 다음과 같다.

1. 어노테이터가 영상의 가운데 slice 라벨링하고 싶은 영역에 bounding box를 그린다.
2. SAM2.1 모델이 초기 마스크를 생성한다.
3. 어노테이터가 마스크를 수정한다.
4. 수정된 마스크를 기반으로 MedSAM2가 모든 slice에 예측 마스크를 생성한다.
5. 어노테이터가 최종적으로 모든 slice의 마스크를 수정하여 최종 정답 마스크를 생성한다.
6. 충분한 데이터가 생성되면 SAM2.1 모델 fine-tuning을 실시한다.

Human-in-the-loop 방식을 통해 각 round가 지날 때마다 라벨링 시간은 단축되고 라벨링하는 CT의 수는 증가하게 되면서 효율적으로 대용량의 데이터셋을 생성하게 된다.

4. Discussion

좋은 성능을 보였음에도 저자는 몇 개의 한계점을 이야기한다. 첫 번째로는 프롬프트 제한이다. 현재 MedSAM2는 bounding box를 입력 프롬프트로 받는데, 혈관과 같이 얇고 긴 장기의 경우 적합하지 않은 프롬프트다. 두 번째로는 고정된 메모리 디자인이다. SAM2와 마찬가지로 8개의 slice를 메모리 뱅크에 저장하는데, 내시경 영상 같은 경우에는 카메라가 빠르고 연속적이게 움직이므로 추적에 실패할 가능성이 올라간다.

Segment Anything Model (SAM)은 이미지에 대해 프롬프트를 기반으로 segmentation을 할 수 있는 foundation model이다.

하지만 이미지는 현실 세계의 정적인 snapshot일 뿐이다. AR/VR, 로보틱스, 자율주행, 비디오 편집, 등의 중요한 어플리케이션은 이미지 레벨을 넘어 비디오 레벨에 대한 분할을 요구한다. 이 문제를 해결하기 위해 이미지와 동영상 모두에서 영상 분할이 가능한 universal segmentation system이 필요하다.

SAM2는 이러한 이미지 분할 task를 비디오 도메인에서 가능하도록 한 모델을 제안한다. 모델의 입력으로는 비디오의 아무 프레임에 대해 프롬프트를 입력하고, 출력으로 하나의 이미지, 혹은 비디오의 모든 프레임에 대한 segmentation mask를 뱉는다.

이미지에서 동영상 level로 가면서 많은 도전과제가 생긴다. 우선, 동영상 내에서 분할해야 하는 객체의 모션이 계속 바뀌고, 형태가 바뀌고, 어떤 프레임에서는 안 보일 수도 있고, 조명의 변화로 인해 픽셀 값에도 큰 변화가 생길 수도 있다. 또한, 이미지가 한 개의 프레임으로 된 동영상이라면, 동영상은 이미지에 비해 처리해야 할 프레임의 갯수가 훨씬 많아지게 된다.

2.Task : promptable visual segmentation (PVS)

PVS task는 비디오의 어떤 프레임에 대해서도 모델에 프롬프트를 입력하여 영상을 분할 할 수 있도록 해준다.

특정 프레임에 프롬프트를 줬을 때, 모델은 바로 해당 프레임에 유효한 segmentation mask를 출력한다. 이 출력된 prompt를 비디오 전체에 프롬프트로 줘서 모든 프레임에 대해 segmentation mask를 출력하게 되고, 만약 어떤 프레임에서 분할해야 할 객체를 놓친다면 여기서 우리는 다시 프롬프트를 줘서 모든 프레임에 대해 객체를 분할 할 수 있게 된다.

3. Model

3.1. Model architecture overview

모델의 전반적인 그림이다. 기존의 SAM1 모델과의 차이가 있다면 memory 부분이 생긴 것이다. 각 파트 별로 자세히 들어가보자.

3.2. Image encoder

SAM2는 이미지 encoder로 MAE task를 통해 pre-train된 Hiera encoder를 사용한다. MAE는 Masked AutoEncoder로, 이미지 내 일부 패치를 masking한 후 encoder와 decoder를 통과시켜 이미지를 복원하는 task이다. Hiera는 ICML 2023에 소개된 논문으로, SAM1에서는 MAE pre-trained ViT 모델을 사용했다면, Hiera는 기존의 ViT보다 계층적인 구조로 encoder를 만들어서 고차원의 정보와 저차원의 정보를 모두 반영하여 좀 더 이미지의 feature를 더 잘 반영할 수 있도록 제안된 모델이다.

3.3. Memory attention

SAM1에서는 이미지 encoder를 통과해서 나온 이미지 embedding을 바로 mask decoder에 입력으로 줬다.

SAM2에서는 encoder를 통과한 이미지 embedding을 unconditioned embedding이라고 한다. 이미지 레벨에서는 해당 이미지에 대해서만 잘 분할하면 됐지만, 동영상에서는 모든 프레임들끼리 정보를 공유하는 것이 분할 작업을 더 잘 수행할 것이라 판단하여 이전 프레임 정보, 예측 결과 등의 정보를 이 embedding에 반영하여 conditined embedding을 생성한다. 그러기 위해 나중에 소개할 memory bank 정보와 이미지 embedding 간의 cross-attention 연산을 수행한다. 이 연산을 memory attention 블록에서 진행하고, memory attention을 통해 얻은 frame embedding을 mask decoder의 입력으로 넣어주게 된다.

3.4. Prompt encoder

Prompt encoder는 SAM1과 동일한 encoder를 사용하므로 이전에 정리한 SAM1을 참고하자! (https://velog.io/@jjun_8177/Segment-Anything-%EC%A7%84%ED%96%89-%EC%A4%91)

3.5. Mask decoder

Mask decoder도 SAM1을 많이 따르고 있다.

SAM1에서 Mask와 IoU score를 출력하는 과정을 아래 페이지에 상세하게 작성해뒀으니 참고하면 좋을 것 같다. (https://velog.io/@jjun_8177/Segment-Anything-%EC%A7%84%ED%96%89-%EC%A4%91)

해당 페이지에서는 SAM1에서 소개한 mask decoder와 다른 부분만 짚어보자.

• SAM1에서는 image to token attention 연산을 통해 image embeddings를 업데이트한 후 바로 transposed convolution 연산을 통해 해상도를 높이고 채널 수를 축소했다면, SAM2에서는 upscaling 과정에서 image encoder (unconditioned) 의 stage 1과 2, 즉 고해상도 정보를 담고 있는 feature map을 저해상도인 image embeddings에 더해준다.

• SAM1에서는 mask decoder의 output이 mask와 IoU score고, 후보 mask들 중 예측 IoU score가 가장 높은 mask를 최종 출력한다. 하지만 SAM2에서는 그 외에 obj ptr, occlusion score도 함께 출력한다.

• Obj ptr는 object pointer로, 마스크를 생성해내는 "출력 토큰", 즉 분할 대상 객체의 고차원 이미 정보를 담고 있는 벡터로, 나중에 소개할 memory bank에 활용된다.

• SAM1은 이미지 레벨을 다루고 있고, 이미지에서는 보이는 객체에 대해 단순히 분할을 할 수 있다면 동영상의 경우에는 분할 대상인 객체가 시야에서 벗어나거나 다른 객체에 가려지면서 특정 프레임에서는 사라질 수 있다. 이를 대비하여 SAM2에서는 occlusion score를 도입하여 해당 프레임에서 분할 대상인 객체가 있는지 없는지를 occlusion score를 통해 표현한다. 이는 scalar 값으로 객체가 보이면 1, 안 보이면 0으로 나타낸다.

3.6. Memory encoder

Memory는 mask decoder에서 최종 출력한 mask를 downsampling한 후, 해당 frame의 frame embedding 과 합쳐서 만들어진다. 여기서 frame embedding은 Hiera encoder를 통과하여 출력된 값이다.

3.7. Memory bank

Memory bank는 말 그대로 memory 은행으로, frame마다 생성된 memory를 모아둔 공간이다. Memory bank는 이 memory를 queue 형식으로 관리하고, 현재 시점을 기준으로 최근 $N$ 개 frame의 memory와, 해당 frame들 중 prompt가 있는 $M$ 개의 memory로 구성된 2 개의 queue로 관리한다. 또한 mask decoder에서 소개한 object pointer 정보도 함께 저장한 후 memory attention layer에서 모두 frame embeddings(conditioned)를 생성하는데 활용된다.

Object pointer가 실제로 모델의 성능을 올렸는지 궁금했는데 마침 ablation study에 정리가 되어있었다.

솔직히 그렇게 드라마틱한 차이가 있는가 생각이 들지만 논문에서는 SA-V val 데이터셋에서 좋은 성능을 보여 최종적으로 사용했다고 한다.

3.8. Training

Training은 pre-training, full-training, fine-tuning 3 개의 과정으로 진행된다.

3.8.1. Pre-training

Pre-training에서는 정적인 이미지만 있는 SA-1B 데이터셋을 활용하여 학습을 하고, 한 이미지 당 7번의 수정 클릭을 하여 8번의 iteration을 돌면서 학습하게 된다. 이미지만 학습하므로 memory 영역이 필요없을 것이고, loss function은 출력 mask와 정답 mask 간의 loss를 focal loss와 dice loss의 선형 조합한 loss, 예측 IoU score와 실제 IoU score 간의 $L_1$ loss를 합쳐서 모델을 학습하게 된다.

3.8.2. Full-training

Full-training은 이미지(SA-1B)와 비디오(SA-V) 데이터셋 모두를 활용한다. 여기서 비디오는 frame 수가 많으므로 8 개의 연속된 frame 선택 후 최대 2 개의 frame까지 수정 클릭이 가능하도록 설정하고, masklet의 수도 최대 3개로 제한한다. 동영상은 50%의 확률로 역재생하도록 하여 모델이 양방향 전파(bi-directional propagation)하여 일반화될 수 있도록 한다. Loss function은 출력 mask와 정답 mask 간의 loss를 focal loss와 dice loss의 선형 조합한 loss, 예측 IoU score와 실제 IoU score 간의 MAE loss, occlusion score는 cross-entropy loss를 합쳐서 모델을 학습한다.

3.8.3. Fine-tuning

Full-training에서 비디오에 대해 8 frame만 샘플링하여 모델을 학습했는데, 동영상은 일반적으로 훨씬 많은 수의 frame으로 구성되어 있을 것이다. 그래서 fine-tuning 단계에서는 보다 긴 16 frame을 샘플링하고, 비디오는 challenging한 비디오로 제한한다.

4. Data

SA-V 데이터셋은 총 5만 9백 개의 비디오, 64만 2천 6백 개의 masklet으로 구성되어있다.

데이터 획득 과정 모델과 인간이 상호작용하며 만들어지고, 총 3개의 단계로 구성되어 있다.

Phase 1 : SAM per frame

Annotator는 비디오 모든 frame에 대해 분할 대상인 객체를 SAM 모델을 이용하여 분할하고, 지우개나 브러쉬같은 편집 tool로 수정한다.

Phase 2 : SAM + SAM2 Mask

동영상의 첫 번째 frame을 SAM 모델을 이용하여 mask를 생성하고, Mask만 prompt로 받을 수 있도록 커스터마이징 된 SAM2 Mask 모델에 해당 mask를 넣어 동영상의 나머지 frame에 대해서도 segmentation을 진행한다. 다음 frame의 예측 결과를 SAM 모델을 이용해서 수정한 후 다시 SAM2 Mask 모델에 넣는 방식으로, 정확한 mask가 나올 때까지 반복한다.

Phase 3 : SAM2

Annotator는 예측 mask를 수정하기 위해 SAM2에게 수정클릭만을 prompt로 줄 수 있도록한다.

5. Experiment

5.1. Semi-supervised video object segmentation

해당 방식은 첫 번째 frame에 대해서만 입력 prompt를 주고, 나머지 frame에 대해 mask를 예측하는 방식이고, 비교 모델은 VOS 모델인 XMem++와 Cutie를 SAM과 합친 모델이다. 모든 prompt에 대해 SAM2가 가장 높은 정확도를 보인다.

5.2. Image segmentation

이미지 데이터로만 학습된 SAM1과 SAM2를 비교해봤을 때, SAM2 모델 자체에서 모델 성능이 향상된 것을 확인할 수 있고, 이미지와 동영상으로 구성된 데이터로 학습했을 때 이미지만으로 학습된 모델보다 더 좋은 성능을 보이는 것을 확인할 수 있다.

5.3. Comparison to state-of-the-art semi-supervised VOS

VOS를 수행하는 최신 method와 비교한 결과 SAM2가 가장 좋은 성능을 보이는 것을 확인할 수 있다.

컴퓨터 비전, 자연어 처리 분야 모두 다양한 downstream task를 zero / few-shot으로 해결하는데 기여하는 foundation model을 개발하기 위해 많은 연구를 진행하고 있다. 본 논문은 어떠한 이미지에 대해서도 depth를 예측할 수 있는 monocular depth estimation (MDE) foundation model을 제안한다. MDE는 로보틱스, 자율주행, VR, 등의 어플리케이션에서 필히 해결해야 할 문제이다. 이전에는 depth를 계산하기 위해 센서, stereo matching, SFM (Structure-From-Motion) 등의 기술을 사용했다. 하지만 저자는 처음으로 대용량의 레이블 없는 데이터에 초점을 둔다. 해결 방식은 이전에 센서를 이용하여 레이블을 얻는 방식이나 SFM 기법과 비교하여 큰 이점을 가지고 있다.

1. Monocular 이미지는 획득하기 쉽고, 어디에든 존재하고, 센서와 같은 비싼 장비를 요구하지 않는다.
2. Monocular 이미지는 방대한 범위의 score를 커버할 수 있다. 이는 모델의 일반화와 확장성에 큰 도움을 준다.
3. MDE 모델은 레이블된 데이터를 사용하여 pre-train한 후 레이블이 없는 데이터에 대해 예측을 진행하면 된다. Stereo matching과 같은 복잡한 계산이 필요없어지게 된다.

모델은 우선 레이블된 데이터를 제공하는 6개의 public dataset을 통해 초기 모델을 training 한다. 그 후 레이블이 없는 데이터에 대해 자동으로 annotation을 생성하고, 레이블된 이미지와 함께 student model을 training 하는데 사용한다. 하지만, 생성된 가짜 레이블 데이터와 레이블된 데이터를 바로 함께 student model에 사용했을 때, 레이블된 데이터만으로 학습했던 teacher model에서 성능이 크게 올라가지 않는다고 한다. 저자는 이 문제를 해결하기 위해 student 모델이 가짜 레이블로 학습할 때 더 최적화하기 어려운 타겟을 통해 학습하도록 하여, student model이 추가적인 시각적 지식을 얻고, 안 본 데이터에 대해서도 robust 해지도록 한다.

2. Method

Label이 있는 데이터셋은 $D^l$, 없는 데이터셋은 $D^u$라고 하자. 우섯 첫 번째로 teacher model $T$를 $D^l$을 이용하여 학습하고, 학습된 $T$를 통해 $D^u$의 가짜 레이블을 생성한다. 다음으로 student model $S$를 $D^l$과 가짜 레이블을 포함한 $D^u$를 통해 학습하게 된다.

2.1. Learning Labeled Images

$T$의 학습과정은 다음과 같다. 입력 이미지 각 픽셀에 대해 예측된 depth 값 $d^*$이 있을 대, 우리가 알지 못하는 scale과 shift를 무시하는 affine-invariant loss를 통해 모델의 최적화가 이루어진다. Labeled loss는 다음과 같이 정의된다.

$$ \mathcal{L}l=\frac{1}{HW}\sum{i=1}^{HW}\rho(d_i^*,d_i) $$

여기서 $\rho$가 affine-invariant loss로, 아래와 같다.

$$ \rho(d_i^,d_i) = |\hat{d}_i^-\hat{d}_i| $$

여기서 $\hat{d}_i^*, \hat{d}_i$는 depth 값의 scale, shift된 버전이다.

$$ \hat{d}i=\frac{d_i-t(d)}{s(d)} \ t(d) = median(d) \ s(d) = \frac{1}{HW}\sum{i=1}^{HW}|d_i-t(d)| $$

나아가, $T$를 더 강력하게 만들기 위해, 저자는 사전 학습된 DINO v2의 가중치로 모델의 인코더의 초깃값을 설정한다. 그리고 사전 학습된 segmentation 모델을 통해 하늘 영역을 인식하여 0으로 값을 설정해둔다.

2.2. Unleashing the Power of Unlabeled Images

이 부분이 본 논문의 핵심이다. 컴퓨터 비전 task에서 레이블이 있는 데이터셋은 얻기 힘들지만, 없는 데이터셋은 쉽게 얻을 수 있다. 또한 한 scene에 대해 여러 각도에서 촬영한 이미지들을 통해 stereo matching이나 SFM 알고리즘을 쓰는 것보다, MDE 모델을 통해 depth map을 쉽게 얻을 수 있다.

$$ \hat{D}^u = \big{(u_i, T(u_i))|u_i \in D^u\big}_{i=1}^N $$

$D^l, \hat{D}^u$의 조합을 이용해서 student model $S$를 학습한다.

하지만 이 과정을 거쳐도 모델 성능에 큰 향상이 없었다고 한다. 충분한 양의 레이블 셋이 있다보니, 레이블이 없는 셋에서 추가적인 지식을 얻는 데에는 한계가 있었다. 레이블이 없는 셋에서 추가적인 시각적 지식을 얻기 위해, 저자는 $S$에 최적화하기 어려운 이미지를 입력으로 주는 방법을 제안한다. 첫 번째 기법은 color jittering, Gaussian blur와 같은 color distortion 기법이고, 두 번째 기법은 CutMix라는 spatial distortion 기법이다. CutMix 방식은 랜덤한 두 개의 레이블이 없는 이미지 $u_a, u_b$를 공간적으로 interpolate하는 방식이다.

$$ u_{ab}=u_a \odot M + u_b \odot (1-M) $$

여기서 $M$은 임의의 직사각형 영역을 1로 설정한 mask이다.

Unlabeled loss는 $M, 1-M$ 영역에 대해 각각 labeled loss를 구했을 때처럼 affine-invariant loss를 계산한 후 weigted average를 통해 합쳐서 구해진다.

$$ \mathcal{L}u^M=\rho(S(u{ab})\odot M, T(u_a) \odot M) \ \mathcal{L}u^{1-M}=\rho(S(u{ab}\odot (1-M), T(u_b) \odot (1-M)) \ $$

$$ \mathcal{L}_u = \frac{\sum M}{HW}\mathcal{L}^M_u + \frac{\sum (1-M)}{HW}\mathcal{L}^{1-M}_u $$

2.3. Semantic-Assisted Perception

이전 연구에서 semantic segmentation task를 depth estimation task에 활용하면 성능이 향상한다는 것을 보였다. 저자는 high-level한 semantic 정보를 MDE 모델에 주입하여 더 좋은 성능을 기대한다. 첫 번째 시도는 RAM, Grounding DINO, HQ-SAM 모델의 조합으로 레이블 없는 셋에 semantic segmentation label을 할당하는 방법이다. 이를 MDE 모델과 함께 training 해서 depth와 segmentation 결과를 예측하게 했지만, 기존의 MDE 모델만 사용했을 때와 성능 면에서 큰 차이가 없었다고 한다. 이는 이미지를 이산적인 클래스 공간에 decoding 하는 것이 많은 semantic 정보를 손실하기 때문에 성능 향상을 기대할 수 없었다고 한다. 두 번째 시도는 image retrieval, semantic segmentation과 같은 semantic 관련 task에서 fine-tuning 없이 학습된 가중치만으로도 좋은 성능을 보인 DINO v2 모델을 사용하는 것이다. 저자는 해당 모델을 통해 얻은 semantic 정보를 MDE 모델에 주입하고, feature alignment loss를 도입하여, 고차원이면서 연속적인 feature space에서 최적화를 진행하여 모델의 성능을 향상시키고자 한다.

$$ \mathcal{L}{feat}=1-\frac{1}{HW}\sum{i=1}^{HW}\cos{f_i, f_i^{'}} $$

위 식에서 $f_i$는 $S$에서, $f_i^{'}$는 DINO v2에서 각각 나온 feature vector이고, 두 벡터 간의 코사인 유사도를 통해 loss 값을 구한다.

여기서 주의할 점은, DINO v2는 하나의 물체에 대해 각자 다른 부분에서 비슷한 벡터를 산출하지만 (예를 들어 자동차의 앞, 옆, 뒤 모두 결국엔 자동차이므로 같은 벡터를 산출), depth estimation에서는 같은 물체에 대해서도 각 픽셀마다 depth가 다양하다는 점이다. 이 문제를 해결하기 위해, 저자는 허용 가능한 margin 값 $\alpha$ 를 설정하여, $f_i$와 $f_i^{'}$의 코사인 유사도가 $\alpha$를 넘지 않는 픽셀에 대해서만 해당 loss term을 적용한다. 전체적인 loss는 $\mathcal{L}l, \mathcal{L}_u, \mathcal{L}{feat}$ 총 3 개의 loss term을 평균 조합하여 모델을 최적화한다.

Transformer 기반의 구조는 번역과 같은 시퀀스 모델에서 최고의 성능을 보여줬지만, image captiong과 같은 multi-modal task에서는 아직 충분한 연구가 수행되지 않았다. 본 논문은 image captioning task를 수행하기 위한 Meshed-memory Transformer 모델을 제안한다.

모델의 구조는 우선 encoder 부분에서 $i$ 개의 이미지 영역과 그들의 관계가 multi-level로 인코딩 되도록 하는데, 이런 구조 덕분에 low-level, high-level 모든 관계를 고려할 수 있게 된다. Decoder 또한 low / high-level 한 시각적 관계를 모두 탐험하는 multi-layer 구조로 되어 문장을 생성해낸다.

2. Method (Meshed-Memory Transformer)

모델은 encoder, decoder 모듈로 나눠서 볼 수 있고, 두 모듈 모두 여러 개의 attention layer가 쌓여있는 구조이다. Intra-modal, 단어와 이미지 간의 cross-modal 모두 recurrence가 아닌 scaled dot-product attention 연산을 수행한다.

$$ Attention(\bold{Q},\bold{K},\bold{V}) = softmax\bigg(\frac{\bold{Q}\bold{K}^T}{\sqrt{d}}\bigg)\bold{V} $$

위의 연산은 Transformer 구조가 처음 제안된 'Attention is all you need (NeurIPS 2017)' 논문에서 제안된 수식이고, $\bold{Q}, \bold{K}, \bold{V}$는 각각 query, key, value로 입력 시퀀스를 임베딩 벡터로 변환하여 만들어진 벡터들이다.

2.1. Memory-augmented encoder

이미지 내에서 뽑은 이미지 영역 집합 $X$가 있을 때, self-attention을 통해 encoding 정보 $S(X)$를 얻을 수 있다.

$$ S(X) = Attention(W_qX, W_kX, W_vX) $$

$W_q, W_k, W_v$는 학습 가능한 가중행렬이고, 이 가중행렬을 통해 input feature인 $X$를 투영시켜 $\bold{Q}, \bold{K}, \bold{V}$를 얻는다. Self-attention을 통해 input feature 간의 관계를 encoding하게 된다.

하지만 self-attention만으로 '사전 지식'까지는 encoding할 수 없다. 예를 들어, 사람과 야구공이 input feature로 뽑혔을 때, 사람은 player, game 등으로 추리할 수 있지만, 컴퓨터는 그런 사전 지식이 없다.

2.1.1. Memory-augmented attention

이러한 self-attention의 한계를 극복하기 위해 본 논문은 memory-augmented attention을 제안한다. Memory-augment, 말 그대로 메모리를 증강하는 방법으로, key와 value에 사전 정보를 encoding 할 수 있는 slot을 추가하여 확장하는 것이다.

$$ \mathcal{M}_{mem}(\bold{X}) = Attention(W_q\bold{X}, \bold{K}, \bold{V}) \ \bold{K} = [W_k\bold{X}, \bold{M}_k] \ \bold{V} = [W_v\bold{X}, \bold{M}_v] $$

위 수식에서 $\bold{M}_k, \bold{M}_v$는 학습 가능한 가중 행렬이고, key와 value는 이 행렬들을 각각 concatenate하여 확장된다. 이렇게 구성된 query, key, value로 self-attention 수행 시, $\bold{X}$에 아직 임베딩 되지 않는 사전 지식을 가져올 수 있다. $W_q, W_k, W_v, \bold{M}_k, \bold{M}_v$는 각 encoder head 마다 다르게 구성된다.

2.1.2. Encoding layer

Memory-augmented attention 연산을 통해 나온 출력값은 두 개의 Feed forward layer를 통과한다. 또한, 기존의 transformer 구조와 같이 self-attention layer, feed forward layer 모두 residual connection과 layer normalization을 적용한다.

2.2. Meshed Decoder

디코더는 encoding된 이미지 영역과 이전 시점에 생성된 단어 모두에 conditioning되어 다음 단어를 예측한다. 인코더를 통해 이미지의 multi-level representation을 얻었는데, 본 논문은 meshed attention 연산을 제안하여 기존의 Transformer에서 제안된 cross-attention 연산과 다르게 단어를 생성하기 위해 모든 encoding layer를 고려하게 된다.

2.2.1. Meshed cross-attention

$i$ 개의 encoder layer에서 나온 출력값 $\bold{X}^i$에 대해, 디코더의 입력 시퀀스 벡터 $\bold{Y}$와 이 모든 $\bold{X}$들 간의 cross-attention 연산을 수행한다. 그 후, 각 cross-attention 결과와 같은 크기의 가중 행렬 $\alpha_i$를 통해 cross-attention 결과를 조정한 후 더하여 출력값을 얻는다.

$$ \mathcal{M}{mesh}(\mathcal{X}, \bold{Y}) = \sum{i=1}^{N}{\alpha_i\odot C(\bold{X}^i,\bold{Y})} \ C(\bold{X}^i,\bold{Y})=Attention(W_q\bold{Y},W_k\bold{X}^i,W_v\bold{X}^i) \ $$

2.2.2. Architecture of decoding layer

이전에 Transformer에서 소개된 내용과 마찬가지로, 본 논문 또한 masked self-attention을 통해 다음 단어를 이전 시점에서 예측된 단어들에만 의존하여 예측할 수 있도록 한다. Masked self-attention을 통해 얻은 query와 encoder layer에서 얻은 $N$ 개의 key, value 간의 cross-attention을 수행한다. Decoder layer 또한 encoder layer와 마찬가지로 feed forward layer, 각 component마다 residual connection, layer normalization을 적용한다.

2.3. Training Details

본 논문은 우선 단어 수준의 cross-entropy loss로 모델을 pre-training한 후, 강화학습을 통해 시퀀스 생성을 fine-tuning한다. 강화학습 적용 시, beam search를 이용해 각 시퀀스마다 확률이 가장 높은 $k$ 개의 후보 단어들을 샘플링 하여, 이들을 통해 보상 평균을 계산해서 안정적이게 업데이트 하도록 한다. 보상으로는 CIDEr-D score를 사용한다. 최종적으로 한 샘플에 대한 gradient는 다음과 같이 계산된다.

$$ \triangledown_\theta L(\theta) = -\frac{1}{K}\sum_{i=1}^{k}(r(\bold{w}^i)-b)\triangledown_\theta \log{p(\bold{w}^i))} $$

본 논문은 2D CNN (Convolutional Neural Network), 3D CNN, Transformer를 조합하여, 3D MRI 이미지를 통해 알츠하이머 증상 (Alzheimer Disease, AD)를 분류해내는 모델을 제안한다. 제안된 모델은 이전 3D image 분류기들과 비교한 결과 가장 좋은 성능을 보이면서, 3D로 구성된 medical image를 위한 CNN, Transformer의 조합이 타당하다는 것을 증명한다.

1. Introduction

CNN은 classification, detection, segmentation, 등 다양한 컴퓨터 비전 task에서 중요한 역할을 수행했다. 나아가 CNN기반 모델은 X-ray, MRI, CT, 초음파, 등의 2D부터 3D까지 다양한 차원으로 구성된 medical image에 대해서도 적용되어 왔다. 이 수많은 modality들 중에서, 특히 3D medical image를 분석하기 위해 2D 및 3D CNN이 사용되었다. 하지만 각 네트워크의 장단점은 명확히 구분된다. 2D CNN은 대규모의 2D 자연 이미지로 학습된 사전 학습 모델을 사용한다는 이점이 있지만, 3D 이미지의 문맥을 이해하기는 힘들다. 3D CNN은 이와는 다르게 3D 이미지의 문맥을 잘 이해할 수 있다. 하지만, 많은 수의 파라미터와 방대한 양의 계산을 요구하므로 네트워크를 깊게 쌓는 것은 불가능하다. 최근에는 Transformer 네트워크가 제안되면서 자연어처리와 컴퓨터비전 분야 모두에 널리 사용되고 있다. CNN은 수용 영역 (receptive field)의 한계로 인해 데이터의 global한 부분만 보는데, Transformer는 넓은 수용 영역을 가지고 있어서 네트워크가 깊어질 수록 수용 영역 또한 넓어지게 된다. 하지만 학습 데이터셋이 충분하지 않다면 CNN보다 정확도가 떨어진다는 단점이 있다. 3D medical image의 경우 다른 domain의 데이터셋에 비해 많이 부족하므로 이러한 단점을 고려하지 않을 수가 없다. 이러한 두 네트워크의 장점을 모두 가져갈 수 있도록 조합된 hybrid network가 많이 연구되었지만, 모두 2D 이미지를 위한 2D CNN과 Transformer를 조합하였다. 본 논문은 3D 이미지를 위해 2D CNN, 3D CNN, Transformer 모두를 조합한 네트워크를 소개한다. 왜 이 3개의 네트워크가 AD를 분류하는데 필요할까?

• AD 환자에게 생기는 대뇌피질 (cerebral cortex) 위축은 뇌에 전반적으로 분포하므로, 수용 영역이 넓은 Transformer가 필요하다.
• 뇌실 (ventricle) 영역과 해마 (hippocampus) 영역의 수축은 뇌의 local한 영역에서 발생하므로 CNN이 필요하다.
• MRI 이미지는 3차원으로 구성되어 있고, slice 단위가 아닌 volume 전체를 사용하기 위해 3D CNN이 필요하다.

2. Methods

모델의 전반적인 구조는 아래와 같다.

2.1. 3D Convolutional Neural Network Block

우선 입력으로 이미지의 depth $L$, width $W$, height $H$가 동일한 MRI 이미지를 받는다. 그 후 feature를 얻기 위해 2개의 $5\times5\times5$ 3D CNN layer와 Batch Normalization, ReLU function으로 구성된 layer를 통과한다.

2.2. Extraction of Multi-plane, Multi-slice Images

보통 CT나 MRI 촬영을 진행할 때, 3개의 view에 대해 필요로 하는 view로 촬영을 진행한다. View의 종류로는 axial, sagittal, coronal view가 있다.

3D CNN block을 통해 얻은 3D feature map에 대해 3개의 view에 대한 feature extractor를 통해 feature를 추출한다.

2.3. 2D Convolutional Neural Network Block

추출된 feature들은 weight를 공유하는 2D CNN layer를 통과한 후, 두 개의 MLP layer와 ReLU function으로 구성된 layer를 통과해서 $d$ 차원의 벡터로 투영된다. 본 논문은 $d$를 256으로 설정한다. 또한 2D CNN block으로 ImageNet으로 pre-train이 되어있는 ResNet50을 사용한다고 한다.

2.4. Position and Plane Embedding Block

Non-linear projection layer를 통해 position token과 plane embedding token을 각 image token에 더해준다. 또한 classfication token인 $z_{cls}$와, BERT의 sep token과 비슷한 $z_{sep}$도 추가한다.

2.4. Transformer Block

위의 embedding 과정을 거친 image token들은 Multi-head Self Attention (MSA), Layer Normalization, MLP block으로 구성된 $k$ 개의 Transformer layer를 통과한다. 본 논문은 $k$를 8로 설정한다. 각 layer를 통과할 때마다 image token의 차원인 $d$는 유지된다. 마지막으로, MLP head인 $z_{cls}^k$를 기반으로 알츠하이머 증상 여부를 분류한다.

3. Experiments

3.1. Implementation details

Preprocessing

다양한 기관에서 받은 MRI 이미지들을 정규화 및 표준화하기 위해, 동일한 전처리 작업을 모든 데이터셋에 적용한다.

• N4 algorithm : Image intensities의 불균일성 제거
• Skull stripping algorithm : Non-brain tissues 제거
• Resize (same voxel spacing & matrix size)
• Normalize (Image intensities to zero-mean unit-variance)

3.2. Comparison study results

다른 연구들과 비교한 결과, 본 논문이 제안한 방법이 가장 높은 정확도를 보인다. 또한, CNN과 Transformer를 조합한 hybrid network가 3D CNN network보다 더 좋은 성능을 보인다. Transformer 만을 사용한 network는 적은 수의 데이터에 대해 낮은 정확도를 보인다는 것 또한 알 수 있다.

3.3. Ablation study results

제안된 모델에서 각각 3D CNN, 2D CNN, Transformer block을 제거하고 실험한 결과, 모든 block을 사용한 network가 제일 높은 정확도를 보인다.

3D CNN을 통해 얻은 feature map에 대해 coronal / sagittal / axial view 에 대한 feature를 추출하는데, 모든 view에 대한 feature를 사용하는 것이 가장 높은 정확도를 보인다.

데이터 수집은 인터넷 상에서 이미지 및 이에 대한 캡션을 모두 긁어서 얻게 된다.

그 후 text encoder를 통해 벡터 형태로 나타낸다면, 이 벡터가 이미지 입장에서 label이 될 수 있을 것이다. 즉, image encoder 및 text encoder를 통해 얻은 image vector 및 text vector 사이의 관계를 학습하는 것이 핵심이 될 것이다. 이 두 vector가 비슷한 위치로 mapping 된다면 해당 text가 해당 image의 label이 될 것이다.

학습 방식은 다음과 같다.

우선 image encoder는 ViT (Vision Transformer), text encoder는 Transformer를 사용해서 각 image vector, text vector를 얻어낸다. 그 후 두 vector의 cosine similarity를 사용해서 많은 vector들 중 한 객체를 표현하는 image와 text가 비슷한 위치에 mapping되도록 학습한다.

이후 학습 때 사용하지 않은 이미지를 학습된 image encoder에 통과시켜 벡터를 추출한 후, N 개의 text를 각각 text encoder에 통과시킨 벡터들과 cosine similarity를 계산하여 inference를 진행한다.

CLIP의 zero-shot transfer의 성능을 가장 직관적으로 나타낸 figure를 보면,

ImageNet에 있는 바나나 데이터로 학습했을 때 결과가 비슷하지만, 학습된 모델로 다른 데이터셋의 바나나에 적용하면 CLIP의 효력을 확실하게 볼 수 있다.

모델의 입력부터 출력까지 위 모델의 그림에 따라 정리해보겠다.

1. Input embedding

우선 입력 시퀀스는 'input embedding'을 통과해서 각 시퀀스마다 크기가 512인 벡터로 변환한다.

2. Positional encoding

기존의 RNN은 모델 자체적으로 입력 시퀀스의 순서를 고려했지만, Transformer는 RNN이 제거되었으므로 위치 정보가 고려되지 않는다. 그래서 아래 그림처럼 입력 시퀀스의 순서를 바꿔도 attention의 결과가 바뀌지 않는다.

그래서 Transformer에서는 입력 시퀀스의 위치 정보를 제공하기 위해 positional encoding(PE)을 사용한다.

하지만 이 수식을 보면 직관적으로 와닿지 않는다. 그리고 위치 정보를 입력하기 위해 positional encoding을 하는건 알겠지만, 왜 이 방법을 쓰는지도 와닿지 않을 수 있다. 이 부분을 한 번 짚고 넘어가보자.

Positional encoding이 도입되기 전 다양한 시도가 있었다.

2.1. Positional encoding using index

단순히 입력 벡터의 인덱스를 PE로 사용하면, 입력의 길이가 길어지면서 값이 너무 커지는 문제가 발생하며, 네트워크를 학습하는데 지장이 생긴다.

2.2. Positional encoding using normalized index

2.1의 문제를 해결하기 위해 정규화된 인덱스를 PE로 사용해보았지만, 이번에는 입력의 길이가 길어지면서 주변 단어들 간의 상대적인 위치 차이가 애매해지는 문제가 생긴다.

2.3. Positional encoding using binary representation

입력 벡터의 인덱스를 2진수로 나타내어 임베딩 값에 합쳐봤지만, 이 방식은 낭비되는 bit가 많고, dot-product를 수행하는 attention에 적합하지 않다.

2.4. Sinsusoidal positional encoding

Sinusoidal PE는 임베딩 벡터와 같은 차원을 갖도록 위치값을 임베딩하여 입력값과 합친다.

또한, 입력 시퀀스의 길이가 길어지면 문제가 있는 2.1, 2.2의 문제를 해결한다. 이는 수학적으로 증명이 가능하다. $\Delta t$ 만큼 차이나는 두 PE vector가 있을 때, 두 벡터의 dot-product를 하면 다음과 같다. 편의상 $\frac{1}{1000^{2*k-1/512}}$ 를 $w_k$로 두겠다. $$ PE_i^TPE_{i+\Delta t}=\begin{bmatrix}\sin{w_0i} \ \cos{w_0i} \ \vdots \ \sin{w_ki} \ \cos{w_ki}\end{bmatrix}^T \begin{bmatrix}\sin{w_0(i+\Delta t)} \ \cos{w_0(i+\Delta t)} \ \vdots \ \sin{w_k(i+\Delta t)} \ \cos{w_k(i+\Delta t)}\end{bmatrix} $$ $$ = \sin{w_oi} \sin{w_0(i+\Delta t)} + \cos{w_0i} \cos{w_0(i+\Delta t)} + \cdots $$ $$ = \sin{w_0i}\bigg(\sin{w_0i}\cos{w_0\Delta t} + \cos{w_0i}\sin{w_0\Delta t}\bigg) + \cos{w_0i}\bigg(\cos{w_0i}\cos{w_0\Delta t} - \sin{w_0i}\sin{w_0\Delta t}\bigg) + \cdots $$ $$ = \sin^2{w_0i}\cos{w_0\Delta t} + \cos^2{w_0i}\cos{w_0\Delta t} + \cdots $$ $$ = \cos{w_0\Delta t}\bigg(\sin^2{w_0i}+\cos^2{w_0i}\bigg) + \cdots $$ $$ = \sum_{i=1}^{k}\cos{w_i\Delta t} $$

이는 결국 positional encoding을 적용하면 어떠한 위치든 한 시점에 대해 linear function으로 설명할 수 있게 된다. 이렇게 단어의 갯수만큼 사이즈 512인 PE 벡터를 만들어서 input embedding과 더해서 모델에 입력하게 된다.

3. Encoder

인코더는 positional encoding과 input embedding layer가 더해져서 입력된 후, 각 벡터 간의 self-attention을 수행하는 것이 핵심이다. 인코더의 구조는 아래 그림과 같다. 이 부분도 그림의 흐름에 따라 적어보겠다.

3.1. Multi-Head Attention

3.1.1. Self-Attention

이전 글에서 작성한 attention은 디코더의 target 시점에서 인코더의 각 시점과의 유사도를 계산하여, 입력 시퀀스의 어느 부분에 초점을 맞춰 출력 시퀀스에 반영할지 학습했다.

Self-attention은 인코더 내 각 입력 시퀀스 간의 관계를 비교하는 것이다. 각 입력 시퀀스마다 key$(K)$, value$(V)$가 있고, 현재 집중할 부분을 query$(Q)$라 했을 때, attention을 구하는 수식은 다음과 같다. $$ Attention(Q,K,V) = softmax(QK^T)V $$ $Q$, $K$, $V$ 모두 $(1\times512)$ 크기로 임베딩된 벡터이므로, $Attention(Q,K,V)$의 크기는 $(1\times512)$ 그대로 유지하게 된다.

3.1.2. Scaled dot-product attention

Dot-product 시 입력값에 따라 값이 매우 크거나 작아질 수 있다. 이 값을 softmax를 그대로 취한 후 backpropagation을 하면 gradient가 모두 0으로 될 것이다. 이를 방지하기 위해 softmax를 취하기 전에 일정한 scalar 값으로 나누어 normalize를 한다. 여기서 scalar 값은 Q,K의 차원 수 $d_k$의 제곱근을 보통 사용한다. $$ Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V $$

3.1.3. Multi-Head Attention

Multi-head self attention은 입력 벡터를 여러 개의 작은 벡터로 나누어, 서로 self-attention을 수행하며 입력 시퀀스 간의 관계와 패턴을 학습할 수 있도록 한다.

3.2. Residual Connection

Kaiming He의 'Deep Residual Learning for Image Recognition'에서 소개된 residual connection을 모델의 구조에 도입했다.

3.3. Add & Norm

Attention layer를 통과한 벡터 $N$개와 입력 임베딩 벡터 $N$개를 residual connection을 통해 더한 후, 안정적인 학습을 위해 $N(0,1)$ 분포를 갖도록 normalize를 한다.

3.4. Feed Forward

Attention layer의 출력에 non-linearity가 적용되지 않았으므로, 해당 layer를 추가하여 비선형성을 추가한다.

해당 layer를 통해 $(1\times512)$ 크기의 벡터를 $(1\times2048)$ 크기로 늘렸다가 다시 $(1\times512)$ 크기의 벡터로 돌려놓는다.

4. Decoder

Output embedding은 다음 단어 예측을 위해 정답 문장을 한 칸씩 오른쪽으로 shift한 문장을 입력으로 받는다. 예를 들어, 인코더의 입력 시퀀스가 'I love you' 였다면, 디코더의 입력 시퀀스는 '$<$start$>$ I love'가 된다. 그 다음은 인코더에서와 같이 각 시퀀스를 $(1\times512)$ 크기의 벡터로 임베딩하고, positional encoding과 함께 더해서 디코더에 입력하게 된다. 디코더의 구조는 다음과 같다. 디코더의 모델 구조에서 masked multi-head attention을 제외한 나머지는 인코더 부분을 설명할 때 한 내용들이므로, 여기서는 masked multi-head attention만 짚어보겠다.

4.1. Masked Multi-Head Attention

일반적인 self-attention 방식은 미래의 출력들까지 고려하므로, $n$번 째 입력에 대한 출력을 구하고 싶다면 $n-1$까지의 입력만을 고려해야 한다. 이를 위해 attention 계산 과정에서 n-1 시점 이후의 입력에 대해서는 masking을 적용한다.

4.2. Masked Multi-Head Cross Attention

기존의 attention이라 생각하면 된다. 인코더의 $n-1$ 시점까지의 입력 시퀀스와 디코더의 입력 시퀀스 간의 연산을 통해 인코더의 어떤 시점에 집중하여 출력값을 결정할 지 알아낸다.

디코더를 통과해서 나온 $N \times 512$의 벡터 ($N$은 vocab size)를 linear layer에 통과하여 사이즈 $N$인 벡터를 얻고, softmax를 통과하여 각 단어에 대한 확률 예측값을 구하게 된다. 이 예측 확률값들로 구성된 벡터와 one-hot encoding 된 실제 정답을 비교하여 모델 학습을 진행하게 된다.

Seq2Seq (Sequence-to-Sequence)는 인코더와 디코더라는 2 개의 모듈로 구성되어 있다. 인코더는 입력 시퀀스를 압축해서 하나의 벡터로 만들고, 이를 context vector라고 부른다. 디코더는 이 context vector를 초기 hidden state로 설정해서 값을 출력하게 된다.

인코더 내부는 RNN계열 (vanilla, LSTM, GRU)로 구성되어 있다. 차이점은 RNN은 각 입력 시퀀스마다 hidden state를 통해 값을 출력했다면, Seq2Seq의 인코더는 모든 입력 시퀀스를 통해 하나의 hidden state를 산출해낸다. 여기서 이 하나의 hidden state가 context vector인 것이다. 기존의 RNN은 시퀀스가 길어질수록 앞 쪽 시점의 정보를 잘 반영하지 못했다. 하지만 Seq2Seq는 입력 시퀀스의 모든 정보를 압축해서 반영했다고 할 수 있다. 디코더는 이 context vector를 첫 번째 hidden state로 활용하여 출력값을 예측한다. RNN은 출력값과 hidden state를 출력하는데, 이 정보를 이용해서 다음 시퀀스를 예측하는데 사용한다.

2. Seq2Seq의 한계

Seq2Seq는 입력 시퀀스의 모든 정보를 고정된 크기의 context vector에 encoding을 하려다 보니 information bottleneck 문제가 발생한다.

3. Seq2Seq with Attention

Attention은 이러한 bottlenect 문제를 해결하기 위해 제안됐다. 핵심 아이디어는 각 시점마다 디코더에서 입력 시퀀스의 특정 부분에 초점을 맞출 수 있도록 인코더로 직접 연결하는 것이다. 이를 위해 디코더의 한 시점에서의 벡터 $\bold{s_t}$와 인코더의 모든 시점에서의 벡터 $[\bold{h_1, h_2}, \dots, \bold{h_N}]$ 간의 dot-product를 취해 attention score $e^t$를 계산하고, softmax를 취해 attention distribution $\alpha^t$를 구해서 어떤 시점에 집중할 지 알아낸다. Attention distribution을 기반으로 인코더의 모든 시점에서의 벡터 $[\bold{h_1, h_2}, \dots, \bold{h_N}]$에 weighted sum을 적용하여 최종 attention output $\bold{a^t}$를 얻는다. Attention output은 attention 값이 높은 시퀀스의 정보를 주로 담는다. Attention output과 디코더의 hidden state를 concatenate 하여 일반적인 Seq2Seq처럼 처리하여 출력 시퀀스를 계산한다.

이 과정에 대한 그림과 수식은 아래와 같다. $$ e^t = [\bold{s_t}^T\bold{h_1}, \bold{s_t}^T\bold{h_2},\dots,\bold{s_t}^T\bold{h_N}] $$ $$ \alpha^t=Softmax(e^t) $$ $$ \bold{a^t}=\sum_{i=1}^{N}{\alpha^t_ih_i} $$

3.1. Why dot-product?

두 벡터 $\bold{v}, \bold{w}$가 있다 해보자. 내적은 수식으로 다음과 같이 표현할 수 있다. $$ \bold{v}\cdot\bold{w}=||\bold{v}||||\bold{w}||\cos{\theta} $$ 여기서 $\theta$는 두 벡터 사이의 사잇각이다. 이 식을 $\cos{\theta}$에 대해 정리하면 다음과 같다. $$ \cos{\theta}=\frac{\bold{v}\cdot\bold{w}}{||\bold{v}||||\bold{w}||} $$ 여기서 $\cos$에 대한 특성을 생각해보면, 사잇각이 작아질수록 $\cos{\theta}$의 값은 1에 가까워지고, 사잇각이 커질수록 $\cos{\theta}$의 값은 -1에 가까워진다. 즉, 두 벡터의 내적값을 토대로 두 벡터의 유사도를 판단할 수 있게 된다.

오늘 다룰 글은 RNN이다.

1. RNN의 등장배경

시계열, 동영상, 텍스트, 등 다양한 시퀀스 데이터가 존재한다.

그리고 이러한 입력 시퀀스가 있을 때, 이전 시점을 통해 다음 시점을 예측할 수 있다. 이론 상 이를 무한히 반복해서 원하는 시점까지의 데이터를 예측할 수 있을 것이다.

하지만 이렇게 단순히 한 시점에만 의존하여 출력을 계산하면, 이전 시점들과는 독립적이므로 부정확한 결과가 나올 확률이 크다. 그래서 등장한 것이 RNN이다. RNN은 시퀀스를 처리할 때마다 바뀌는 hidden state가 있는데, 이를 통해 이전 시점을 반영하여 다음 시점을 예측한다.

2. RNN의 연산과정

연산과정을 좀 더 자세히 나타내면 다음 그림과 같다.

$$ O_t = h_tW_{ho} \ h_t = \tanh(h_{t-1}W_{hh} + x_tW_{xh}) $$

여기서 $x_t$, $o_t$는 입력과 출력이고, $W_{xh}$, $W_{hh}$, $W_{ho}$ 는 0~1 사이의 랜덤 값으로 초기화된 weight matrix이다. Weight matrix 각각의 shape은 (input 길이 $\times$ hidden 길이), (hidden 길이 $\times$ hidden 길이), (hidden 길이 $\times$ output 길이) 이다.

예시로 'h', 'e', 'l', 'l'이 입력으로 들어올 때 'o'를 예측하는 방법을 수식으로 풀어보자.

1. One-hot encoding 각 단어인 'h', 'e', 'l', 'o'를 벡터로 encoding한다.

   이 벡터들이 입력값인 $x_t$가 된다.

2. Hidden state 계산 입력값 $x_t$와 이전 시점의 hidden state $h_{t-1}$, 그리고 다음 hidden state로 보내는데 사용되는 두 개의 weight matrix $W_{xh}$, $W_{hh}$ 를 이용해서 다음 hidden state를 계산할 수 있다. $$ h_t = \tanh{(h_{t-1}W_{hh} + x_tW_{xh})} $$ 이 식을 따르면 $h_1 = \tanh{(h_{0}W_{hh} + x_1W_{xh})}$ 가 된다. $x_1$은 $(1,0,0,0)$이고, 초기 은닉 상태값은 영행렬로 설정한다. $h$의 shape은 하이퍼파라미터이고, $h_0$는 $(0,0,0)$이라 하겠다. 그럼 $W_{hh}$, $W_{xh}$는 각각 $(4\times3)$, $(3\times3)$이 될 것이다. $$ W_{hh}=\begin{bmatrix}0.3152 & 0.5083 & 0.9454 \ 0.3008 & 0.6058 & 0.2999 \ 0.9741 & 0.3419 & 0.9133\end{bmatrix} $$ $$ W_{xh}=\begin{bmatrix}0.2973 & 0.2766 & 0.7974 \ 0.3869 & 0.9170 & 0.4125 \ 0.5538 & 0.5646 & 0.5026 \ 0.2206 & 0.6801 & 0.3880\end{bmatrix} $$ 이 값들을 가지고 계산해보면, $$ h_0W_{hh}=\begin{bmatrix}0 & 0 & 0\end{bmatrix} \begin{bmatrix}0.3152 & 0.5083 & 0.9454 \ 0.3008 & 0.6058 & 0.2999 \ 0.9741 & 0.3419 & 0.9133\end{bmatrix} = \begin{bmatrix}0 & 0 & 0\end{bmatrix} $$ $$ x_1W_{xh}=\begin{bmatrix}1 & 0 & 0 & 0\end{bmatrix} \begin{bmatrix}0.2973 & 0.2766 & 0.7974 \ 0.3869 & 0.9170 & 0.4125 \ 0.5538 & 0.5646 & 0.5026 \ 0.2206 & 0.6801 & 0.3880\end{bmatrix} = \begin{bmatrix}0.2973 & 0.2766 & 0.7924\end{bmatrix} $$ $$ \tanh{(h_0W_{hh}+x_1W_{xh})}=\tanh{(\begin{bmatrix}0.2973 & 0.2766 & 0.7924\end{bmatrix})}=\begin{bmatrix}0.2899 & 0.2689 & 0.6626\end{bmatrix} \ $$

3. 출력값 계산 출력값 $o_1$은 다음과 같이 계산된다. $$ o_1 = h_1W_{ho} $$ $W_{ho}$의 shape은 (3 $\times$ 4)가 될 것이다. $$ W_{ho}= \begin{bmatrix}0.3189 & 0.6216 & 0.5164 & 0.7501 \ 0.0260 & 0.6230 & 0.0179 & 0.6123 \ 0.1320 & 0.9009 & 0.3873 & 0.8142\end{bmatrix} $$ $$ h_1W_{ho}=\begin{bmatrix}0.2899 & 0.2698 & 0.6626\end{bmatrix}\begin{bmatrix}0.3189 & 0.6216 & 0.5164 & 0.7501 \ 0.0260 & 0.6230 & 0.0179 & 0.6123 \ 0.1320 & 0.9009 & 0.3873 & 0.8142\end{bmatrix}=\begin{bmatrix}0.1866 & 0.9446 & 0.4106 & 0.9213\end{bmatrix} $$

4. Softmax & Argmax 출력값 $o_1$에 대해 softmax를 취해 각 인덱스에 대한 확률값을 구한다. $$ softmax(o_1)=\begin{bmatrix}0.1546 & 0.3298 & 0.1934 & 0.3222\end{bmatrix} $$ 여기에 argmax를 적용하면 1번 인덱스 값인 'e'가 최종 output이 된다. 이후 다음 output을 계산할 때의 weight matrix는 동일한 weight matrix를 사용해서 다음 시점의 output을 구해나간다.

이 모든 과정을 시각화하면 다음과 같다.

이 과정을 모든 시점에 적용하여 마지막 시점까지의 출력값을 계산한 후, 각 시점에서의 출력값과 실제값의 차이를 모두 합하여 cross-entropy에 적용하여 얻은 loss값을 토대로 backpropagation을 진행해서 weight matrix를 업데이트 한다.

3. RNN의 한계

• Short-term memory

  RNN은 입력 시퀀스가 길어지면 앞 쪽 시점에 대한 정보를 잘 반영하지 못한다는 단점이 있다.

• Gradient vanishing RNN은 activation function으로 tanh를 사용하지만, 값이 거의 항상 1보다 작아서 기울기 소실 문제가 발생한다.

이 문제를 해결하기 위해 gate를 추가하여 선택적으로 정보를 기억할 수 있도록 개선하는 논문들도 등장하였다.

하지만 위에서 든 단어 예측 예시는 입력 시퀀스와 출력 시퀀스가 1:1 매핑이 되서 사용이 가능한거였고, 번역과 같은 task의 경우에는 입력 시퀀스와 출력 시퀀스의 길이가 다를 수 있어서 한계가 있다.

이 한계점들을 극복하기 위해 다음 글에서 소개할 Sequence 2 Sequence가 등장한다. 이에 대한 내용은 다음 게시글에 작성하도록 하겠다.

방대한 언어 데이터를 활용한 NLP foundation model의 등장으로, 다양한 데이터 및 task에 활용되고 있다. 본 논문은 segmentation task에서의 foundation model을 만들어, 새로운 데이터 분포에서도 프롬프트를 이용하여 segmentation을 잘 해내는 것을 목표로 한다.

Task

SAM의 task는 어떠한 프롬프트가 주어져도 유효한 segmentation mask를 반환하는 것이다.

Model

SAM은 다음의 제약조건을 만족하는 모델을 개발하는 것이 목표이다.

1. 다양한 프롬프트 지원
2. 실시간 segmentation mask 계산
3. Ambiguity-aware

본 논문은 다음 제약조건들을 만족하는 간단한 모델을 개발했다. 모델의 구조는 위의 그림과 같다. Image encoder를 통해 image embedding을 계산하고, prompt encoder를 통해 prompt embedding을 계산하고, mask decoder에 이 두 embedding 값을 조합하여 segmentation mask를 뱉는 것이다. Image encoder의 경우 프롬프트가 바뀔 때마다 embedding을 하는 것이 아닌, 한 번 embedding해둔 이미지를 계속 활용하는 것이다. 이렇게 연산과정을 줄여 계산시간을 줄일 수 있게 된다. 또한, ambiguity-aware하기 위해 프롬프트가 주어졌을 때 모델은 여러 개의 mask를 예측하도록 한다.

Data engine

모델을 훈련하기 위해서는 대규모의 데이터셋이 필요하다. 그래야 새로운 데이터 분포에서도 strong generalization을 유지할 수 있다. 하지만 온라인에서는 그렇게 많은 segmentation 데이터셋을 찾을 수 없다. 이를 위해 본 논문은 data engine을 구축한다. 본 논문의 data engine은 assisted-manual, semi-manual, fully automatic 세 단계를 통해 데이터를 구축해나간다.

3. Segment Anything Model

3.1. Image encoder

입력 이미지는 encoder를 통해 $256$차원의 vector로 mapping한다. Encoder는 'Masked Autoencoders Are Scalable Vision Learners'에서 소개한 MAE pre-trained ViT를 사용한다. Encoder를 통해 출력된 image embedding은 $64 \times 64 \times 256$이 된다.

3.2. Prompt encoder

Prompt의 종류는 sparse, dense가 있고, sparse에는 points, boxes, texts가, dense에는 masks가 있다. Prompt encoder를 통해 모든 입력 prompt는 256차원으로 mapping된다.

• Point : 지정된 point의 좌표에 positional encoding을 한 결과와 points의 위치가 피사체 / 뒷배경 중 어디에 속하는 지 나타내는 학습된 embeddings가 합쳐져 256차원으로 mapping
• Box : 좌상단, 우하단 좌표에 positional encoding을 한 결과와 좌상단, 우하단을 나타내는 학습된 embeddings가 합쳐져 256차원으로 mapping
• Mask : 입력 mask를 $\frac{1}{4}$배의 해상도로 바꿔 {$4\times2\times2,stride=2$ convolution}, {$16\times2\times2,stride=2$ convolution} 연산을 수행하여 총 4배의 추가 downsampling을 진해한다. 그 후 {$256\times1\times1$ convolution} 연산을 통해 최종 256차원으로 mapping

• Text : CLIP(Learning Transferable Visual Models from Natural Language Supervision)에서 소개된 text encoder를 통해 256차원으로 mapping

3.3. Mask decoder

Image embeddings와 token embeddings를 통해 segmentation map을 출력하게 되는데, 과정은 다음과 같다. 위 그림에서 output tokens는 target object에 대해 출력할 mask의 수, prompt tokens는 prompt encoder를 통해 embedding된 prompt이다. 이 둘을 합쳐서 'tokens'라 한다.

Decoder에서는 two-way attention을 통해 token/image embedding을 image/token embedding과 cross-attention을 수행하여 각각의 embedding 값을 업데이트한다. 각 layer를 통과하면서 출력되는 shape을 이해하기 위해 attention의 수식에 대해 간단한 복습을 하고 넘어가자.

★ Attention = $Sofmax(Q\cdot K^T) \cdot V$

여기서 $Q$는 query, $K$는 key, $V$는 value이다. 만약 $(64 \times 64 \times 256)$의 image embedding를 $(N \times 256)$의 tokens과 cross-attention을 수행한다면 image embedding이 query, tokens가 key/value가 된다. $$ Q \cdot K^T =(64 \times 64 \times 256) \cdot (256 \times N) \ =(64 \times 64 \times N) \ Softmax(Q \cdot K^T) \cdot V = (64 \times 64 \times N) \cdot (N \times 256) \ = (64 \times 64 \times 256) $$ 즉, attention은 상대와의 관계를 바라보기 위한 것이므로 shape에는 변함이 없다는 것을 알 수 있다.

1. Token embeddings self-attention $(N \times 256)$

2. Token embeddings와 image embeddings 간의 cross-attention $(N \times 256)$

-Image 정보로 token embeddings 업데이트

3. point-wise MLP로 token 업데이트

4. Image embeddings와 token embeddings 간의 cross-attention $(64 \times 64 \times 256)$

-Prompt 정보로 image embeddings 업데이트

이 과정을 총 2회 수행한다. 그 후

5. Image embeddings에 up-convolution 연산을 2번(channel 수는 각각 64 / 32) 통과시켜 해상도는 올리고 채널은 축소 $(256 \times 256 \times 32)$

6. Token embeddings와 image embeddings 간의 cross-attention 연산을 한 번 더 수행 $(N \times 256)$ 한 후 3-layer MLP에 통과시켜 up-convolution된 image embeddings의 채널 수와 맞춤 $(N \times 32)$

7. Up-convolution된 image embeddings와 token embeddings 간의 dot product를 통해 최종 $(N \times 256 \times 256)$ mask 출력

8. 업데이트된 token embeddings를 MLP layer에 통과시켜 mask 별 예측 IoU score 출력 $( N\times 1)$

여기서 IoU score가 가장 높은 mask를 최종 output mask로 선정한다.

4. Segment Anything Data Engine

4.1. Assisted-manual stage

첫 번째 단계에서는, 전문가들이 이미지 각 픽셀을 직접 클릭하여 피사체와 배경을 구분한다. 여기서 본인이 알거나 설명할 수 있는 객체만 라벨링 하되, 해당 객체의 이름이나 설명은 모으지 않는다.

1. 공개된 데이터셋을 이용하여 SAM 모델 초기 훈련
2. 새롭게 annotation된 mask를 이용하여 재훈련
3. 모델 구조 진화 (image encoder ViT-B to ViT-H)

첫 번째 단계를 통해 이미지 당 평균 마스크 수는 20에서 44로 증가한다.

4.2. Semi-automatic stage

두 번째 단계에서의 목표는 모델이 모든 것을 segment할 수 있도록 mask의 다양성을 증가하는데에 있다. 이를 위해 모델이 먼저 다양한 mask를 뽑아내고 사람이 후처리를 한다.

1. 모델이 신뢰도 높은 mask 자동으로 감지
2. 감지된 마스크에 전문가가 추가 라벨링

두 번째 단계를 통해 이미지 당 평균 마스크 수는 44에서 72로 증가한다.

4.3. Fully automatic stage

두 개의 단계를 통해 모델을 향상시키기 충분한 데이터셋을 수집했고, 충분히 ambiguity-aware한 모델을 만들었다. 마지막 단계에서는 이 모델을 통해 다음과 같은 최종 작업을 진행한다.

1. Point가 가리키는 객체의 subpart, part, whole part를 반환
2. IoU (Intersection over Union) 값을 계산하여 신뢰도가 가장 높은 mask 선택
3. NMS (Non Maximum Suppression)을 통해 중복 mask 제거

CNN 기반의 segmentation 기법은 좋은 성능을 내기 위한 특수한 네트워크 구조 및 다양한 훈련 기법을 요구한다. 하지만 이러한 요소 없이 다른 종류의 다양한 데이터셋에 대해 일반적인 결과를 산출할 수 있는 알고리즘이 필요하다. U-Net은 네트워크 구조를 조금씩 수정하여 특정 문제에 과적합하도록 한다. 하지만 본 논문은 네트워크 구조 외에 다른 방면에서 더 영향력이 있다고 제안한다. nnU-Net 또한 'No New U-Net'의 줄임말로, 이전의 U-Net과 모델 구조에 대해서는 다른게 없다고 말하고 있는 것을 알 수 있다.

2. Methods

2.1. Network architectures

위에서 말했듯, 네트워크 구조 외의 다른 면에 초점을 맞추므로 네트워크는 U-Net과 크게 다르지 않다. 다른 점이 있다면 U-Net에서 활성화 함수로 ReLU 함수를 사용했다면 nnU-Net은 Leaky ReLU 를 사용하여 weight가 음수의 값을 갖는 부분에서도 어느정도 살렸고, U-Net에서 batch normalization을 사용했다면 nnU-Net은 instance normalization을 사용하여, batch 단위가 아닌 이미지 단위로 normalize를 하여 각 이미지 고유 정보를 잃는 것을 방지하였다. MRI, CT와 같은 이미지들은 3차원으로 구성되어 있는데, 각 slice로 저장된 2차원 이미지 여러장을 합친 꼴과 같다. 이러한 측면에서 instance normalization을 통해 각 이미지 고유 정보를 가져가는 것이 어쩌면 medical image segmentation에 더 적합하다고 생각된다.

2D U-Net

2D U-Net의 경우 z축에 대한 중요한 정보를 고려하지 않을 수 있으므로 3차원인 medical image를 segmentation하기에는 적합하지 않지만, 데이터셋이 anisotropic하다면 충분히 사용할 가치가 있다.

3D U-Net

3D U-Net이야말로 medical image segmentation에 적합한 모델이지만, GPU 메모리의 한계로 인해 이미지를 패치 단위로 나누어 모델에 입력하게 된다. 이는 큰 이미지의 경우 Field of View의 한계로 모든 context 정보를 못 모을 수 있다.

U-Net cascade

이 모델은 이러한 3D U-Net의 단점을 보완하기 위해 제안되었다. U-Net을 통과해서 얻은 segmentation map을 ground-truth segmentation map에 one-hot encoding 형태로 두 번째 U-Net에 통과시켜 최종 output을 얻어낸다.

Dynamic adaption

GPU를 최대한 활용하여 모델이 입력 이미지의 공간 정보를 충분히 담아내기 위해 입력 patch의 크기, 층 별 pooling 횟수를 잘 설정해야 한다. 또한 네트워크의 수용능력에 따라 batch size 조절이 필요하다. 본 논문은 이러한 hyperparameter를 자동적으로 결정될 수 있도록 한다. 3D U-Net의 경우 입력 이미지가 지정된 patch size보다 작다면, 입력 이미지의 median shape로 patch size를 줄이고, 최종 처리되는 voxel의 수가 기존의 patch size와 batch size로 진행했을 때와 같도록 batch size를 늘리게 된다.

2.2. Preprocessing

다음은 모델에 데이터를 입력하기 전에 진행하는 전처리 과정에 대해 설명한다.

Cropping

모든 입력 데이터셋에 대해 픽셀 값이 0이 아닌 영역을 crop한 후 진행한다.

Resampling

모든 입력 데이터셋에 median voxel spacing을 통한 resampling을 진행하여 네트워크가 공간 정보를 학습할 수 있도록 설정한다. Resampling 과정을 통해 생기는 빈 voxel 값은 interpolation을 통해 채우게 된다.

Normalization

정규화 기법으로는 Z-score normalization을 사용하여, 모든 voxel의 intensity 값의 평균과 표준편차를 기반으로 데이터를 정규화한다. $$ x_{norm} = \frac{x-\mu}{\sigma} $$

2.3. Training procedure

Loss function

Loss function은 U-Net에서 소개된 cross entropy loss와 함께 dice loss를 합쳐 학습하게 된다.

$$ L_{total} = L_{dice} + L_{CE} $$ $$ L_{dc} = -\frac{2}{|K|}\sum_{k \in K}{\frac{\sum_{i\in I}u_i^kv_i^k}{\sum_{i \in I}u_i^k + \sum_{i \in I}v_i^k}} $$

여기서 $K$는 총 클래스의 수, $I$는 총 픽셀 수, $u$는 softmax output, $v$는 ground-truth segmentation map에 softmax output을 기반으로 one-hot encoding을 적용한 값이다. Cross entropy의 경우 각 픽셀마다 weight를 주면서 픽셀 단위의 분류를 돕고, dice loss의 경우 두 segmentation map의 겹치는 정도를 전체적인 segmentation의 품질을 올리게 된다. Dice loss의 경우 예측된 segmentation map과 ground truth segmentation map의 겹쳐지는 정도를 계산하고, dice score가 클수록 예측을 잘 했다고 할 수 있다. 하지만 loss function의 경우 loss 값을 최소화하는 방향으로 학습해야 하므로, dice score는 1에 가까워지도록, dice loss는 $1-L_{dc}$를 통해 0게 가까워지도록 만든다.

Data augmentation

Data augmentation은 rotation, scaling, elastic deformation, gamma correction, mirroring 등의 기법을 사용했고, 훈련 도중에 실시간으로 적용된다. U-Net cascade의 경우 강한 상호작용을 방지하기 위해 2번째 stage 전에 무작위로 연결 요소를 제거한다. 이는 같은 클래스로 분류된 픽셀 덩어리에서 일부러 연결 요소를 제거하여 훈련 데이터의 다양성을 증가하는 것이다.

2.4. Inference

각 patch에 대한 output 값에 대해, 경계 영역이 네트워크의 정확도를 떨어뜨리는 경향을 고려해서 경계에 가까운 영역보다 중앙에 가까운 영역에 가중치를 더 높게 준다. 또한 patch는 이전 패치의 절반 영역이 겹치도록 하여 다음 patch를 설정한다. 이렇게 함으로서 이전 patch에서 경계 영역에 있던 픽셀이 다음 patch에서 중앙 영역에 오게 되면서 모든 픽셀이 모델 학습에 고려되도록 한다. 또한 5-fold cross validation 기법을 통해 얻은 5개의 네트워크를 앙상블하여 test case에 적용한다.

2.5. Postprocessing

학습 데이터에서 각 클래스가 몇 개의 연결성분으로 나타나는지 분석하여, 특정 클래스가 1개로 나타나면 (ex : 간 1개), 그것을 일반적인 성질로 간주한다. 만약 예측 결과 해당 클래스가 여러 개의 클래스로 분할되면, 가장 큰 연결 성분을 띄는 클래스를 제외하고 나머지 연결 성분을 지워 학습한다.

2.6. Ensembling

2D U-Net, 3D U-Net, U-Net cascade 모든 네트워크 중 2개를 조합해보면서 성능을 끌어올리고, 가장 좋은 결과를 내는 모델을 최종 모델로 선정한다.

3. Results

해당 대회 test step에서 보여진 각 문제에 대한 3개의 데이터셋에 대해, BrainTumour를 제외한 나머지 문제에서 최고의 점수를 보이는 것을 알 수 있다.

1. Introduction

U-Net은 이전에 제안됐던 FCN (Fully Convolutional Network) 네트워크 구조를 수정하고 확장한 모델이다. 그 결과 FCN 모델보다 더 적은 데이터로 학습해도 더 정확한 결과를 산출할 수 있었다고 한다.

U-Net이 제안한 method는 총 세 가지이다.

• FCN은 downsampling에서만 feature channel 수를 늘렸다면, U-Net은 upsampling에서도 feature channel의 수를 확대하여 네트워크가 네트워크가 더 높은 해상도를 가진 layer에 위치 정보를 전달할 수 있게 하였다.
• 적은 수의 훈련 데이터가 주어져도, data augmentation 기법을 통해 모델이 좋은 결과를 낼 수 있도록 하였다.
• Weighted loss 라는 새로운 loss term을 통해 객체를 분할할 때, 맞닿는 두 개의 object의 경계선이 더 잘 표현될 수 있도록 하였다.

2. Network

U-Net의 네트워크 구조는 다음과 같다.

Contracting path는 이미지의 해상도를 줄이면서 feature channel의 수를 늘려가고, Expansive path는 이미지의 해상도를 늘려가며 feature channel의 수를 줄여가는 방식으로 진행된다.

2.1. Contracting path

Contracting path는 각 layer마다 두 번의 $3 \times 3$ convolution 연산과 ReLU 함수를 통해 비선형성을 추가한다. 그 후 $2 \times 2$ max pooling을 통해 이미지의 해상도를 줄이게 된다. 각 downsampling마다 feature channel의 수는 2배가 된다.

2.2. Expanding path

Expanding path에서는 $2 \times 2$ up-convolution을 통해 upsampling을 진행한 후, 마찬가지로 두 번의 $3 \times 3$ convolution 연산과 ReLU 함수를 통해 비선형성을 추가한다. 각 upsampling마다 feature channel의 수는 $\frac{1}{2}$배가 된다. 하나의 특이한 구조는 아래 그림의 파란색 타원으로 표시된 부분을 보면, upsampling된 feature map과 contracting path에서 대응되는 feature map을 concatenate한다는 것이다. Feature map을 concatenate하는 이유는 downsampling 시 feature map의 크기가 작아져 위치 정보가 손실되고, upsampling 시 해상도가 올라가지만 원래의 세밀한 정보를 복원하지 못한다. 그래서 skip connection을 통해 global한 정보와 local한 정보를 합치는 것이라 보면 되겠다. 최종 output은 class 수만큼의 1x1 convolution filter와의 연산을 통해 segmentation map을 출력하게 된다.

3. Training

훈련 이미지를 통해 학습하기 위해, 저자는 입력 patch의 사이즈를 키우고 batch 사이즈를 줄여 overhead를 최소화 하고 GPU 메모리 사용을 최대화하였다. 또한, 높은 momentum 값을 설정하여 모델이 바로 전에 본 훈련 이미지를 통해 optimization step를 업데이트할 수 있도록 하였다. 마지막으로, weighted loss를 적용하여 더 나은 segmentation map을 출력할 수 있도록 하였다.

$$ E = \sum_{x \in \Omega}{\omega(\bold{x})\log{(p_{l(\bold{x})}(\bold{x}))}} $$

수식을 뜯어보면, 우선 $\bold{x}$는 각 픽셀의 위치를 나타내고, $l(\bold{x})$는 해당 픽셀의 정답 클래스를 나타낸다. $p_{l(\bold{x})}(\bold{x})$는 softmax 함수로, $$ p_{k}(\bold{x})=\frac{\exp(a_{k}(\bold{x}))}{\sum_{k^{'}=1}^{K}\exp(a_{k^{'}}(\bold{x}))} $$ $a_k(\bold{x})$는 픽셀 위치 $\bold{x}$에서 $k$번째 class의 activation 값이고, $p_k(\bold{x})$는 모든 class의 activation 값으로 나누어 각 class에 대해 해당 픽셀이 어떤 class에 대응하는지 예측한 확률값을 나타내게 된다. 결국 $p_{l(\bold{x})}(\bold{x})$는 정답 class에 대해 어느정도의 확률을 예측했는지 구한 꼴이 된다. $\omega(\bold{x})$는 class 간의 균형을 위해 사전에 계산된 weight map이고, 수식은 다음과 같다. $$ \omega(\bold{x}) = \omega_c(\bold{x}) + \omega_0 \cdot \exp{(-\frac{(d_1(\bold{x})+d_2(\bold{x}))^2}{2\sigma^2})} $$ 위 수식에서 $\omega_c(\bold{x})$는 해당 픽셀의 중요도에 따라 사전에 계산해둔 weight 값이고, $d_1(\bold{x})$, $d_2(\bold{x})$는 해당 픽셀과 제일 가까운 두 객체까지의 최단 거리이다. $exp(-x)$ 이므로 만약 거리가 적으면 경계를 나누는데 중요한 픽셀이므로 가중치를 더 주고, 멀면 값이 작아져서 경계를 나누는데에 중요하지 않다 판단하게 된다.

3.1. Data augmentation

Data augmentation 기법은 학습에 사용할 훈련 데이터의 수가 적을 경우, 데이터의 약간의 변화를 주어 새로운 데이터를 만들어 훈련 데이터의 수를 늘리는데 사용한다. 저자는 탄성 변형 (elastic deformation)을 적용하여 데이터를 증강하였다. 탄성 변형의 예시는 아래 그림을 보면 이해가 될 것이다.

Multi-View Stereo (MVS)는 이미지 간의 대응점을 통한 3차원 재구성을 위해 카메라 내부 파라미터와 외부 파라미터를 요구하지만, 이 파라미터를 구하는 과정은 상당히 번거롭다. DUSt3R은 Dense and Unconstrained Stereo 3D Reconstruction의 약자로, 문장에서 예상할 수 있듯이 카메라 파라미터나 포즈에 대한 사전 정보 없이 이미지를 통해 카메라 파라미터와 포즈를 복원한다. 논문에서 소개한 네트워크는 일반적인 transformer 구조로 이루어져있고, 이미지 쌍으로부터 scene을 3D pointmap을 계산한다. 이 pointmap은 scene의 가하학 정보, 픽셀과 scene 간의 관계, 두 시점 간의 관계를 포함한 풍부한 정보를 가지고 있고, 계산된 pointmap을 이용하여 카메라 파라미터와 scene의 가하학 정보를 복원할 수 있다. 복원된 정보를 이용하여 다양한 3D vision task를 수행할 수 있고, 다양한 분야에서 SOTA의 성능을 보이고 있다. 지금까지 대부분의 3D vision 분야의 논문들이 실험을 위해 아직까지도 COLMAP을 사용하였는데, 이 방법론의 등장으로 새로운 패러다임이 제시되었다고 생각한다.

2. Method

- Pointmap

Pointmap $\bold{X}$는 이미지 각 픽셀에 대응되는 3차원 좌표로, 해상도가 $W \times H$인 이미지가 있다면, $W \times H \times 3$ 의 pointmap을 형성하게 된다. 본 논문은 카메라에서 scene을 향하는 ray가 하나의 3차원 점을 맞춘다고 가정한다.

- Cameras and scene

카메라의 내부 파라미터 $K$와 외부 파라미터 $[R|t]$, world 좌표계에서 정의된 3차원 좌표 $\bold{X}$를 이용하여 이미지 좌표계에 대응하는 픽셀 좌표를 계산할 수 있다.

$$ \begin{bmatrix}i \ j \ 1\end{bmatrix} = K[R|t]\bold{X} $$

카메라의 내부 파라미터 $K$가 주어져 있고, 실제 depth map $D$가 있다면, pointmap $\bold{X}$는 다음과 같이 계산할 수 있다 (pointmap은 카메라 좌표계에 정의되어있다 가정).

$$ \bold{X}{i,j} = K^{-1}D{i,j}\begin{bmatrix}i \ j \ 1\end{bmatrix} $$

$n$ 카메라와 $m$ 카메라가 있다면, $m$ 카메라 좌표계에서 정의된 $n$ 카메라에 대응되는 pointmap $\bold{X}^{n,m}$은 $n$ 카메라의 외부파라미터 $P_n$ 과 $m$ 카메라의 외부파라미터 $P_m$ 을 통해 다음과 같이 계산할 수 있다.

$$ \bold{X}^{n,m} = P_mP_n^{-1}h(\bold{X}^n) $$ 여기서 $h$는 homogeneous 좌표계를 뜻한다. $P_m, P_n$은 $[R|t]$ 이다. 즉, $P_n^{-1}h(\bold{X}^n)$을 통해 pointmap을 world 좌표계에 정의하고, $P_m$과의 연산을 통해 $m$ 카메라 좌표계에 정의된 pointmap을 계산하게 되는 것이다.

2.1. Overview

네트워크 $f$는 두 개의 이미지 $I^1, I^2$를 입력으로 받아 각 이미지 픽셀에 대응하는 pointmap $X^{1,1}, X^{2,1} \in \real^{W \times H \times 3}$과 각 픽셀에 대응하는 pointmap의 확신 각 픽셀에 대응하는 pointmap의 신뢰도를 나타내는 confidence map $C^{1,1}, C^{2,1} \in \real^{W \times H}$ 를 출력한다. 이렇게 나타낸 이유는 두 pointmap 모두 1번 카메라 좌표계에 대해 나타내므로 다음과 같이 표현한다.

네트워크 구조는 다음과 같다.

우선 기존의 ViT와 마찬가지로 이미지들을 패치로 나누어 flatten한 후 linear projection을 통해 일정한 크기의 latent vector로 나타낸 후 각각의 encoder에 입력하게 된다. 이 때 두 개의 encoder는 weight를 공유하여 동일한 feature 표현을 보장하여 비슷한 feature space에 임베딩하도록 한다. $$ F^1 = Encoder(I^1) \ F^2 = Encoder(I^2) $$

Decoder는 일반적인 Transformer 구조를 이용하여 각 decoder block에서 self-attention과 cross-attention을 순차적으로 진행한다. $$ G_i^1 = DecoderBlock(G_{i-1}^1, G_{i-1}^2) \ G_i^2 = DecoderBlock(G_{i-1}^2, G_{i-1}^1) $$

$i$는 $B$ 개의 decoder block에 대해 $1, \dots, B$ 번 수행하고, $G_0^1, G_0^2$ 는 encoder를 통해 나온 $F^1, F^2$로 정의한다. $$ G_0^1, G_0^2 = F^1, F^2 \ G_1^1, G_1^2 = DecoderBlock(G_0^1, G_0^2), DecoderBlock(G_0^2, G_0^1) \ \vdots \ G_B^1, G_B^2 = DecoderBlock(G_{B-1}^1, G_{B-1}^2), DecoderBlock(G_{B-1}^2, G_{B-1}^1) $$

마지막으로, 각 branch에서의 head가 decoder의 token들을 받아서 pointmap과 confidence map을 출력하게 된다. $$ X^{1,1}, C^{1,1} = Head^1(G_0^1, \dots, G_B^1) \ X^{2,1}, C^{2,1} = Head^2(G_0^2, \dots, G_B^2) $$

2.2 Training Objective

학습을 위한 loss function으로는 3D Regression loss와 Confidence-aware loss를 사용한다.

2.2.1. 3D Regression loss

실제 pointmap을 $\bold{\bar{X}}^{1,1}, \bold{\bar{X}}^{2,1}$이라 하자. 실제 pointmap은 이전에 이야기했듯, 카메라의 내부 파라미터와 실제 depth를 알고 있다면 쉽게 구할 수 있다. 네트워크는 실제 pointmap에 대응하는 depth map $D^1, D^2$에서 유요한 픽셀 $i \in D^v \enspace ( v \in {1,2}$)에 대해 실제 pointmap과 예측된 pointmap 간의 Euclidian distance를 최소화하는 방향으로 학습한다. $$ \mathcal{L}{regr}(v,i) = \bigg| \bigg|\frac{1}{z}\bold{X}^{v,1} - \frac{1}{\bar{z}}\bold{\bar{X}}^{v,1}\bigg| \bigg| \ z = \frac{1}{|D^1|+|D^2|}\sum{v \in {1,2}}\sum_{i \in D^v}||\bold{X}_i^v|| $$ 여기서 z는 예측된 pointmap과 실제 pointmap간의 scale을 normalize하기 위해 사용된다.

2.2.2. Confidence-aware looss

하늘처럼 무한한 곳이거나, 투명한 물체와 같은 부분에 대해서는 잘못 계산된 3차원 좌표들이 있을 수 있다. 따라서 이 부분에 해당하는 픽셀에 대해 얼마나 확신하는지를 나타내는 점수를 함께 학습한다. 최종 학습의 목적 함수는, 모든 유효한 픽셀에 대해 regression loss에 confidence를 곱한 것이다. $$ \mathcal{L}{conf} = \sum{v \in {1,2}} \sum_{i \in D^v} C_i^{v,1}\mathcal{L}_{regr}(v,i) - \alpha \log{C_i^{v,1}} $$

2.3. Downstream Applications

2.3.1. Point matching

2차원 이미지에서는 SIFT와 같은 feature를 얻는 기법을 이용하여 keypoint를 찾았다. 본 논문은 두 카메라에 대해 얻은 pointmap을 이용하여, pointmap간의 가장 가까운 점들을 찾는 Nearest Neighbor 기법을 이용하여 두 카메라 모두 대응하는 3차원 좌표를 계산한다.

2.3.2. Recovering intrinsics

카메라의 내부 파라미터는 다음과 같이 정의된다. $$ K = \begin{bmatrix}f_x & 0 & C_x \ 0 & f_y & C_y \ 0 & 0 & 1\end{bmatrix} $$ 우선 이미지의 높이와 너비를 알고 있고, 높이와 너비가 같으므로 주점인 $C_x, C_y$는 쉽게 구할 수 있다. $$ C_x = C_y = \frac{H}{2} $$ 내부 파라미터를 구하기 위해서는 카메라의 focal length만 계산하면 된다. 이미지 평면의 좌표는 카메라 좌표계에서 정의된 pointmap과 카메라의 내부 파라미터 간의 내적을 통해 계산된다. (1번 카메라에 대한 수식) $$ \begin{bmatrix}i \ j \ 1 \end{bmatrix} = \begin{bmatrix}f & 0 & C \ 0 & f & C \ 0 & 0 & 1 \end{bmatrix}\begin{bmatrix}X^{1,1}{i,j,0} \ X^{1,1}{i,j,1} \ X^{1,1}_{i,j,2}\end{bmatrix} $$

Focal length인 $f$는 다음과 같은 수식을 최소화하는 $f$를 구하여 카메라의 focal length를 계산할 수 있다. $$ f_1^* = \argmin_{f_1}{\sum_{i=0}^{W} \sum_{j=0}^{H} C_{i,j}^{1,1}\bigg|\bigg| (i,j)-f\frac{(X_{i,j,0}^{1,1}, X_{i,j,1}^{1,1})}{X_{i,j,2}^{1,1}} \bigg|\bigg|} $$ 여기서 $C_{i,j}^{1,1}$은 주점이 아닌 1번 카메라 pointmap의 각 픽셀에 대응하는 신뢰도를 계산한 confidence map이다.

2.3.3. Relative pose estimation

두 카메라 간의 상대적인 위치를 추정하는 방법은 세 가지가 있다.

1. 두 이미지 간의 2D matching을 진행하고, 위의 방법을 통해 두 카메라의 내부 파라미터를 계산한다면, epipolar geometry를 통해 두 카메라 간의 변환관계를 설명하는 $R, \bold{t}$를 계산할 수 있다.
2. Procrustes alignment를 이용해서 두 pointmap을 비교할 수 있다. $$ P^* = \argmin_{\sigma,R,\bold{t}}{\sum_{i}{C_i^{1,1}C_i^{1,2}||\sigma(RX_i^{1,1}+\bold{t})-X_i^{1,2}||^2}} $$
3. 더 robust한 솔루션은 RANSAC과 PnP (Perspective n point) 알고리즘을 사용하는 것이다.

2.4. Global Alignment

NeRF는 3D object 또는 scene을 촬영한 이미지들로 향하는 ray 위의 점을 샘플링해서 학습한 모델을 이용하여 새로운 시점에 대한 이미지를 렌더링하는 기술이다. 이 기법은 일정한 거리에서 촬영한 이미지들에 대해 좋은 렌더링 결과를 보여줬지만, 가까이에서 바라보는 시점에 대한 이미지는 흐릿한 현상이 발생했고, 거리가 먼 시점에 대한 이미지는 계단현상이 발생하는 한계를 보였다. 또한, NeRF는 fine network와 coarse network를 통해 최종 픽셀 값을 결정하게 되는데, 큰 해상도의 이미지는 ray 수도 많을 뿐더러 샘플링하는 점의 수까지 많아지다보니 모델에 입력할 query가 많아서 하나의 이미지를 렌더링하는데 오랜 시간이 걸린다는 단점이 있었다.

본 논문은 이러한 한계점을 극복하고자 새로운 방식을 제안한다.

1. Cone tracing

우선, 기존의 NeRF는 카메라에서 object나 scene의 3D point를 향해 발사되는 ray 위의 점들을 샘플링한 후 누적하여 픽셀의 최종 색을 결정했다. Mip-NeRF는 이런 얇은 광선 대신 콘(cone) 형태의 ray를 발사한 후 일정 구간으로 쪼개진 conical frustum 안의 좌표들을 사용하게 된다. $t_0, t_1$ 시점 사이의 Conical frustum 안에 있는 3D 좌표들인 $\bold{x}$는 다음과 같다. $$ F(\bold{x}, \bold{o}, \bold{d}, \dot{r}, t_0, t_1)=1 \bigg{\bigg(t_0 < \frac{\bold{d}^T(\bold{x}-\bold{o})}{||\bold{d}^2||} < t_1 \bigg) \wedge \bigg(\frac{\bold{d}^T(\bold{x}-\bold{0})}{||\bold{d}||_2||\bold{x}-\bold{o}||_2} > \frac{1}{\sqrt{1+(\frac{\dot{r}}{||\bold{d}||_2})^2}} \bigg) \bigg} $$ Wedge 기호를 기준으로 왼쪽 부분은 ray marching이 수행되는 거리의 범위, 오른쪽 부분은 각도의 범위를 나타낸 것으로, 해당 범위에 속한 좌표들만을 사용하는 것이다.

하지만, 얻어낸 conical frustum 위의 좌표들에 단순하게 positional encoding($\gamma(\bold{x})$)을 해서 계산을 한다면 다음의 적분식은 closed form 형태가 아니므로 계산을 할 수 없다. $$ \gamma^*(\bold{o}, \bold{d}, \dot{r}, t_0, t_1) = \frac{\int{\gamma(\bold{x})F(\bold{x}, \bold{o}, \bold{d}, \dot{r}, t_0, t_1)d\bold{x}}}{\int{F(\bold{x}, \bold{o}, \bold{d}, \dot{r}, t_0, t_1)}d\bold{x}} $$

그래서 이 conical frustum을 추후에 다룰 integrated positional encoding (IPE) 를 위해 다변수 가우시안으로 근사화한다. 이를 위해 $F(\bold{x}, \cdot)$의 평균과 공분산을 구해야 한다. conical frustum은 원형이고 대칭이므로, 이 Gaussian 분포는 광선과 평행한 방향으로의 평균 거리 $\mu_t$, 광선과 평행한 방향으로의 분산 $\sigma^2_t$, 광선과 수직인 방향으로의 분산 $\sigma^2_r$ 세 개의 값으로 설명할 수 있다.

우선 기존의 Cartesian space에 정의된 좌표 $\bold{x} = (x,y,z)$는 광선이 퍼지는 효과를 제대로 반영하기 어려우므로, $(r, t, \theta)$ 좌표계로 변환해준다. 여기서 $r$은 radius로 광선의 퍼짐 정도를, $t$는 광선의 깊이(depth), $\theta$는 광선의 기울기를 나타낸다. 여기서 $x, y, z$는 각각 $rt\cos{\theta}, rt\sin{\theta}, t$로 나타낼 수 있고, 이렇게 변환된 변수는 미분가능하게 해준다. 그 이유는 : $$ dxdydz = |J|drd\theta dt $$ 여기서 Jacobian 행렬 $J$는 다음과 같이 구성되고, $$ J = \begin{bmatrix}\frac{dx}{dr} & \frac{dx}{dt} & \frac{dx}{d\theta} \ \frac{dy}{dr} & \frac{dy}{dt} & \frac{dy}{d\theta} \ \frac{dz}{dr} & \frac{dz}{dt} & \frac{dz}{d\theta}\end{bmatrix} \ = \begin{bmatrix}t\cos{\theta} & r\cos{\theta} & -rtsin{\theta} \ t\sin{\theta} & r\sin{\theta} & rt\cos{\theta} \ 0 & 1 & 0\end{bmatrix} $$

그럼 $|J| = rt^2$이 되므로 $dxdydz = (rt)drd\theta dt$가 되기 때문이다.

Conical frustum $F(\bold{x}, \cdot)$의 부피 $V$는 다음과 같다. $$ V = \int_{0}^{2\pi}\int_{t_0}^{t_1}\int_{0}^{\dot{r}}rt^2drdtd\theta \ = \frac{\dot{r}^2}{2}\cdot\frac{t_1^3-t_0^3}{3}\cdot2\pi \ = \pi\dot{r}\frac{t_1^3-t_0^3}{3} $$

따라서 conical frustum에서 균일하게 샘플링된 점들의 확률 밀도 함수는 $\frac{rt^2}{V}$ 이 되고, t의 1차 모멘트(moment)와 2차 모멘트는 다음과 같다. $$ E[t] = \frac{1}{V}\int_{0}^{2\pi}\int_{t_0}^{t_1}\int_{0}^{\dot{r}}t\cdot rt^2drdtd\theta \ = \frac{1}{V} \cdot \pi\dot{r}^2\frac{t_1^4-t_0^4}{4} \ = \frac{3(t_1^4-t_0^4)}{4(t_1^3-t_0^3)} $$ $$ E[t^2] = \frac{1}{V}\int_{0}^{2\pi}\int_{t_0}^{t_1}\int_{0}^{\dot{r}}t^2\cdot rt^2drdtd\theta \ = \frac{1}{V} \cdot \pi \dot{r}^2 \frac{t_1^5 - t_0^5}{5} \ = \frac{3(t_1^5-t_0^5)}{5(t_1^3-t_0^3)} $$

이를 활용하여 $\sigma_t^2= E[t^2]-E[t]^2$을 구할 수 있다. Conical frustum의 반경 $r$에 대한 분산은 대칭성에 의해 균일하게 점을 샘플링하면 기댓값이 중심인 $0$으로 오게 된다. 즉, $\sigma_r^2 = E[x^2]$이 된다. $$ E[x^2] = \frac{1}{V}\int_{0}^{2\pi}\int_{t_0}^{t_1}\int_{0}^{\dot{r}}(rt\cos{\theta})^2\cdot rt^2drdtd\theta \ = \int_{t_0}^{t_1}\int_{0}^{\dot{r}}r^3t^4\int_{0}^{2\pi}\cos^2{\theta}d\theta drdt \ = \frac{1}{V}\cdot\frac{\dot{r}^4}{4}\cdot\frac{t_1^5-t_0^5}{5}\cdot\pi \ = \dot{r}^2\cdot\frac{3(t_1^5-t_0^5)}{20(t_1^3-t_0^3)} $$

이렇게 구한 $\mu_t, \sigma_r^2, \sigma_t^2$을 이용하면, $t_0$와 $t_1$이 큰 차이를 낸다면 괜찮지만, 서로 가까이 위치해있을 때 이 차이를 큰 지수 값으로 계산하게 되면 값이 $0$ 또는 $NaN$을 뱉는 경우가 있으므로, $t_0, t_1$ 대신 $t_0$와 $t_1$의 중점인 $t_\mu = \frac{t_0 + t_1}{2}$, $t_0$와 $t_1$사이의 거리의 절반인 $t_\delta = \frac{t_0 - t_1}{2}$ 로 파라미터를 바꾸어 사용한다고 한다. 최종 변수들은 다음과 같다.

$$ \mu_t = t_u + \frac{2t_\mu t_\delta^2}{3t_\mu^2 + t_\delta^2} \ \sigma^2_t = \frac{t_\delta^2}{3} - \frac{4t_\delta^4(12t_\mu^2-t_\delta^2)}{15(3t_\mu^2+t_\delta^2)^2} \ \sigma^2_r = \dot{r}^2\bigg(\frac{t_\mu^2}{4} + \frac{5t_\delta^2}{12} - \frac{4t_\delta^4}{15(3t_\mu^2+t_\delta^2)} \bigg) $$

여기서 $t_\mu$는 $t_0$와 $t_1$의 중점이고, $t_\delta$는 $t_0$와 $t_1$ 사이의 거리의 반을 나타낸다. $$ t_\mu = \frac{(t_0+t_1)}{2} \ t_\delta = \frac{(t_0-t_1)}{2} $$

Conical frustum 좌표계에서 정의된 가우시안 분포는 다음과 같이 world coordinate으로 변형할 수 있다.

$$ \mu = \bold{o} + \mu_t\bold{d} \ \Sigma = \sigma^2_t(\bold{d}\bold{d}^T) + \sigma^2_r\bigg(\bold{I}-\frac{\bold{d}\bold{d}^T}{||\bold{d}||_2^2} \bigg) $$

여기서 $\mu$는 평균 벡터로, 데이터의 중심을 나타내고, $\Sigma$는 공분산 행렬로, 데이터의 퍼짐 정도를 나타낸다. 이를 통해 world coordinate에서 데이터의 분포를 정확히 모델링할 수 있게 된다.

2. Integrated positional encoding (IPE)

다음 그림은 1차원 공간에서 positional encoding(PE)와 integrated positional encoding(IPE)의 차이를 보여준다. NeRF는 각 ray 위에서의 점들을 샘플링하고 모든 주파수를 동일하게 인코딩한다. Aliasing이라고 불리는 계단현상은 고주파 신호가 저주파 신호로 왜곡되는 현상인데, 이는 샘플 간격에 비해 고주파 특성들이 너무 많이 표현되면서 발생하는 것이다. IPE는 주파수 구간을 통합함으로써 고주파 특성들이 구간의 크기에 비례해 감소하게 하여 anti-aliasing 효과를 얻을 수 있다.

우선 coordinate space에서의 좌표 $\bold{x}$를 frequency space로 Fourier feature를 통해 embedding할 수 있다.

$$ \gamma(\bold{x}) = \begin{bmatrix}\sin{(\bold{P}\bold{x})}\\cos{(\bold{P}\bold{x})}\end{bmatrix} \quad where ; \bold{P}=\begin{bmatrix}2^0 & 0 & 0 & 2^1 & 0 & 0 & \cdots & 2^{L-1} & 0 & 0 \ 0 & 2^0 & 0 & 0 & 2^1 & 0 & \cdots & 0 & 2^{L-1} & 0 \ 0 & 0 & 2^0 & 0 & 0 & 2^1 & \cdots & 0 & 0 & 2^{L-1}\end{bmatrix}^T $$

이를 이용해 Gaussian 분포의 데이터를 기존의 좌표계에서 positional encoding을 적용하여 고차원 공간으로 변환할 수 있다. 공분산은 $Cov[A\bold{x}, B\bold{y}] = A,Cov[\bold{x}, \bold{y}]B^T$라는 성질을 이용하여 다음과 같이 변환된다.

$$ \mu_{\gamma} = \bold{P}\mu \ \Sigma_\gamma = \bold{P}\Sigma\bold{P}^T $$

최종적으로 IPE feature를 얻는 단계는 이 다변량 가우시안에서 무작위로 샘플링된 값들의 평균을 나타내는 기댓값을 계산하는 것이다. $$ \gamma(\mu, \Sigma) = E_{\bold{x}\sim \mathcal{N}(\mu_\gamma, \Sigma_\gamma)}[\gamma(\bold{x})] $$ $$ = \begin{bmatrix}\sin{(\mu_\gamma)}\circ \exp{(-(\frac{1}{2})diag(\Sigma_\gamma))} \ \cos{(\mu_\gamma)}\circ \exp{(-(\frac{1}{2})diag(\Sigma_\gamma))}\end{bmatrix} $$

3. Optimization

NeRF는 coarse network, fine network 두 개의 네트워크를 사용하여 최종 이미지 픽셀 값을 정하였다. 하지만, Mip-NeRF는 multi-scale representation이 가능하므로 단 하나의 네트워크만을 사용하여 이미지 픽셀 값을 결정하므로 모델 사이즈도 절반으로 줄일 수 있고, 샘플링도 더 효과적이고, 전체적인 알고리즘도 단순해진다.

$$ \min_{\Theta}\sum_{\bold{r}\in \mathcal{R}}\bigg(\lambda||C^(\bold{r}) - C(\bold{r};\Theta,\bold{t}^c)||^2_2 + ||C^(\bold{r}) - C(\bold{r};\Theta,\bold{t}^f)||^2_2\bigg) \quad where ; \lambda=0.1 $$

3차원 scene과 이 scene을 보고 있는 두 개의 카메라가 있을 때, 우리는 3D 좌표와 각 카메라의 이미지 평면에 투영된 2D 좌표를 이용하여 두 카메라 사이의 기하학적 관계를 설명할 수 있다. 이 두 카메라 사이의 기하학적 관계를 찾는 것을 epipolar geometry라고 한다.

위 그림에서 표시된 변수들은 다음과 같다.

• $\bold{C}, \bold{C^{'}}$ : 카메라 중점

• $\bold{e}, \bold{e}^{'}$ : Epipoles

• $\overline{\bold{x}\bold{e}}, \overline{\bold{x}^{'}\bold{e}^{'}}$ : Epipolar lines

• $\triangle \bold{X}\bold{C}\bold{C}^{'}$ : Epipolar plane

Epipolar geometry를 통해 두 카메라 및 두 이미지 사이의 변환관계를 구할 수 있다면, 두 개의 카메라에서 각 이미지 위의 좌표를 향해 ray를 쏴서 3D reconstruction을 진행할 수 있으므로 3D vision을 공부한다면 꼭 한 번 짚고 넘어가야 하는 개념이다.

2. Fundamental matrix

Fundamental matrix는 $\bold{C}, \bold{C}^{'}, \bold{X}$ 가 있을 때, 대응점 쌍 $(\bold{x}, \bold{x}^{'})$에 대해 유일하게 결정되는 $3 \times 3$ 행렬이다. $$ \bold{x^{'}}^{T}F\bold{x}=0 $$ 이 행렬을 안다면, 두 이미지 사이의 대응점 쌍을 찾을 수 있다. 즉, 한 점 $\bold{x}$와 $F$가 주어지면 $\bold{x}^{'}$을 추론하는게 가능해진다.

2.1. 8-point algorithm

위 식을 다음과 같이 변형할 수 있다.

$$ \bold{x^{'}}^TF\bold{x}=0 $$

$$ \begin{bmatrix}x^{'}n & y^{'}_n & 1\end{bmatrix} \begin{bmatrix}F{11} & F_{12} & F_{13} \ F_{21} & F_{22} & F_{23} \ F_{31} & F_{32} & F_{33}\end{bmatrix} \begin{bmatrix}x_n \ y_n \ 1\end{bmatrix} = 0 $$

$$ (F_{11}x_n + F_{12}y_n + F_{13})x_n^{'} + (F_{21}x_n + F_{22}y_n + F_{23})y_n^{'} + F_{31}x_n+ F_{32}y_n + F_{33} = 0 $$

$$ \begin{bmatrix}x_nx_n^{'} & x_n^{'}y_n & x_n^{'} & x_ny_n^{'} & y_ny_n^{'} & y_n^{'} & x_n & y_n & 1\end{bmatrix} \begin{bmatrix}F_{11} \ F_{12} \ F_{13} \ F_{21} \ F_{22} \ F_{23} \ F_{31} \ F_{32} \ F_{33}\end{bmatrix} = 0 $$

Fundamental matrix는 scale이 정해져있지 않다. 즉, $F$나 임의의 상수 $k$가 곱해진 $kF$나 같은 결과를 뱉는다.

$$ \bold{x^{'}}^{T}F\bold{x} = \bold{x^{'}}^{T}(kF)\bold{x} = 0 $$

그래서 총 9개의 값 중 scale인 $F_{33}$을 제외한 8개의 파라미터만 추정하면 되므로, 8개의 feature point가 주워진다면 fundamental matrix를 구할 수 있다.

$$ \begin{bmatrix}x_1x_1^{'} & x_1^{'}y_1 & x_1^{'} & x_1y_1^{'} & y_1y_1^{'} & y_1^{'} & x_1 & y_1 & 1 \ \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots & \vdots \ x_8x_8^{'} & x_8^{'}y_8 & x_8^{'} & x_8y_8^{'} & y_8y_8^{'} & y_8^{'} & x_8 & y_8 & 1\end{bmatrix} \begin{bmatrix}F_{11} \ F_{12} \ F_{13} \ F_{21} \ F_{22} \ F_{23} \ F_{31} \ F_{32} \ F_{33}\end{bmatrix} = 0 \ AF = 0 $$

행렬 $A$는 언제나 full rank일 것이다. $A$에 Singular Value Decomposition (SVD) 를 적용하여 $UDV^T$로 분리할 수 있는데, 여기서 $V^T$의 마지막 열을 다시 $3 \times 3$ 으로 reshape 해준다면 우리가 구하는 fundamental matrix $F$가 된다.

3. Essential matrix

Fundamental matrix는 uncalibrated인 상황에서 사용된다면, essential matrix는 calibrated인 상황에서 사용되므로 카메라의 파라미터를 알고있다고 가정한다. Fundamental matrix를 구하기 위해 사용되는 $\bold{x}$는 world coordinate에서 정의된 $\bold{X}$에서 camera coordinate으로 변환하는 extrinsic parameters $R, \bold{t}$, camera coordinate에서 pixel coordinate으로 변환하는 intrinsic parameters $K$ 를 이용하여 구해진다. $$ \bold{x} = K[R|\bold{t}]\bold{X} $$

Essential matrix는 camera coordinate에서의 대응점 $\overline{\bold{x}}, \overline{\bold{x}^{'}}$이 주어졌을 때의 fundamental matrix가 된다. $K[R|\bold{t}]$를 $P$라 하면 $$ \overline{\bold{x}}=K^{-1}P\bold{X} $$ $$ \overline{\bold{x^{'}}}^{T}E\overline{\bold{x}}=0 $$

캘리브레이션이 되어있으므로 우리는 $K$를 알고 있고, fundamental matrix와 essential matrix는 $E=K^{'^T}FK$라는 관계를 가지고 있어서 쉽게 계산할 수 있다.

4. Pose estimation

Essential matrix $E$를 분해하여 두 카메라 좌표계 사이의 변환관계를 설명하는 $R, \bold{t}$를 구할 수 있다. 우선 $E$는 SVD를 통해 분해할 수 있다. 여기서 essential matrix의 특성으로 인해 두 개의 singular value는 같고, 하나는 0이 된다. $$ E = U\begin{bmatrix}1 & 0 & 0 \ 0 & 1 & 0 \ 0 & 0 & 0\end{bmatrix}V^T \quad (up \ to \ scale) $$

여기서 $\begin{bmatrix}0 & 1 & 0 \ -1 & 0 & 0 \ 0 & 0 & 1\end{bmatrix}$로 구성된 행렬 $W$를 이용하여

$ R_1 = UWV^T \ R_2 = UW^TV^T \ \bold{t} = U\begin{bmatrix}0 \ 0 \ 1\end{bmatrix} or -U\begin{bmatrix}0 \ 0 \ 1\end{bmatrix} $

이렇게 최종 $R, \bold{t}$ 후보를 구하게 된다. 이들 중 복원되는 3D 좌표가 가장 많은 $R, \bold{t}$가 최종 $R, \bold{t}$가 된다.

0. Abstract

NeRF (Neural Radiance Field)는 복잡한 scene에 대한 새로운 시점에서의 이미지를 렌더링하는 task (novel view synthesis)를 수행한다. 모델에 5차원의 좌표 정보 (ray의 방향 정보 ($\theta , \phi$) 와 ray의 위의 3차원 좌표 ($x, y, z$)) 가 입력되면, 출력된 color와 density를 통해 volume rendering을 거쳐 이미지를 구성하게 된다. 본 논문은 효과적인 모델 최적화를 통해 새로운 시점에 대한 현실적인 이미지를 렌더링하는 방법을 소개한다.

1. Introduction

Novel view synthesis를 위한 학습 방법은 다음과 같다.

1. 카메라에서 scene으로 향하는 ray 위의 3차원 점 샘플링.
2. 샘플링된 점과 ray의 방향을 모델에 입력하여 color ($r,g,b$)와 density $\sigma$ 출력
3. Volume rendering을 통해 ray 위의 각 점의 color와 density를 누적하여 픽셀 값 결정
4. 예측된 픽셀 값과 실제 픽셀 값과의 차이를 이용하여 loss 값 계산

이 과정은 미분이 가능하므로 gradient descent 기법을 통해 실제 이미지와 rendering된 이미지 간의 차이를 최소화하도록 모델을 학습할 수 있다.

하지만, 이렇게 간단한 구현으로는 몇 가지의 문제점이 있다.

• 복잡한 scene을 나타낸다면 고화질로 나타내기에 모델이 수렴하지않는다.
• 카메라 ray마다의 비효율적인 샘플링을 수행하게 된다.

이러한 이슈를 해결하기 위해 본 논문은

• Positional encoding을 도입하여 모델이 고화질의 이미지를 나타낼 수 있도록 한다.
• Hierarchical sampling을 통해 최종 렌더링에 기여하는 점을 샘플링한다.

2. Neural Radiance Field Scene Representation

본 논문은 scene을 연속적으로 표현하기 위해 ray 위의 좌표 $\bold{x}$와 ray의 방향 $\bold{d}$를 입력했을 때 color $\bold{c}$와 density $\sigma$를 얻을 수 있도록 MLP 네트워크의 가중치 $\Theta$를 최적화한다. $$ F_\Theta : (\bold{x}, \bold{d}) \rightarrow (\bold{c}, \sigma) $$ 또한, 단일 view에서 뿐만아니라 다양한 view에서 일관되기 위해 $\sigma$는 $\bold{x}$에 대해서만 계산하고, $\bold{c}$는 $\bold{x}$와 $\bold{d}$ 모두에 대해 계산하도록 한다. 이를 만족하기 위해 $F_\Theta$는 먼저 positional encoding 된 $\bold{x}$를 입력으로 받아 8개의 fully-connected layer로 구성된 네트워크에 통과 (다섯번째 layer에서 skip connection을 통해 입력값 한 번 더 입력) 시켜서 $\sigma$와 256차원의 feature vector를 얻고, 이 feature vector를 $\bold{d}$와 합쳐서 한 개의 fully-connected layer를 통과시켜 $\bold{c}$를 얻도록 한다. $\gamma(\bold{x}) ,\gamma(\bold{d})$ 옆에 있는 숫자는 positional encoding에서 다루겠다.

3. Volume Rendering with Radiance Fields

하나의 ray $\bold{r} = \bold{o}+t\bold{d}$에서 샘플링된 좌표 $\bold{x}$들과 해당 ray의 방향 $\bold{d}$를 통해 얻은 density $\sigma$들과와 color $\bold{c}$들이 있을 때, volume rendering을 통해 픽셀의 최종 색 $\bold{C}{(\bold{r})}$을 정할 수 있다. $$ \bold{C}{(\bold{r})}=\int_{t_n}^{t_f}T(t)\sigma(\bold{r}(t))\bold{c}(\bold{r}(t),\bold{d})dt, \quad where \enspace T(t) = exp\bigg(-\int_{t_n}^{t}\sigma(\bold{r(s)})ds\bigg) $$

여기서 $\sigma$와 $\bold{c}$는 위에서 이야기했듯 모델을 통해 얻은 출력이고, $t_n, t_f$는 ray의 시작과 끝, $T$는 누적 투과율 (accumulated transmittance)를 나타낸다. 투과도가 필요한 이유는 ray가 scene을 향해 뻗어나가다가 물체에 부딪혔을 때, 물체 내부의 점은 필요성이 떨어지므로 이를 표현하기 위해 추가한 것이다. 즉, ray가 부딪히기 전인 $t$까지의 정보를 누적하여 색을 결정하는데 사용한다. 위 식은 구적법을 이용하여 수치적으로 근사화하기 위해 strtified sampling을 이용하여 $[t_n, t_f]$ 구간을 N개의 균등한 구간으로 나눈 후, 각 구간에서 하나의 sample을 random하게 샘플링하여 모델에 입력할 $\bold{x}$를 얻는다. $$ t_i \sim \mathcal{U}\bigg[t_n+\frac{i-1}{N}(t_f-t_n), t_n+\frac{i}{N}(t_f-t_n) \bigg] $$

최종 렌더링 식은 다음과 같다. $$ \bold{C}{(\bold{r})} \approx \sum{i=1}^{N}T_i(1-exp(-\sigma_i\delta_i))\bold{c}i, \quad where \enspace Ti = exp \bigg( -\sum{j=1}^{i-1}\sigma_i\delta_i \bigg) $$

위 식에서 $\delta$는 샘플링된 두 인접한 좌표 간의 거리를 나타낸 것이다. $\sigma$는 밀도, 즉 해당 점이 물체와 부딪혔는지 통과했는지를 scalar 값으로 나타낸 것이다. 수식의 각 부분을 'Depth-supervised NeRF'에서 시각화를 잘 해두어서 가져와봤다.

그림을 보면, 결국 누적투과율에 의해 물체와 처음 부딪히는 점이 큰 기여를 하게 되고, 이후의 점들은 죽는 것을 볼 수 있다.

모델을 통해 얻은 $\sigma$와 $\bold{c}$와 이 volume rendering 식을 이용하여 $\bold{C}$를 얻는다.

4. Optimizing a Neural Radiance Field

Introduction에서 이야기했듯이 이 방법만으로는 SOTA 급의 퀄리티를 내기는 힘들다. 본 논문은 두 가지의 방법을 추가하여 성능을 개선했는데, 그것은 바로 positional encoding과 hierarchical sampling이다.

4.1. Positional encoding

모델 $F_\Theta$에 $\bold{x}$와 $\bold{d}$를 바로 입력한다면 렌더링 결과가 color나 geometry가 고주파 요소, 즉 세밀한 디테일을 잘 표현하지 못한다고 한다. 이는 딥러닝이 저주파 함수를 학습하는 경향이 있어서 그렇다. 네트워크가 세밀한 디테일을 표현할 수 있도록 하기 위해 본 논문은 입력값을 더 높은 차원의 공간으로 mapping하는 positional encoding을 도입했다. $$ F_\Theta=F^{'}_\Theta \cdot \gamma $$ $$ \gamma(p)=(sin(2^{0}\pi p), cos(2^{0}\pi p), \dots, sin(2^{L-1}\pi p), cos(2^{L-1}\pi p)) $$

Positional encoding에 대한 자세한 설명은 'Fourier Features Let Networks Learn High Frequency Features in Low Dimensional Domains'에서 볼 수 있다.

4.2. Hierarchical volume sampling

단순히 ray 위의 점 $N$개를 샘플링하는 것은 빈 공간이나 물체 뒤의 가려진 영역처럼 렌더링에 기여하지 않는 점들도 있으므로 비효율적이다. 본 논문은 hierarchical sampling을 통해 최종 렌더링에 기여하는 부분을 집중적으로 샘플링하여 렌더링의 효율성을 올렸다.

이를 위해 하나의 네트워크가 아닌 coarse, fine 총 두 개의 네트워크를 최적화한다. 첫음에는 일반적인 샘플링 방식을 통해 $N_c$ 개의 샘플을 얻은 후 모델을 통해 얻은 출력값으로 volume rendering을 수행하여 모델의 성능을 평가한다. $$ \hat{C}c(\bold{r})=\sum{i-1}^{N_c}w_ic_i, \quad w_i = T_i(1-exp(-\sigma_i\delta_i)). $$ Coarse 네트워크의 출력값을 얻은 후, 최종 렌더링에 기여하는 영역에서 추가로 $N_f$샘플링을 진행한다. 추가적인 샘플링은 $w_i$를 normalize한 후 ($\hat{w}i=\frac{w_i}{\sum{j=1}^{N_c}w_j}$) 각 샘플 구간별로 얻은 일정한 확률 밀도 함수에서 inverse transform sampling을 통해 얻는다. 이렇게 얻은 총 $N_c + N_f$ 개의 샘플링을 이용하여 최종 렌더링된 이미지를 얻게 된다. $$ \hat{C}f(\bold{r})=\sum{i=1}^{N_c+N_f}T_i(1-exp(-\sigma_i\delta_i))c_i, \quad where \enspace T_i = exp\bigg(-\sum_{j=1}^{i-1}\sigma_i\delta_i \bigg) $$

4.3. Loss function

Loss function은 간단하게 렌더링된 이미지와 실제 이미지 간의 차이를 이용한다. $$ \mathcal{L}=\sum_{r\in \mathcal{R}}\bigg[\bigg|\bigg|\hat{C}_c(\bold{r})-C(\bold{r})\bigg|\bigg|^2_2 + \bigg|\bigg|\hat{C}_f(\bold{r})-C(\bold{r})\bigg|\bigg|^2_2 \bigg] $$ 최종 렌더링은 fine network를 통해 얻은 값을 통해 수행되지만, coarse network를 통해 얻은 weight의 분포가 fine sampling에 사용되므로 coarse network 또한 사용하여 loss를 최소화하는 방향으로 모델을 최적화하게 된다.

5. Results

5.1. Synthetic dataset

5.2. Real-world dataset

Synthetic 데이터셋은 아무래도 blender와 같은 프로그램을 통해 물체의 정확한 3차원 위치를 알 수 있고, 세상 좌표계 및 카메라 좌표계 간의 변환 관계도 정확하게 알 수 있으므로, real-world 데이터셋보다 높은 정확도를 보여준다고 생각한다.

5.3. Ablation studies

positional encoding, view direction, hierarchical sampling, 더 적은 수의 이미지, positional encoding 시 frequency band, 등 다양한 요소를 비교한 결과 모든 요소를 고려했을 때 가장 정확한 결과를 내는 것을 확인할 수 있다.

DETR architecture

DETR의 모델 구조는 아래 그림처럼 정말 간단하다.

1. Backbone

우선 입력 이미지에 대해 CNN으로 구성된 backbone network 를 통해 더 낮은 해상도의 feature map을 획득한다. 여기서 backbone으로 ResNet을 사용한다. 이미지의 차원이 $H \times W \times 3$ 이라면 backbone network를 통과하여 $\frac{H}{32} \times \frac{W}{32} \times 2048$ 인 activation map을 얻도록 설정한다.

2. Encoder & Decoder

2.1. Encoder

Backbone network를 통과해서 얻은 activation map은 $1 \times 1$ 컨볼루션을 통해 채널의 차원 수를 줄여 새로운 activation map으로 만든다. Transformer의 encoder는 1차원의 고정된 사이즈의 벡터를 입력으로 받으므로 이 activation map을 flatten하여 위치 정보와 함께 입력으로 넣어주게 된다. Encoder의 구조는 ViT에서 소개했듯이 normalization layer, Multi-head self-attention layer, 그리고 feed-forward neural network layer로 구성되어 있다.

2.2. Decoder

Decoder의 입력으로는 객체 class 수보다 많은 N 개의 object query가 들어간다. Encoder의 입력과 마찬가지로 1차원의 고정된 사이즈의 벡터 N 개를 입력하게 되면 self-attention, 그리고 encoder의 출력과 decoder의 입력 간의 cross-attention을 수행하여 모델이 모든 객체들 간의 전체적인 관계를 알 수 있고, 이미지 전체를 사용할 수 있게 된다.

3. Prediction feed-forward networks (FFNs)

모델의 최종 출력은 MLP로 구성된 FFN을 통해 이뤄진다. 첫 번째 FFN은 입력 이미지에 대해 예측된 bounding box의 중점 과 너비, 그리고 폭을 예측하고, 두 번째 FFN은 softmax 함수를 통해 객체의 레이블을 출력한다. 그 어떤 객체도 인식되지 않았을 때를 위해 $\emptyset$ class를 추가한다.

4. Loss function

DETR은 auxiliary loss를 사용하여 모델의 파라미터를 최적화한다. Transformer의 decoder layer는 보통 6개로 구성되어 있다. 이렇게 deep한 구조는 학습이 어려울 수 있으므로, 각 decoder에서 출력한 값들에 대한 loss를 모두 계산하여 학습할 수 있도록 한다. 즉, 각 decoder layer에서 FFN이 class와 bounding box를 예측하면서 parameter를 공유한다.

4.1. Object detection set prediction loss

정답들로 구성된 집합을 $y$, N 개의 예측으로 구성된 집합을 $\hat{y}$이라 하면, 두 집합 사이의 bipartite matching을 찾기 위해 N 개의 순열을 탐색한다. $$ \hat{\sigma} = \argmin_{\sigma \in \mathcal{G}N} \sum{i=1}^{N}{\mathcal{L}{match}(y_i,\hat{y}{\sigma(i)})} $$ 최적의 match를 찾기 위해 본 논문은 Hungarian algorithm을 통해 효율적으로 계산한다.

4.2. Hungarian loss

Hungarian algorithm을 통해 얻은 최적의 match를 이용하여 loss 값을 계산하게 되는데, 수식은 아래와 같다. $$ \mathcal{L}{Hungarian}(y,\hat{y})=\sum{i=1}^{N}{-c_i\log\hat{p}{\hat{\sigma}(i)}} + 1{{c_i\neq\emptyset}}\mathcal{L}{box}(b_i,\hat{b}{\sigma(i)}) $$

4.2.1. Classification loss

$$ -c\log\hat{p}{\hat{\sigma}} $$ 여기서 $C$는 정답 클래스, $\hat{p}{\hat{\sigma}}$는 예측된 클래스의 확률값이다. 여기서 정답 클래스의 확률을 1, 즉 100%로 예측하면 $\log(1)=0$이므로 loss 값은 0이 되고, 정답 클래스의 확률을 0.01, 즉 1%로 예측하면 $\log(0.01)=-4.6$이 된다. 이렇게 되면 잘못된 예측해 대한 loss 값이 정답에 대한 loss 값보다 작아지므로 $-$를 붙여주는 것이다.

4.2.2. Bounding box loss

$$ \mathcal{L}{box}(b,\hat{b}{\sigma}) = \lambda_{iou}\mathcal{L}{iou}(b,\hat{b})+\lambda{L_1}||b-\hat{b}||_1 $$ Bounding box loss는 L1 loss 와 generalized IoU (GIoU) loss를 선형결합한 loss를 사용한다.

• L1 loss L1 loss는 정답 bounding box의 꼭짓점 좌표 4개와 예측된 bounding box의 꼭짓점 좌표 4개의 차이를 계산하는 것이고, GIoU는 기존의 IoU에서 유용한 gradient를 제공하기 위해 사용된 것이다.

• GIoU (Generalized Intersection of Union) IoU는 정답 bounding box와 예측된 bounding box가 얼마냐 겹쳐있는지를 나타내는 metric이다.

  겹치는 부분이 클수록 IOU 값은 커질 것이고, 반대로 겹치는 부분이 작아질수록 IOU 값은 작아질 것이다. IOU의 최댓값은 1, 최솟값은 0이다. GIOU는 기존의 IOU가 겹치지 않는 영역을 0으로 loss 값을 계산하는 것이 학습 과정에서 유용한 gradient를 제공하지 않는다고 판단하여, 겹치는 영역이 없는 경우에도 차이를 반영할 수 있도록 확장한 개념이다. $$ GIoU = IoU - \frac{C-(A \cup B)}{|C|} $$ (C는 두 bounding box A, B를 포함하는 최소한의 bounding box) $\mathcal{L}_{iou}$ 앞에 있는 $1_{\{c\neq\emptyset\}}$는 정답 class와 일치하는 class에 대해서만 진행한다는 뜻이다. 즉, 정답 class와 일치한다면 해당 GIoU 값을 이용하고, 일치하지 않는다면 0을 곱하여 해당 GIoU 값을 무력화시킨다.

이 loss function을 이용해서 각 decoder layer에서 예측한 class와 bounding box를 정답 class와 bounding box와 비교하여 각 layer에서 얻은 loss 값을 모두 이용하여 모델의 parameter를 최적화하게 된다.

Transformer 구조는 자연어처리 task에서 성공적인 결과를 낸 반면에, 컴퓨터비전은 Convolutional Neural Network (CNN)에 의존하고 있다. 본 논문은 Transformer를 컴퓨터비전 분야에 적용하여 classification에서 이전 연구보다 더 좋은 결과를 낼 수 있는 Vision Transformer (ViT)를 제안한다.

1. Introduction

Transformer는 대규모의 text corpus를 통해 훈련을 한 후, fine-tuning을 통해 특정 task의 데이터셋에 대한 좋은 결과를 낼 수 있었다. 이는 Transformer의 computational efficiency와 scalability가 훌륭하다는 장점 덕분이다. 자연어처리에서의 성공에 영감을 받아, 본 논문은 Transformer 구조를 바로 이미지에 적용해본다. 이를 위해, 이미지를 여러개의 patch로 나누고 이 patch를 자연어처리에서의 token으로 다뤄 모델을 훈련시킨다.

2. Method

모델의 구조는 아래 그림과 같다.

2.1. Vision Transformer (ViT)

2.1.1. Image patching

Transformer는 1차원의 벡터를 입력으로 받는다. 입력값을 맞춰주기 위해 $H\times W\times 3$ 인 이미지를 $n \times n \times 3$ 인 패치로 분할한 후, flatten을 하여 1차원으로 만들게 된다. 또한 Transformer는 모든 layer에서 일정한 크기의 latent vector를 사용하므로, flatten한 패치를 학습 가능한 linear projection을 통해 일정한 크기의 latent vector로 mapping하게 된다.

2.1.2. CLS token

여기서 CLS는 Class가 아닌 Classification이다. 저자는 입력 패치들 앞에 학습 가능한 임베딩 값을 붙여 모델에 입력한다. 이 classification head는 pre-training, fine-tuning 과정 모두에 부착되어 사용되고, pre-training에서는 한 개의 hidden layer를 가진 MLP로 구현되고 fine-tuning에서는 한 개의 선형 layer로 구현된다.

2.1.3. Positional embedding

각 패치마다의 위치 정보를 유지하기 위해 1차원의 학습 가능한 position embedding을 더해준다.

2.1.4. Transformer encoder

Encoder는 임베딩된 패치들이 들어오면 첫 번째로 정규화를 진행한 후 Multi-Head Attention layer를 통과한 값과 정규화 전의 입력값을 residual connection 과정을 값을 더해주게 된다.

2.1.4.1. Layer normalization

한 layer를 통과할 때 분산이 바뀌면 gradient가 발산할 수 있으므로, layer의 output이 $N(0,1)$을 갖도록 정규화를 해준다.

2.1.4.2. Multi-Head Attention

입력 시퀀스의 입력값은 서로 여러 가지 방식으로 관련이 가능하다. 문장을 예로 들면 구문론적, 의미론적, 담화적 관계가 있을 수 있고, 음성을 예로 들면 주파수 대역별로 관계가 있을 수 있다. Transformer는 이 Multi-Head Attention layer를 통해 이를 해결한다. 여기서 query(Q), key(K), value(V) 개념이 등장하는데, Q는 현재 위치에서 다른 위치의 정보를 얼마나 가져올지, K는 Q와 얼마나 관련 있는지, V는 최종적으로 가져올 정보이다. 아래 그림에서 초록색을 K, 파란색을 Q, 빨간색을 V로 나타냈다. 계산 과정은 입력 벡터 간의, 즉 Q와 K 간의 dot-product를 계산한 후 softmax를 통해 attention weight를 얻어서 최종적으로 V와 곱해 최종 출력을 얻게 된다.

2.1.4.3. Residual Connection

ResNet에서 소개된 것으로, 네트워크가 깊어질수록 gradient vanishing 문제가 발생하는데, 이전 뉴런에서의 출력값을 더해줌으로서 이를 방지하면서 네트워크를 깊게 쌓을 수 있도록 해준다.

2.2. Inductive Bias

ViT는 CNN과 비교했을 때 귀납적 편향이 발생하여, 새로운 데이터에 대해 내리는 결과의 정확도가 떨어지게 된다. 그 이유는 CNN이 이미지의 local한 영역을 많이 보는 반면 Transformer는 global하게 보기 때문이다. 그래서 부족한 데이터셋을 통해 훈련을 하게 된다면, 당시에 이미지 분류에서 SOTA였던 ResNet과 비교하면 살짝 떨어지는 정확도를 보인다. 하지만, 자연어처리에서 처럼 대규모의 데이터셋을 통해 훈련을 한 후에 특정 데이터셋에 대한 정확도를 보면 높은 정확도를 자랑한다.