Denoising Vision Transformers (DVT)

Positional Embedding에 의해 유발되는 noise artifact를 제거하는 네트워크

 

[Project Page]

[Github]

[arXiv](2024/01/05 version v1)

 

Abstract

ViT의 출력에서 나타나는 noise artifact를 분리하고 제거할 수 있는 Denoising Vision Transformers (DVT) 제안

 

 

 

Introduction

아래 그림은 원시 ViT 출력에 클러스터링 알고리즘을 적용하면 노이즈가 많은 클러스터가 생성된다는 것을 보여준다.

 

연구진은 3가지 이유로 위치 임베딩이 이러한 현상에 기여한다고 가정했다.

• Zero-tensor를 입력해도 유사한 노이즈 패턴이 발생한다.
• 위치 임베딩 없이 훈련한 모델에서는 노이즈 패턴이 발생하지 않는다.
• 입력 프레임의 context에 상당한 차이가 있음에도 artifact의 위치가 일정하게 유지된다.

 

 

Preliminaries

ViT의 forward pass:

 

 

Denoising Vision Transformers

• Factorizing ViT Outputs
• Per-image Denoising with Neural Fields
• Generalizable Denoiser

Factorizing ViT Outputs

이상적인 visual feature는 translation, reflection invariant 이어야 한다.

하지만 위치 임베딩이 이러한 불변성을 깨트리고 있다.

 

ViT의 출력은 입력 의존적인 semantic term, 공간 위치와 관련된 입력 독립적인 artifact term, 둘 사이의 상호 의존성을 설명하는 residual term으로 분해할 수 있다.

 

Per-image Denoising with Neural Fields

먼저 위에 제시한 ViT 출력 분해에서 semantic term을 F라 한다. F는 feature 그 자체를 의미한다.

다른 함수나 모델과 관련 없는 그냥 feature matrix이다.

F

또한 artifact term은 G.

G

 

학습 가능한 feature F_θ, G_ξ, function h_ψ에 대해서 y는 빨간 원, ŷ'는 파란 원이다.

ŷ은 y의 예측값이다.

coords()는 픽셀 좌표를 추출하는 함수, sg = stop_gradient

 

다음과 같은 손실을 최적화한다.

 

이상적인 결과는 ∆가 0이 되고 ViT(x)가 F와 G로 완벽히 분리되어 ViT(x) = F + G가 되는 것이다.

 

최적화는 2 stage로 구성되며, stage-1에서 먼저 L_distance만을 사용하여 F와 G를 최적화한 뒤 G를 동결하고 L_recon을 사용하여 F와 h를 훈련한다.

 

Generalizable Denoiser

이미지마다 매번 수동 최적화를 수행할 수는 없기에, denoising network D를 증류를 통해 학습시킨다.

 

구체적으로 D는 단일 transformer block으로 구현되며, ViT의 forward pass 후에 학습 가능한 위치 임베딩을 추가하고 D를 통해 처리한다.

 

 

 

Experiments