[논문 리뷰] U2-Former: A Nested U-shaped Transformer for Image Restoration

Haobo Ji, Xin Feng외 3명|arXiv (Cornell University)|2021. 12. 04.

Image Enhancement Techniques인용 수 14

한 줄 요약

U2-Former는 다중 척도 특징 상호작용과 특징 필터링 메커니즘을 활용하여 계산 효율성을 향상시키는 깊고 중첩된 U자형 트랜스포머 아키텍처를 제안한다. 다중 시점 대비 학습을 통합하여 잡음과 청소화된 이미지 간의 분리 가능성을 높여 반사 제거, 비 제거, 안개 제거 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

While Transformer has achieved remarkable performance in various high-level vision tasks, it is still challenging to exploit the full potential of Transformer in image restoration. The crux lies in the limited depth of applying Transformer in the typical encoder-decoder framework for image restoration, resulting from heavy self-attention computation load and inefficient communications across different depth (scales) of layers. In this paper, we present a deep and effective Transformer-based network for image restoration, termed as U2-Former, which is able to employ Transformer as the core operation to perform image restoration in a deep encoding and decoding space. Specifically, it leverages the nested U-shaped structure to facilitate the interactions across different layers with different scales of feature maps. Furthermore, we optimize the computational efficiency for the basic Transformer block by introducing a feature-filtering mechanism to compress the token representation. Apart from the typical supervision ways for image restoration, our U2-Former also performs contrastive learning in multiple aspects to further decouple the noise component from the background image. Extensive experiments on various image restoration tasks, including reflection removal, rain streak removal and dehazing respectively, demonstrate the effectiveness of the proposed U2-Former.

연구 동기 및 목표

• 고비용의 계산 비용과 낮은 크로스 스케일 커뮤니케이션으로 인해 트랜스포머의 이미지 복원 적용 깊이가 제한된 문제를 해결한다.
• 표준 인코더-디코더 트랜스포머의 비효율성을 극복하기 위해 개선된 스케일 간 상호작용을 통해 더 깊은 특징 학습을 가능하게 한다.
• 패치 수준, 콘텐츠 수준, 잡음 수준의 표현 간 다중 시점 대비 학습을 도입하여 노이즈-배경 분리 성능을 향상시킨다.
• 낮은 품질의 토큰을 압축하는 특징 필터링 메커니즘을 통해 자기주의 어텐션 블록의 계산 효율성을 최적화한다.
• CNN 기반 사전 지식에 의존하지 않고도 반사 제거, 비 제거, 안개 제거 등 다양한 이미지 복원 작업에서 뛰어난 성능을 달성한다.

제안 방법

• 외부 인코더-디코더 프레임워크와 내부 U자형 트랜스포머 블록을 갖춘 중첩된 U자형 아키텍처를 사용하여 깊고 다중 척도의 특징 상호작용을 가능하게 한다.
• 자기주의 어텐션 계산 이전에 낮은 품질의 토큰을 제거하는 특징 필터링 메커니즘을 도입하여 계산 부담을 줄이면서도 핵심 특징을 유지한다.
• 세 가지 양성 쌍 구성 방식을 사용한 다중 시점 대비 학습을 구현한다: (1) 동일한 복원된 이미지에서의 패치, (2) 복원된 이미지와 진짜 청소화된 이미지, (3) 복원된 이미지와 무작위 진짜 이미지로 노이즈 불변 특징을 학습한다.
• 대비 학습과 함께 진짜 이미지에 대한 표준 감독을 사용하여 복원 및 특징 분리 최적화를 동시에 수행한다.
• 모델을 인코딩 및 디코딩 경로에 걸쳐 깊이 있는 트랜스포머 블록 스택으로 구성하여 스케일 간 전역적 맥락 모델링을 가능하게 한다.
• 모든 레이어에서 자기주의 어텐션 메커니즘을 활용하여 장거리 의존성을 캐치하고, 복원된 이미지의 공간 일관성을 향상시킨다.

실험 결과

연구 질문

• RQ1 중첩된 U자형을 가진 깊은 트랜스포머 기반 아키텍처가 이미지 복원에서 스케일 간 특징 통신을 개선할 수 있는가?
• RQ2 트랜스포머 블록에 특징 필터링 메커니즘이 도입될 경우 계산 비용은 크게 감소하지만 성능 유지 여부는 어떻게 되는가?
• RQ3 다중 시점 대비 학습이 이미지 복원에서 노이즈 성분과 청소화된 이미지 콘텐츠를 효과적으로 분리하는 데 기여하는가?
• RQ4 다양한 이미지 복원 작업에서 제안된 U2-Former의 PSNR 및 SSIM 측정치에서 최신 기술 수준의 방법들과의 비교는 어떻게 되는가?
• RQ5 표준 U-Net 또는 Uformer 스타일 아키텍처에 비해 중첩된 U자형 설계가 성능 향상에 얼마나 기여하는가?

주요 결과

• 실제 반사 제거 벤치마크에서 U2-Former는 Real20 세트에서 23.67 dB PSNR와 0.835 SSIM을 기록하여 Uformer(22.02 dB PSNR) 및 다른 최신 기술 수준 방법들을 능가했다.
• Rain100L에서의 비 제거 작업에서 U2-Former는 39.31 dB PSNR를 달성하여 Uformer(37.68 dB)보다 1.63 dB 향상되었으며, 시각적 결과에서 더 나은 세부 사항 복원 성능을 보였다.
• Rain100H 데이터셋에서 U2-Former는 30.87 dB PSNR를 기록하여 Uformer보다 1.50 dB 향상되었으며, 강한 비 조건에서도 뛰어난 성능을 보였다.
• 이미지 안개 제거 작업에서 U2-Former는 실내 세트에서 36.42 dB PSNR와 0.988 SSIM, 실외 세트에서 31.10 dB PSNR와 0.976 SSIM을 기록하여 Uformer보다 각각 4.51 dB와 4.58 dB 향상되었다.
• 시각적 결과에서 U2-Former는 색상 왜곡과 과도한 부드러움을 줄였고, DCP, DA 및 기타 기준 방법들보다 더 나은 세부 사항 보존 성능을 보였다.
• 제거 분석 결과 중첩된 U자형과 다중 시점 대비 학습이 성능 향상에 크게 기여했으며, 특징 필터링 메커니즘이 성능 저하 없이 추론 효율성을 향상시켰다.