[논문 리뷰] Attention for Image Registration (AiR): an unsupervised Transformer approach

Zihao Wang, Hervé Delingette|arXiv (Cornell University)|2021. 05. 05.

Advanced Neural Network Applications참고 문헌 17인용 수 4

한 줄 요약

이 논문은 이미지 레지스트레이션을 시퀀스-투-시퀀스 번역 작업으로 간주하는 비지도 학습 기반의 트랜스포머 기반 프레임워크인 AiR를 제안한다. 컨볼루션 백본 없이 다중 척도 자기주의 메커니즘을 활용함으로써 AiR는 기준 데이터셋에서 최신 기술 수준의 성능을 달성하였으며, 단일 척도 트랜스포머 및 CNN 기반 방법보다 변형 정확도와 구조 유사도 측면에서 뛰어난 성능을 보였다. 특히 큰 변형에 있어서 두드러진 성능을 보였다.

ABSTRACT

Image registration is a crucial task in signal processing, but it often encounters issues with stability and efficiency. Non-learning registration approaches rely on optimizing similarity metrics between fixed and moving images, which can be expensive in terms of time and space complexity. This problem can be exacerbated when the images are large or there are significant deformations between them. Recently, deep learning, specifically convolutional neural network (CNN)-based methods, have been explored as an effective solution to the weaknesses of non-learning approaches. To further advance learning approaches in image registration, we introduce an attention mechanism in the deformable image registration problem. Our proposed approach is based on a Transformer framework called AiR, which can be efficiently trained on GPGPU devices. We treat the image registration problem as a language translation task and use the Transformer to learn the deformation field. The method learns an unsupervised generated deformation map and is tested on two benchmark datasets. In summary, our approach shows promising effectiveness in addressing stability and efficiency issues in image registration tasks. The source code of AiR is available on Github.

연구 동기 및 목표

• 유사도 측도 최적화에 기반한 전통적인 비학습 기반 이미지 레지스트레이션 방법의 불안정성과 높은 계산 비용 문제를 해결하기 위해.
• CNN 기반 딥 러닝 방법이 장거리 의존성과 전반적인 구조적 변형을 포착하는 데에 한계가 있음을 극복하기 위해.
• 이전에 자연어 처리(NLP) 분야에서 우세했던 트랜스포머 아키텍처가 지도 학습이 필요 없는 종단 간 비지도 학습 기반 이미지 레지스트레이션에 어떻게 활용될 수 있는지 탐색하기 위해.
• 복잡한 변형을 더 잘 모델링하기 위해 트랜스포머 내에서 다중 척도의 자기주의 메커니즘을 개발하여 서로 다른 수준의 특징 표현을 향상시키기 위해.
• 자기주의 메커니즘이 약한 지도 학습 방식으로 변형 필드를 효과적으로 학습할 수 있으며, 주석이 없는 학습 데이터가 필요 없음을 입증하기 위해.

제안 방법

• 프레임워크는 고정 이미지와 이동 이미지의 패치를 입력 토큰으로 사용하여 이미지 레지스트레이션을 시퀀스-투-시퀀스 번역 작업으로 간주한다.
• 트랜스포머 인코더는 다중 헤드 자기주의를 사용하여 고정 이미지 패치를 처리하여 원본 이미지의 메모리 표현을 학습한다.
• 트랜스포머 디코더는 인코딩된 메모리 특징과 이동 이미지 패치에 대한 자기주의를 모두 사용하여 변형 필드를 예측한다.
• 예측된 변형 필드를 사용하여 이동 이미지를 미분 가능 공간 샘플링을 통해 왜곡하며, 이로 인해 재샘플링 연산을 거쳐 역전파가 가능해진다.
• 다양한 수신 필드 크기에서의 특징 표현을 향상시키기 위해 다중 척도 병렬 자기주의(MAPT) 모듈을 도입하였다.
• 모델은 이동 이미지와 고정 이미지 간 픽셀 단위 L2 거리 기반 복원 손실을 사용하여 비지도 방식으로 학습된다.

실험 결과

연구 질문

• RQ1 컨볼루션 특징 추출기 없이 순수한 트랜스포머 아키텍처가 이미지 레지스트레이션을 위한 변형 필드를 효과적으로 학습할 수 있는가?
• RQ2 단일 척도 트랜스포머와 다중 척도 자기주의 메커니즘 간의 성능 비교에서 복잡한 변형을 모델링할 때 어떤 것이 더 우수한가?
• RQ3 지도 학습이 필요 없는 비지도 트랜스포머 기반 방법이 진정한 변형 레이블 없이도 경쟁 가능한 정확도를 달성할 수 있는가?
• RQ4 트랜스포머의 자기주의 메커니즘이 CNN보다 이미지 레지스트레이션 작업에서 전반적인 구조적 관계를 더 잘 포착할 수 있는가?
• RQ5 다중 척도 특징 학습이 생성된 변형 지도의 품질과 이미지 복원 품질에 어떤 영향을 미치는가?

주요 결과

• 다중 척도 병렬 자기주의 트랜스포머(MAPT-AiR)는 MNIST 데이터셋에서 모든 지표에서 최고 성능을 기록하였으며, 평균 PSNR는 17.311 ± 2.522, 스무스 DICE 스코어는 0.827 ± 0.034를 기록하였다.
• 단일 척도 AiR는 CNN 기반 DIRNets보다 PSNR에서 뛰어난 성능(16.966 ± 2.897 vs. 12.941 ± 2.218)을 보였지만, 다중 척도 변종에 비해 성능이 열등하였다.
• MAPT-AiR는 '9', '7', '8'과 같은 큰 변형을 효과적으로 포착한 반면, 단일 AiR는 이러한 복잡한 형태를 정렬하지 못했다.
• MNIST에서 모델은 평균 MSE 0.027 ± 0.015를 기록하여 매우 높은 복원 정밀도와 최소한의 픽셀 수준 오차를 보였다.
• 강력한 구조적 정렬에도 불구하고, 모든 트랜스포머 기반 방법은 왜곡된 이미지에서 기하학적 세부 정보 결함과 노이즈를 보였으며, 이는 픽셀 수준 정밀도에 대한 한계를 시사한다.
• 결과적으로 트랜스포머가 전반적인 변형 모델링에 효과적임을 확인했지만, 미세한 무늬와 에지 보존 측면에선 여전히 CNN보다 뒤처지고 있음을 확인하였다.