Progressive Split Mamba: Effective State Space Modelling for Image Restoration

이 논문은 이미지 복원 시 국소 구조 보존과 장기적 공간 일관성을 동시에 달성하기 위해, 2D 토폴로지 왜곡과 장기 의존성 감쇠 문제를 해결하는 위계적 상태 공간 모델인 'Progressive Split-Mamba (PS-Mamba)'를 제안하고 다양한 복원 작업에서 기존 모델보다 우수한 성능을 입증합니다.

핵심

이 논문은 **"이미지 복원 **(사진을 깨끗하게 만드는 기술)에 대해 설명합니다.

기존의 AI 모델들은 사진을 고칠 때 두 가지 큰 고민이 있었습니다.

1. **세부 묘사 **(국소성) 작은 구석구석의 디테일을 잘 살리려면 넓은 시야가 필요하고, 넓은 시야를 가지려면 디테일이 흐려지는 모순이 있었습니다.
2. 계산의 비효율성: 전체 사진을 한 번에 분석하려면 컴퓨터가 너무 많은 일을 해야 했습니다.

최근 'Mamba'라는 새로운 AI 기술이 등장하며 이 문제를 해결할 희망이 생겼지만, Mamba 를 사진 (2 차원) 에 적용할 때 두 가지 치명적인 결함이 있었습니다.

이 논문은 이 결함을 해결한 **'PS-Mamba **(Progressive Split-Mamba)라는 새로운 방법을 제안합니다.

---

🧩 핵심 비유: "거대한 도서관의 책 정리법"

이 기술의 원리를 이해하기 위해 거대한 도서관을 상상해 보세요.

1. 기존 Mamba 의 문제점: "책장을 한 줄로 늘린 도서관"

기존 Mamba 는 도서관의 모든 책 (픽셀) 을 바닥에 일렬로 늘어뜨려서, 맨 앞의 책부터 맨 뒤의 책까지 한 줄로만 읽는 방식입니다.

• **문제 1 **(위치 왜곡) 원래 책장 옆에 있던 책 (인접한 픽셀) 이 일렬로 늘어지면 서로 멀리 떨어지게 됩니다. "이 책의 옆에는 어떤 책이 있었지?"를 기억하기가 매우 어려워집니다. (세부 묘사 손실)
• **문제 2 **(정보 소실) 도서관이 너무 길어지면, 맨 앞의 책에서 읽은 정보가 맨 뒤까지 전달되다 보면 중간에 사라져 버립니다. (긴 거리의 정보 연결 실패)

2. PS-Mamba 의 해결책: "작은 구역으로 나누어 정리하는 도서관"

이 논문은 "일렬로 늘어뜨리지 말고, **작은 구역 **(조각)을 제안합니다.

• **단계 1: 지능적인 분할 **(Progressive Split)
  도서관을 거대한 공간으로 보지 않고, **반 (Half) → 사분면 (Quadrant) → 팔분면 **(Octant) 순서로 작게 나눕니다.

  • 비유: 전체 도서관을 한 번에 훑어보는 대신, **작은 책상 **(패치)을 만들어 그 안에서만 책을 정리합니다.
  • 효과: 책상 안의 책들은 원래 위치를 유지하므로, "옆에 있는 책"을 기억하는 것이 매우 쉬워집니다. (국소적 세부 묘사 보존)

• **단계 2: 정보의 고속도로 **(Symmetric Cross-Scale Shortcut)
  작은 책상들끼리만 대화하면, 도서관 전체의 큰 흐름 (전체적인 분위기) 을 놓칠 수 있습니다.

  • 비유: 작은 책상들 사이로 **전체 도서관을 연결하는 '고속도로 **(건너뛰는 통로)를 만듭니다.
  • 효과: 먼 곳의 정보도 바로바로 전달되어, 책이 사라지지 않고 전체적인 구조가 무너지지 않습니다. (긴 거리의 정보 연결 강화)

---

🚀 이 기술이 가져온 변화

이 'PS-Mamba'를 사용하면 다음과 같은 놀라운 일이 일어납니다.

1. 선명한 디테일: 흐릿해진 사진의 머리카락, 나뭇잎, 글자 같은 미세한 부분도 뾰족하게 살아납니다. (기존 Mamba 가 놓쳤던 '옆에 있는 것'을 잘 기억하기 때문)
2. 자연스러운 전체감: 사진의 한쪽 끝에서 다른 쪽 끝까지 색감이나 구조가 자연스럽게 이어집니다. (정보的高速도로 덕분)
3. 빠른 속도: 전체를 한 번에 분석하는 대신 작은 조각으로 나누어 처리하므로, 컴퓨터가 훨씬 덜 피곤해지고 빠릅니다.

📊 실제 성과 (실험 결과)

연구진은 이 기술을 다양한 테스트에 적용했습니다.

• **초고해상도 **(Super-Resolution) 낮은 화질의 사진을 고화질로 바꿀 때, 기존 최고의 AI 들보다 더 선명하고 자연스러운 결과를 냈습니다.
• **노이즈 제거 **(Denoising) 사진 속의 잡티 (노이즈) 를 제거할 때, 사진의 본래 모습을 더 잘 복원했습니다.
• JPEG 압축 손상 복구: 압축되면서 깨진 사진의 모서리나 얼룩을 더 깔끔하게 고쳤습니다.

💡 결론: "가장 간단한 방법이 가장 효과적일 수 있다"

이 논문은 복잡한 기술을 더 복잡하게 만드는 것이 아니라, **사진의 자연스러운 구조 **(위치)를 유지하면서 Mamba 의 강력한 힘을 끌어내는 **'지혜로운 분할 전략'**을 제시했습니다.

마치 거대한 퍼즐을 한 번에 맞추려다 실패하는 대신, 작은 덩어리별로 맞춰가다가 마지막에 전체 그림을 연결하는 방식으로, 사진 복원 기술의 새로운 기준을 세웠다고 할 수 있습니다.

기술 요약

1. 문제 정의 (Problem)

이미지 복원 (Image Restoration) 작업은 노이즈, 흐림, 아티팩트 등 열화된 입력으로부터 고품질의 이미지를 복원하는 것으로, 국소적인 미세 구조 (fine-grained local structures) 와 장기적인 공간적 일관성 (long-range spatial coherence) 을 동시에 보존해야 하는 난제입니다. 기존 방법론들은 다음과 같은 한계를 가집니다:

• CNN (Convolutional Neural Networks): 국소적인 수용 영역 (receptive field) 에 제한되어 장기적인 의존성을 포착하기 어렵습니다.
• Transformer: 전역적 어텐션 (global attention) 을 통해 장기 의존성을 모델링하지만, 계산 복잡도가 $O(N^2)$으로 높아 고해상도 이미지에 비효율적이며, 국소적인 구조 세부 사항을 놓칠 수 있습니다.
• 기존 Mamba (State Space Models, SSM): 선형 시간 복잡도 ($O(N)$) 로 장기 의존성을 효율적으로 모델링하지만, 2D 이미지를 1D 시퀀스로 단순히 펼치는 (flattening) 방식은 두 가지 근본적인 문제를 야기합니다:
  1. 국소성 왜곡 (Locality Distortion): 2D 특징 맵을 1D 시퀀스로 변환하면 인접한 픽셀들이 시퀀스 상에서 멀리 떨어지게 되어, 국소적인 텍스처 복원이 어려워집니다.
  2. 장거리 감쇠 (Long-range Decay): SSM 의 안정성 제약으로 인해 정보가 시퀀스를 따라 전파될 때 지수적으로 감쇠하여, 먼 거리의 공간적 정보가 약화되고 전역적 일관성이 떨어집니다.

2. 제안 방법론 (Methodology: PS-Mamba)

저자들은 위 문제들을 해결하기 위해 **Progressive Split-Mamba (PS-Mamba)**를 제안합니다. 이는 2D 이미지의 위상적 구조를 보존하면서 SSM 의 효율성을 활용하는 계층적 프레임워크입니다.

핵심 구성 요소:

1. 점진적 분할 및 병합 (Progressive Split and Merge):

   • 전체 이미지를 1D 시퀀스로 펼치는 대신, **기하학적 일관성을 유지하는 패치 (patches)**로 분할합니다.
   • 특징 맵을 반 (halves), 4 분면 (quadrants), 8 분면 (octants) 등으로 점진적으로 분할하여 각 패치 내에서 독립적으로 Mamba 를 적용합니다.
   • 이를 통해 패치 내에서는 국소적인 인접성이 보존되고, 시퀀스 길이가 짧아져 상태 공간 모델의 감쇠 문제를 완화합니다.

2. 대칭적 크로스-스케일 스킵 연결 (Symmetric Cross-scale Skip Pathways):

   • 장거리 감쇠를 상쇄하기 위해, 계층 구조의 서로 다른 단계 (mirrored stages) 간에 대칭적인 스킵 연결을 도입합니다.
   • 이 경로는 저주파 전역 컨텍스트 (global context) 를 직접 전송하여 긴 의존성 거리를 단축하고, 정보 흐름의 안정성을 확보합니다.

3. 하이브리드 아키텍처:

   • 컨볼루션 전처리: 국소적인 연속성과 구조적 패턴을 강화하기 위해 Mamba 입력 전에 경량 컨볼루션 레이어를 적용합니다.
   • 어텐션 기반 퓨전: 컨볼루션 특징과 Mamba 특징을 콘텐츠 적응형 게이팅 (content-adaptive gating) 으로 융합하고, 채널 및 공간 어텐션 (Dual Attention) 을 통해 전역 일관성을 강화합니다.

3. 주요 기여 (Key Contributions)

1. 위상 인식 상태 공간 모듈: 2D 공간 구조를 유지하면서 선형 시간 복잡도를 가진 점진적 분할 (점, 4 분면, 8 분면) 을 통해 국소성 왜곡을 제거하고 안정적인 미세 텍스처 모델링을 가능하게 합니다.
2. 대칭적 크로스-스케일 스킵 링크: 선형 SSM 의 지수적 감쇠를 상쇄하기 위해 전역 구조 정보를 직접 전달하는 경로를 설계하여, 장기 의존성과 전역 일관성을 강화합니다.
3. 통합 복원 아키텍처: 상태 공간 모델링의 전역성, 컨볼루션의 국소성, 어텐션의 정제 능력을 통합하여, 토큰 재배열 (token reordering) 이나 다방향 스캔 없이도 이미지 복원 성능을 극대화합니다.

4. 실험 결과 (Results)

저자들은 초해상도 (Super-Resolution, SR), 이미지 노이즈 제거 (Denoising), JPEG 아티팩트 감소 (JPEG CAR) 작업에서 PS-Mamba 를 평가했습니다.

• 초해상도 (SR):
  • Lightweight 모델: Set5, Set14, BSDS100, Urban100, Manga109 등 모든 벤치마크에서 기존 Mamba 기반 (MambaIR, MambaIRv2) 및 어텐션 기반 (SwinIR-light 등) 모델들을 일관되게 능가했습니다. 특히 ×2, ×3, ×4 스케일 모두에서 SOTA(최고 성능) 를 기록했습니다.
  • Classic 모델: MambaIRv2-Large(34.2M 파라미터) 보다 적은 파라미터 (21.2M) 로 Urban100 에서 0.20dB, Manga109 에서 0.35dB 이상의 PSNR 향상을 달성하며, 더 적은 자원으로 더 높은 성능을 보였습니다.
• 노이즈 제거 및 JPEG 아티팩트 감소:
  • Gaussian 노이즈 제거 (σ=15) 와 JPEG 아티팩트 감소 작업에서도 Restormer, SwinIR, MambaIRv2 등 기존 SOTA 모델들을 능가하는 성능을 보였습니다.
• 효율성:
  • 복잡한 토큰 조작 없이도 선형 복잡도를 유지하면서, 국소성과 전역성 사이의 최적 균형을 이루어 계산 비용 대비 높은 복원 정확도를 달성했습니다.

5. 의의 및 결론 (Significance)

이 논문은 State Space Models (SSM) 을 2D 이미지 복원에 적용할 때 발생하는 국소성 왜곡과 장거리 감쇠라는 두 가지 근본적인 한계를 체계적으로 해결했습니다.

• 기술적 혁신: 이미지를 단순히 1D 로 펼치는 대신, 위상적 구조를 보존하는 '점진적 분할' 전략을 도입하여 SSM 의 효율성을 유지하면서도 CNN 의 국소성 강점을 살렸습니다.
• 실용적 가치: 복잡한 토큰 재배열이나 다방향 스캔 없이도 고품질 복원이 가능하여, 향후 이미지 및 비디오 복원 작업에 대한 새로운 표준 아키텍처로 자리 잡을 수 있는 잠재력을 가집니다.
• 성능: 다양한 복원 작업에서 파라미터 수와 계산 비용을 줄이면서도 기존 최첨단 모델들을 압도하는 성능을 입증하여, 효율적이고 강력한 이미지 복원 솔루션임을 증명했습니다.