ATD: Improved Transformer with Adaptive Token Dictionary for Image Restoration

이 논문은 학습 가능한 토큰 사전과 이를 활용한 교차 어텐션 메커니즘을 도입하여 선형 복잡도에서 전역 의존성을 모델링함으로써 이미지 복원 성능을 극대화하는 새로운 트랜스포머 아키텍처인 ATD 를 제안합니다.

핵심

이 논문은 **"이미지 복원 (Image Restoration)"**이라는 어려운 문제를 해결하기 위해 개발된 새로운 인공지능 모델, **ATD(Adaptive Token Dictionary)**에 대한 이야기입니다.

쉽게 말해, **"흐리거나 깨진 사진을 다시 선명하게 만드는 기술"**을 더 빠르고 똑똑하게 만든 연구입니다.

이 복잡한 기술을 일상적인 비유로 설명해 드리겠습니다.

---

1. 문제 상황: "창문 밖만 보는 경찰관"

기존의 최신 AI 모델들 (Transformer) 은 사진을 복원할 때 **'자기 자신 (Self-Attention)'**을 많이 봅니다.
하지만 이 모델들은 계산량이 너무 많아, **작은 창문 (Local Window)**을 통해 주변만 볼 수 있었습니다.

• 비유: imagine you are a police officer trying to solve a crime in a huge city. You are only allowed to look out of a small window in your office. You can see the people right next to your building, but you can't see the person who committed the crime on the other side of the city, even if they look exactly like the suspect you saw earlier.
• 한국의 비유: 마치 **작은 창문 (Local Window)**을 통해 주변 10 미터만 볼 수 있는 경찰관과 같습니다. 비록 멀리 떨어진 곳에 똑같은 옷을 입은 용의자가 있더라도, 창문 밖은 볼 수 없기 때문에 그 정보를 활용할 수 없습니다. 그래서 복잡한 구조 (예: 도시의 빌딩, 만화 속 선) 를 복원할 때 흐릿해지거나 잘못 그리는 문제가 생깁니다.

2. ATD 의 해결책: "전체 도시를 훑어보는 똑똑한 지도"

이 연구팀은 **"창문 밖을 보지 않아도, 전 세계의 정보를 활용할 수 있는 방법"**을 고안해냈습니다. 바로 ATD입니다.

핵심 아이디어 1: "외부 사전 (Token Dictionary)"을 만듭니다

기존 방식은 사진 속 정보만 썩어먹으며 고민했지만, ATD 는 **학습 과정에서 '이미지의 기본 구조'를 담은 외부 사전 (Dictionary)**을 만들어냅니다.

• 비유: 이 사전은 **"이미지의 알파벳"**이나 **"건축 자재의 샘플북"**과 같습니다.
  • "벽돌은 이런 모양이야", "나무 질감은 저렇게 생겼어", "만화 선은 이런 패턴이야"라는 정보를 미리 학습해 둔 거대한 책자입니다.
  • AI 는 흐린 사진을 볼 때, 이 책자를 펼쳐서 "아, 이 부분은 '벽돌' 패턴이랑 비슷하네!"라고 바로 찾아냅니다.

핵심 아이디어 2: "유사한 친구들끼리 모으기 (Category-based Attention)"

이제 AI 는 사진을 작은 창문으로 나누는 대신, 유사한 특징을 가진 부분끼리 모아서 처리합니다.

• 비유: 학교에서 **학급 (Window)**을 나누는 대신, **동호회 (Category)**를 만드는 것과 같습니다.
  • 창문 방식: "1 반은 1 층, 2 반은 2 층"처럼 위치만 보고 모읍니다. (멀리 떨어져 있어도 같은 동호회일 수 있음)
  • ATD 방식: "축구 좋아하는 애들", "그림 좋아하는 애들"처럼 **관심사 (특징)**가 비슷한 애들끼리 모읍니다.
  • 효과: 멀리 떨어져 있어도 '축구'를 좋아하는 친구끼리 모여 정보를 공유할 수 있습니다. 그래서 사진의 한쪽 구석에 있는 선과 다른 구석에 있는 선이 연결되어 선명하게 복원됩니다.

핵심 아이디어 3: "맞춤형 요리사 (Category-aware FFN)"

모인 친구들 (유사한 특징) 에게는 각자 맞는 정보를 제공합니다.

• 비유: 요리사가 손님이 "이건 매운 걸 좋아해, 이건 싱거운 걸 좋아해"라고 말하면, 그에 맞춰 요리를 다르게 해주는 것처럼, AI 도 이미지의 특징에 맞춰 정보를 더 정교하게 섞어줍니다.

---

3. 왜 이것이 대단한가요?

1. 전체적인 시야 (Global View) + 빠른 속도:

   • 기존 방식은 "전체 도시를 다 보려면 너무 느려서 창문만 볼 수밖에 없었다"면, ATD 는 **"전체 도시를 보면서도 속도는 빠르다"**는 기적을 이뤘습니다.
   • 비유: 전체 도시를 한눈에 보는 드론을 쓰되, 필요한 정보만 빠르게 필터링해서 처리하는 방식입니다.

2. 압도적인 성능:

   • 실험 결과, ATD 는 기존 최고의 모델들 (SwinIR, HAT, MambaIR 등) 보다 더 선명하고 디테일한 사진을 만들어냈습니다.
   • 특히 **만화 (Manga)**나 **복잡한 도시 풍경 (Urban)**처럼 선이 많고 구조가 복잡한 이미지를 복원할 때 빛을 발했습니다.

3. 다양한 작업 가능:

   • 단순히 사진 크기를 키우는 것 (초해상도) 뿐만 아니라, **노이즈 제거 (Denoising)**나 JPEG 압축 손상 복구 같은 다른 작업에서도 최고의 성능을 냈습니다.

4. 요약

이 논문은 **"작은 창문만 통해 세상을 보던 AI 에게, '외부 사전'이라는 거대한 지도를 주고, '유사한 것끼리 모으는 동호회'를 만들어 전 세계 정보를 빠르게 공유하게 했다"**는 내용입니다.

그 결과, 흐릿하고 깨진 사진이 선명하고 생생한 고화질 이미지로 다시 태어날 수 있게 되었습니다. 이는 앞으로 우리가 스마트폰이나 카메라로 찍은 모든 사진을 더 완벽하게 복원할 수 있는 길을 열어준 중요한 연구입니다.

기술 요약

논문 요약: ATD (Adaptive Token Dictionary)

1. 문제 제기 (Problem)

• 배경: 최근 이미지 복원 (초해상도, 노이즈 제거 등) 분야에서 Transformer 기반 아키텍처가 CNN 을 능가하는 성능을 보이고 있습니다. 이는 입력 이미지 전체의 장기적 의존성 (long-range dependencies) 을 모델링할 수 있는 자기 주의 (self-attention) 메커니즘 덕분입니다.
• 한계:
  • 계산 복잡도: 표준 자기 주의 메커니즘은 입력 크기에 대해 이차 (quadratic) 복잡도를 가지므로, 고해상도 이미지에 적용 시 계산 비용이 너무 큽니다.
  • 수용 영역의 제한: 이를 해결하기 위해 기존 방법들 (SwinIR 등) 은 '로컬 윈도우 (local window)' 내의 자기 주의만 계산합니다. 그러나 이는 수용 영역 (receptive field) 을 제한하여 이미지 전체에 분산된 유사한 구조 간의 상호작용을 포착하지 못하게 하고, 결과적으로 성능이 최적화되지 않는 문제가 발생합니다.
  • 전역적 상호작용의 부재: 윈도우 크기를 늘리면 성능은 향상되지만 계산 비용이 급증하여 실용성이 떨어집니다.

2. 제안된 방법론 (Methodology)

저자들은 **적응형 토큰 사전 (Adaptive Token Dictionary, ATD)**이라는 새로운 Transformer 아키텍처를 제안하여, 선형 (linear) 복잡도 유지하면서 전역적 의존성을 모델링하는 것을 목표로 합니다.

• 학습 가능한 토큰 사전 (Learnable Token Dictionary):

  • 전통적인 딕셔너리 학습 (Dictionary Learning) 에서 영감을 받아, 훈련 데이터에서 일반적인 이미지 구조 (외부 사전 지식) 를 요약하는 가변적인 토큰 사전 $D$를 학습합니다.
  • 이 사전은 모델이 입력 이미지의 특징을 외부 지식과 연결하는 데 사용됩니다.

• 토큰 사전 교차 주의 (Token Dictionary Cross-Attention, TDCA):

  • 입력 이미지 토큰과 학습된 토큰 사전 간의 교차 주의 메커니즘을 도입합니다.
  • 입력 토큰이 사전의 어떤 원자 (atom) 와 가장 유사한지 계산하여, 해당 사전 정보를 입력 특징에 주입합니다.
  • 스케일링 인자 재파라미터화: 사전 크기가 커질 때 주의 가중치가 희석되는 문제를 해결하기 위해, 사전 크기 $M$에 비례하는 로그 스케일링 인자 ($\tau' = 1 + \tau \ln(M)$) 를 도입하여 가장 관련성 높은 토큰에 대한 주의를 강화하고 희소성을 유지합니다.

• 적응형 카테고리 기반 자기 주의 (Adaptive Category-based Self-Attention, AC-MSA):

  • 기존 윈도우 기반의 공간적 분할 대신, TDCA 에서 생성된 주의 맵을 기반으로 토큰을 **카테고리 (Category)**로 그룹화합니다.
  • 동일한 사전 원자와 가장 유사한 토큰들을 하나의 그룹으로 묶어, 이미지 내 멀리 떨어져 있더라도 구조적으로 유사한 영역끼리 상호작용하게 합니다.
  • 각 카테고리 내부를 고정된 서브-카테고리 (sub-category) 크기로 나누어 병렬 처리를 가능하게 하므로, 전역적 상호작용을 유지하면서도 계산 복잡도를 선형으로 유지합니다.

• 카테고리 인식 피드포워드 네트워크 (Category-aware FFN, CFFN):

  • 각 토큰이 속한 카테고리 정보 (가장 유사한 사전 원자의 임베딩) 를 피드포워드 네트워크 (FFN) 에 주입합니다.
  • 이를 통해 FFN 이 지역적 특징 변환 시 카테고리별 특성을 고려하여 적응적으로 특징을 융합할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 Transformer 프레임워크: 외부 사전 지식 (학습된 토큰 사전) 을 활용하고, 콘텐츠 인식 카테고리 기반 분할 전략을 통해 선형 복잡도로 전역 자기 주의를 가능하게 하는 ATD 모델 제안.
2. 아키텍처 개선:
   • TDCA 의 스케일링 인자 재파라미터화를 통해 주의 가중치 희석 문제 해결.
   • 카테고리 정보를 FFN 에 통합하여 특징 융합 능력 향상.
3. 범용성 및 성능: ATD(초해상용) 와 ATD-U(다중 스케일 U-Net 구조, 노이즈 제거 및 JPEG 아티팩트 제거용) 를 개발하여 다양한 이미지 복원 작업에서 SOTA 성능 달성.

4. 실험 결과 (Results)

• 이미지 초해상도 (Image Super-Resolution):
  • Set5, Set14, BSD100, Urban100, Manga109 등 주요 벤치마크에서 EDSR, SwinIR, HAT, MambaIRv2 등 기존 SOTA 모델들을 능가하는 성능을 기록했습니다.
  • 특히 Urban100 과 Manga109 에서 HAT 대비 0.290.40 dB, MambaIRv2 대비 0.270.35 dB 의 PSNR 향상을 보였습니다.
  • 경량화 버전 (ATD-light) 또한 경량 모델들 중 최고의 성능을 기록했습니다.
• 이미지 노이즈 제거 및 JPEG 아티팩트 제거:
  • ATD-U 는 컬러/그레이스케일 노이즈 제거 및 JPEG 압축 아티팩트 제거 작업에서도 SwinIR, Restormer, SCUNet 등 최신 모델들보다 우수한 정량적 (PSNR/SSIM) 및 정성적 결과를 보였습니다.
• 계산 효율성:
  • 전역 의존성 모델링을 달성하면서도 GPU 메모리 사용량을 HAT 대비 약 30% 절감하고, MambaIRv2 대비 추론 속도를 25~50% 향상시켰습니다.

5. 의의 및 결론 (Significance)

• 효율적인 전역 모델링: 이 논문은 Transformer 의 핵심인 자기 주의 메커니즘의 계산 복잡도 병목 현상을 해결하면서도, 로컬 윈도우의 한계를 극복하는 새로운 패러다임을 제시했습니다.
• 외부 지식의 통합: 단순한 내부 자기 유사성 (self-similarity) 탐색을 넘어, 학습된 사전 (dictionary) 을 통해 외부 지식을 명시적으로 모델에 주입함으로써 복원 성능을 극대화했습니다.
• 실용성: 다양한 이미지 복원 작업 (초해상, 노이즈 제거, 아티팩트 제거) 에 적용 가능한 범용 아키텍처를 제공하며, 고해상도 이미지 처리에 있어 계산 효율성과 성능의 균형을 성공적으로 달성했습니다.

이 연구는 이미지 복원 분야에서 전역적 자기 주의 (global self-attention) 를 더 효과적이고 효율적으로 구현하는 길을 열었다는 점에서 중요한 의의를 가집니다.