UNet-AF: An alias-free UNet for image restoration

이 논문은 이미지 복원 작업에서 기존 UNet 의 번역 불변성 (equivariance) 한계를 해결하기 위해 최첨단 번역 불변 계층을 선택하여 설계한 'UNet-AF'를 제안하고, 이를 통해 경쟁력 있는 성능 유지와 함께 측정된 불변성을 크게 향상시켰음을 입증합니다.

핵심

이 논문은 **'UNet-AF'**라는 새로운 인공지능 모델에 대해 이야기합니다. 이 모델은 흐릿하거나 노이즈가 낀 사진을 선명하게 만드는 '이미지 복원' 작업을 잘해내도록 설계되었습니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

1. 문제점: "사진을 옮기면 왜 달라질까?" (아일링 현상)

기존의 UNet 이라는 인공지능은 사진을 흐리게 만들거나 다시 선명하게 할 때 아주 잘 작동합니다. 하지만 숨겨진 치명적인 약점이 하나 있었습니다.

비유:
가상의 '사진 정리 로봇'이 있다고 상상해 보세요. 이 로봇은 책상 위에 놓인 사진을 보고 "이 사진은 흐리니까 선명하게 만들어줘"라고 명령을 받습니다.

• 기존 로봇 (기존 UNet): 책상 위 사진을 오른쪽으로 1cm 살짝 밀었을 때, 로봇이 만들어낸 결과물이 완전히 달라져 버립니다. 사진의 위치가 조금만 바뀌어도 로봇은 "아, 이건 완전히 다른 사진이네!"라고 착각해서 엉뚱한 결과를 내놓습니다.
• 원인: 이 로봇은 사진을 확대/축소하거나 잘라낼 때, 마치 저화질로 찍은 사진을 보는 것처럼 '계단 현상'이나 '흐림' (전문 용어로 아일링/Aliasing) 이 생깁니다. 이 왜곡 때문에 로봇이 사진의 위치를 정확히 인식하지 못해, 입력된 사진이 조금만 움직여도 출력 결과가 불안정해지는 것입니다.

2. 해결책: "아일링 없는 UNet-AF"

저자들은 이 문제를 해결하기 위해 **'아일링 (Aliasing) 이 없는 UNet-AF'**를 만들었습니다.

비유:
이 새로운 로봇은 고급 렌즈와 정밀한 필터를 장착했습니다.

• 정밀한 필터 (Anti-aliasing): 사진을 확대하거나 줄일 때 생기는 '계단 현상'을 미리 잡아주는 필터를 모든 단계에 넣었습니다. 마치 안경을 써서 흐릿한 경계선을 부드럽게 만들어주는 것과 같습니다.
• 부드러운 활성화 함수: 로봇이 판단할 때 사용하는 '뇌의 스위치'를 기존에 날카롭고 거친 것 (ReLU) 에서, 부드럽고 매끄러운 것 (Filtered GELU) 으로 바꿨습니다. 이렇게 하면 작은 변화에도 로봇이 덜 놀라고 안정적으로 작동합니다.
• 원형 패딩: 사진 가장자리를 다룰 때도, 사진을 잘라내지 않고 구슬처럼 둥글게 감싸서 처리합니다. (예: 사진 오른쪽 끝이 왼쪽 끝과 이어지는 방식).

3. 실험 결과: "안정적이고 똑똑해진 로봇"

연구진은 이 새로운 로봇을 기존 로봇들과 비교해 봤습니다.

• 위치 변화에 강한가?
  • 기존 로봇: 사진을 살짝만 움직여도 결과가 크게 달라지고, 심지어 adversarial(적대적) 공격에 취약했습니다.
  • UNet-AF: 사진을 아주 미세하게 (0.01 픽셀 단위) 움직여도 결과가 거의 변하지 않았습니다. 위치에 상관없이 일관된 결과를 내는 '진정한 번역 (Equivariance)' 능력을 갖췄습니다.
• 화질은 어때?
  • 놀랍게도, 이 정밀한 필터를 넣었음에도 화질 (PSNR, SSIM 등) 은 오히려 더 좋아졌습니다. 특히 훈련 과정이 훨씬 안정적이어서, 로봇이 학습할 때 흔들림이 없었습니다.
• 단점은?
  • 속도: 정밀한 필터를 계속 사용하다 보니, 기존 로봇보다 약 7 배 정도 느려졌습니다. (고성능 GPU 가 필요함).

4. 핵심 요약 (한 줄 정리)

"기존의 사진 복원 AI 는 사진을 조금만 움직여도 결과가 뒤틀리는 버그가 있었는데, UNet-AF 는 모든 단계에 '정밀 필터'를 달아 위치가 바뀌어도 똑똑하고 일관된 결과를 내도록 만들었습니다. 속도는 조금 느려졌지만, 결과물의 품질과 안정성은 훨씬 뛰어납니다."

이 기술은 의료 영상 (MRI, CT 스캔 등) 이나 고화질 사진 복원처럼 작은 변화에도 민감하게 반응해야 하는 분야에서 큰 도움이 될 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

• 배경: UNet 아키텍처는 이미지 복원, 분할, 확산 모델 등 다양한 분야에서 표준적으로 사용되고 있습니다. 많은 응용 분야에서 UNet 은 입력의 병진 이동 (translation) 에 대해 출력도 동일하게 이동하는 **병진 불변성 (Translation-Equivariance)**을 갖는다고 가정됩니다.
• 한계: 그러나 기존 UNet 은 풀링 (pooling), 업샘플링 (upsampling), 활성화 함수 (activation functions) 등 여러 층에서 발생하는 에일리어싱 (aliasing) 현상으로 인해 실제 병진 불변성이 저하됩니다.
  • 기존 연구들은 풀링/업샘플링에 안티-앨리어싱 필터를 적용하거나 적응형 샘플링을 제안했으나, 이는 주로 픽셀 단위 이동에 국한되거나 서브픽셀 (연속) 이동에 대한 완벽한 불변성을 보장하지 못했습니다.
  • ReLU 와 같은 비선형 활성화 함수의 불연속성 또한 에일리어싱의 원인이 됩니다.
• 목표: 연속적인 서브픽셀 병진 이동에 대해 완벽한 (또는 수치적으로 근사한) 불변성을 가지면서도 이미지 복원 성능을 유지하는 새로운 UNet 아키텍처 개발.

2. 방법론 (Methodology)

저자들은 기존 Ronneberger UNet 과 Jin et al. 의 UNet-Jin 아키텍처를 기반으로, 에일리어싱을 제거하기 위해 각 구성 요소를 상태 최첨단 (SOTA) 의 병진 불변 층으로 교체한 UNet-AF를 제안합니다.

주요 변경 사항 (에일리어싱 제거 전략):

1. 합성곱 (Convolutions): 패딩 (padding) 모드를 병진 이동의 특성에 맞게 조정합니다. 원형 이동 (circular translations) 에는 **원형 합성곱 (circular convolution)**을, 크롭 이동 (crop-translations) 에는 **제로 패딩 (zero-padded convolution)**을 사용하여 경계 오류를 방지합니다.
2. 풀링 (Pooling): 최대 풀링 (Max-pooling) 을 BlurPool로 대체합니다. FFT 를 이용한 주기적 sinc 안티-앨리어싱 필터를 적용하여 수치적 정밀도 내에서 에일리어싱을 완전히 제거합니다.
3. 업샘플링 (Upsampling): 단순한 제로 삽입 방식 대신, 고주파수 대역의 에일리어싱을 제거하는 **필터링된 업샘플링 (Filtered Upsampling)**을 적용합니다.
4. 활성화 함수 (Activation Functions): ReLU 를 **필터링된 GELU (Filtered GELU)**로 대체합니다.
   • 입력을 더 미세한 그리드로 리샘플링 $\rightarrow$ 기본 활성화 함수 적용 $\rightarrow$ 고주파수 제거 $\rightarrow$ 원래 그리드로 리샘플링 과정을 거칩니다.
   • GELU 를 선택한 이유는 부드러운 함수 특성이 고주파수 생성을 제한하고, 다항식 활성화 함수와 달리 수치적 오차를 증폭시키지 않기 때문입니다.
5. 정규화 (Normalization): 배치 정규화 (BatchNorm) 를 **에일리어싱 없는 레이어 정규화 (Alias-free LayerNorm)**로 교체합니다. 기존 레이어 정규화는 병진 불변성이 없으며, BatchNorm 은 훈련 안정성과 성능 면에서 우세하지만 제안된 Alias-free LayerNorm 이 더 나은 성능을 보입니다.
6. 잔차 연결 (Residual Connections): 전체적인 아키텍처에 잔차 연결을 포함하되, 성능과 훈련 안정성에 기여하는 경우에만 선택적으로 활성화합니다.

3. 주요 기여 (Key Contributions)

• 새로운 아키텍처 제안: SOTA 병진 불변 층들을 통합하여 에일리어싱이 없는 UNet-AF 를 최초로 제안했습니다.
• 종합적 평가: 이미지 복원 작업 (원형 디블러링, 유효 디블러링, 가우시안 노이즈 제거) 에서 기존 UNet(Ronneberger, Jin) 과의 성능 비교를 통해 제안된 아키텍처의 유효성을 입증했습니다.
• 심층적 애블레이션 연구: 각 구성 요소 (풀링, 활성화 함수, 정규화 등) 가 최종 모델의 병진 불변성에 미치는 영향을 정량적으로 분석하여, 모든 변경 사항이 필수적임을 입증했습니다.

4. 실험 결과 (Results)

실험은 DIV2K 데이터셋을 기반으로 수행되었으며, PSNR, SSIM, LPIPS(성능 지표) 와 EQUIV(병진 불변성 지표), STAB(훈련 안정성), FPS(계산 효율성) 를 측정했습니다.

• 병진 불변성 (Equivariance):
  • UNet-AF 는 기존 베이스라인에 비해 약 40dB 더 높은 병진 불변성 점수를 기록했습니다.
  • 특히 원형 이동 (circular translation) 과 서브픽셀 이동에 대해 매우 안정적이고 강건한 성능을 보였습니다 (그림 1 참조).
• 성능 (Performance):
  • 병진 불변성을 높이는 과정에서 이미지 복원 성능 (PSNR, SSIM) 도 기존 베이스라인보다 향상되었습니다.
  • 특히 필터링된 GELU 와 BlurPool 사용은 성능 저하 없이 불변성을 크게 개선했습니다.
• 훈련 안정성 (Training Stability):
  • UNet-AF 는 훈련 중 검증 PSNR 변동이 훨씬 적어 (Fig. 4), 더 안정적이고 부드러운 수렴 경향을 보였습니다.
• 계산 비용 (Computational Cost):
  • 추가 필터링과 연산으로 인해 추론 속도 (FPS) 는 기존 UNet-Jin 대비 약 7 배 느려졌습니다. 이는 병진 불변성 확보를 위한 트레이드오프입니다.
• 애블레이션 연구 결과:
  • 어떤 층이라도 에일리어싱이 발생하는 대안 (예: 일반 MaxPool, ReLU, 일반 LayerNorm) 으로 교체하면 병진 불변성 점수가 급격히 하락 (약 40dB 감소) 하여, 제안된 모든 변경 사항이 필수적임을 확인했습니다.
  • 필터링된 다항식 활성화 함수는 수치 오차를 증폭시켜 성능과 불변성을 저하시켰습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 이론적/실용적 기여: UNet 과 같은 인코더 - 디코더 아키텍처에 연속적인 서브픽셀 병진 불변성을 성공적으로 적용한 첫 번째 사례 중 하나입니다.
• 성능과 불변성의 상관관계: 실험을 통해 병진 불변성이 높을수록 이미지 복원 성능과 훈련 안정성도 함께 향상됨을 관찰했습니다.
• 한계 및 향후 과제: 현재 높은 계산 비용이 단점이나, 특정 층에 최적화된 안티-앨리어싱 필터나 전용 GPU 커널을 활용하면 효율성을 개선할 수 있을 것으로 기대됩니다.

이 논문은 이미지 복원 분야에서 모델의 기하학적 일관성 (기하학적 대칭성) 을 확보하는 것이 단순히 이론적 우아함을 넘어, 실제 성능과 안정성 향상으로 이어질 수 있음을 강력하게 시사합니다.