Analyzing the Training Dynamics of Image Restoration Transformers: A Revisit to Layer Normalization

이 논문은 이미지 복원 트랜스포머의 학습 과정에서 기존 레이어 정규화 (LayerNorm) 가 특징의 발산과 채널 엔트로피 붕괴를 유발한다는 문제를 규명하고, 이를 해결하기 위해 공간 상관관계를 보존하고 입력별 통계를 반영하는 새로운 정규화 기법인 i-LN 을 제안하여 학습 안정성과 성능을 동시에 향상시켰음을 보여줍니다.

핵심

🍳 1. 문제 발견: "요리사가 재료를 너무 많이 섞어버렸다!"

최근 AI 는 흐릿하거나 깨진 사진을 고화질로 되돌리는 일을 잘합니다. 그런데 연구자들이 이 AI 의 내부 작동 원리를 자세히 들여다보니, 정말 이상한 현상이 발견되었습니다.

• 현상: AI 가 사진을 처리할 때, 내부에서 사용하는 '데이터의 크기 (Feature Magnitude)'가 100 만 배, 1000 만 배까지 불어나서 폭주했습니다. 마치 반죽을 만들 때 재료를 너무 많이 넣어서 그릇이 터질 듯이 부풀어 오른 것과 같습니다.
• 결과: 데이터가 너무 커지면 AI 는 중요한 디테일 (날카로운 가장자리, 미세한 질감) 을 잃어버리고, 오히려 엉뚱한 부분만 강조하게 됩니다. 마치 소금기를 잃어버린 요리가 맛이 없듯이, AI 가 만든 사진도 흐릿해지거나 아티팩트 (화면 찌그러짐) 가 생깁니다.

🔍 2. 원인 분석: "왜 폭주했을까?"

이 폭주의 원인은 AI 가 사용하는 **'LayerNorm (레이어 노멀라이제이션)'**이라는 도구 때문이었습니다. 이 도구의 역할은 AI 가 학습할 때 데이터를 일정하게 유지시켜 주는 것이지만, 이미지 복원 작업에는 두 가지 치명적인 단점이 있었습니다.

1. 개별 처리의 함정 (Per-token Normalization):

   • 비유: 한 그릇의 국을 만들 때, 국수 하나하나를 따로따로 맛을 보고 소금을 치는 것과 같습니다.
   • 문제: 이미지는 픽셀들이 서로 연결되어 있어야 합니다. 그런데 AI 가 각 픽셀을 따로따로 처리하다 보니, 이미지 속의 공간적 관계 (예: 눈과 코의 위치 관계) 가 깨져버렸습니다. 국수 하나하나가 맛있어도, 국 전체의 맛이 망가진 셈입니다.

2. 입력 무시 (Input-independent Scaling):

   • 비유: 비가 오는 날과 맑은 날, 모든 날에 똑같은 양의 소금을 뿌리는 것과 같습니다.
   • 문제: 이미지의 상태 (흐린지, 깨끗한지) 에 따라 필요한 처리가 다릅니다. 그런데 이 도구는 입력된 이미지의 특성을 무시하고 무조건 같은 방식으로 데이터를 정규화해버립니다. 중요한 정보 (원본 이미지의 통계적 특징) 를 버리게 되는 것입니다.

💡 3. 해결책: "i-LN (아이-엘엔)"이라는 새로운 조리법

저자들은 이 문제를 해결하기 위해 **'i-LN (Image Restoration Transformer Tailored Layer Normalization)'**이라는 새로운 도구를 만들었습니다. 기존 도구를 그대로 갈아끼울 수 있을 만큼 간단하지만 효과가 뛰어납니다.

• 전체적인 맛보기 (Spatial Holisticness):

  • 변화: 이제 국수 하나하나를 따로 보지 않고, 그릇 전체의 국물을 한 번에 맛봅니다.
  • 효과: 픽셀들 사이의 공간적 관계 (이미지의 구조) 를 그대로 유지하면서 데이터를 정리해줍니다. 눈과 코의 위치 관계를 망가뜨리지 않는 것입니다.

• 상황에 맞는 조절 (Input-adaptive Rescaling):

  • 변화: 비가 오면 소금을 적게, 맑으면 많이 뿌리듯이 입력된 이미지의 상태에 따라 데이터의 크기를 유연하게 조절합니다.
  • 효과: 중요한 원본 정보 (저수준 특징) 를 잃어버리지 않고, AI 가 필요할 때 자유롭게 활용할 수 있게 해줍니다.

🏆 4. 결과: 더 맛있고 정확한 요리

이 새로운 방법 (i-LN) 을 적용한 결과:

1. 폭주 멈춤: 데이터 크기가 100 만 배까지 불어나는 일이 사라졌습니다.
2. 더 선명한 사진: 흐릿한 사진이 선명해지고, 비가 낀 사진이 맑아지는 등 복원 성능이 크게 향상되었습니다.
3. 안정성: AI 가 학습하는 과정이 훨씬 안정적이 되어, 어떤 컴퓨터에서 실행하더라도 일관된 좋은 결과를 냅니다.
4. 저전력 장치에서도 가능: 데이터가 너무 커지지 않아서, 성능이 낮은 스마트폰 같은 작은 기기에서도 잘 작동합니다.

📝 요약

이 논문은 **"기존의 AI normalization(정규화) 방식은 이미지 복원 작업에 맞지 않아서 데이터를 폭주시키고 중요한 정보를 잃게 했다"**는 것을 발견했습니다.

그리고 **"이미지 전체를 한 번에 보고, 상황에 맞게 크기를 조절하는 새로운 방식 (i-LN)"**을 제안했습니다. 이는 마치 개별 재료만 보던 요리사가, 이제 그릇 전체의 균형을 보고 상황에 맞게 맛을 조절하는 요리사가 된 것과 같습니다. 그 결과, AI 가 만든 사진이 훨씬 더 선명하고 자연스러워졌습니다.

기술 요약

논문 요약: 이미지 복원 트랜스포머의 학습 동역학 분석: 레이어 정규화 (LayerNorm) 재검토

이 논문은 이미지 복원 (Image Restoration, IR) 작업에 널리 사용되는 트랜스포머 아키텍처의 학습 동역학을 심층적으로 분석하고, 기존 **레이어 정규화 (LayerNorm, LN)**가 가진 치명적이지만 간과된 문제를 제기합니다. 저자들은 기존 LN 이 특징의 크기를 수백만 배까지 발산시키고 채널별 엔트로피를 붕괴시킨다는 사실을 발견했으며, 이를 해결하기 위해 **이미지 복원 트랜스포머 전용 레이어 정규화 (i-LN)**를 제안합니다.

---

1. 문제 제기 (Problem)

기존의 이미지 복원 트랜스포머 (예: SwinIR, HAT 등) 는 Vision Transformer 의 표준인 Per-token LayerNorm을 사용합니다. 그러나 저자들은 IR 작업에서 이 방식이 다음과 같은 심각한 문제를 일으킨다고 분석했습니다.

• 특징 발산 (Feature Divergence): 학습 과정에서 내부 특징 (feature) 의 크기가 급격히 증가하여 수백만 (million) 스케일까지 발산합니다.
• 엔트로피 붕괴 (Entropy Collapse): 채널별 특징 엔트로피가 학습 초기에 급격히 감소합니다. 이는 특정 채널에서만 극단적인 값이 발생하여 통계적 분포가 왜곡됨을 의미합니다.
• 근본 원인: IR 작업은 저수준 (low-level) 의 공간적 상관관계와 입력에 의존적인 통계 정보를 보존해야 하는데, 기존 LN 의 제약이 이를 방해합니다. 네트워크는 LN 의 제약을 우회하기 위해 특징을 비정상적으로 증폭시키는 방향으로 학습하게 됩니다.

기존 LN 의 두 가지 주요 불일치 (Misalignments):

1. 토큰 단위 정규화 (Per-token Normalization): 각 토큰 (픽셀) 을 독립적으로 정규화하므로, 픽셀 간의 **공간적 상관관계 (Spatial Correlations)**가 파괴됩니다.
2. 입력 독립적 스케일링 (Input-independent Scaling): 모든 입력을 동일한 정규화 공간으로 매핑하므로, 입력 고유의 통계적 변동성 (Input-dependent statistics) 이 손실됩니다.

---

2. 제안 방법: i-LN (Methodology)

저자들은 위 문제를 해결하기 위해 **i-LN (Image Restoration Transformer Tailored Layer Normalization)**을 제안합니다. 이는 기존 LN 을 대체할 수 있는 간단한 'drop-in' 솔루션입니다.

i-LN 의 핵심 구성 요소:

1. 공간적 전체적 정규화 (Spatial Holistic Normalization, LN):*

   • 기존 LN 이 토큰 (픽셀) 단위로 평균과 분산을 계산하는 대신, 전체 공간 (Spatial) 과 채널 (Channel) 차원을 합쳐서 전역 평균과 분산을 계산합니다.
   • 효과: 토큰 간의 상대적 차이 (공간적 구조) 를 보존하며, 구조를 보존하는 변환 (Homothety) 을 가능하게 합니다.

2. 입력 적응적 리스케일링 (Input-adaptive Rescaling):

   • 정규화 과정에서 손실된 전역 스케일 (Global Scale) 정보를 복구하기 위해, Attention 과 FFN 레이어의 출력에 이전 정규화 과정의 표준편차를 곱하여 리스케일합니다.
   • 효과: 입력에 의존적인 통계 정보를 보존하고, 내부 표현의 범위 유연성 (Range Flexibility) 을 확보합니다.

수식적 표현:
기존 블록 $B(x)$ 대신 다음과 같이 적용됩니다:
$$B(x; f, \text{i-LN}) = x + \sqrt{\sigma^2 + \epsilon} \cdot f(\text{LN}^*(x))$$
여기서 $\text{LN}^*$는 전체 공간 - 채널 차원에서의 정규화이며, $\sqrt{\sigma^2 + \epsilon}$는 입력 적응적 리스케일링 인자입니다.

---

3. 주요 기여 및 이론적 통찰 (Key Contributions)

• 학습 동역학의 새로운 발견: IR 트랜스포머에서 특징 발산과 엔트로피 붕괴가 발생하며, 이는 네트워크가 LN 의 제약을 우회하려는 결과임을 규명했습니다.
• 구조 보존 이론: 기존 Per-token LN 은 토큰 간의 공간적 구조를 보존하지 못하지만 (Proposition 1), 제안한 공간적 전체적 정규화 (LN*) 는 전역 스케일링을 제외하고 구조를 보존함을 수학적으로 증명했습니다 (Proposition 2).
• 간단하고 효과적인 솔루션: 복잡한 아키텍처 변경 없이, 정규화 방식만 변경하여 학습 안정성과 성능을 동시에 개선했습니다.

---

4. 실험 결과 (Results)

다양한 이미지 복원 작업 (초해상도, 노이즈 제거, 비 제거, JPEG 아티팩트 제거) 에서 HAT, SwinIR, DRCT 등 여러 백본 모델을 사용하여 실험했습니다.

• 학습 안정성:
  • 기존 LN 은 학습 중 특징 크기가 $10^7$ 수준까지 발산하는 반면, i-LN 은 안정적인 범위 ($O(1)$) 를 유지합니다.
  • 채널 엔트로피가 i-LN 에서 잘 분포되어 유지되며, 학습 초기부터 안정화됩니다.
• 성능 향상:
  • 초해상도 (SR): Urban100, Manga109 등 다양한 벤치마크에서 PSNR/SSIM 지표가 기존 LN 대비 유의미하게 향상되었습니다 (예: HAT1 + i-LN 은 Urban100 에서 27.17 dB 달성).
  • 다양한 작업: 비 제거 (Deraining), 노이즈 제거 (Denoising), JPEG 아티팩트 제거 (CAR) 등 모든 IR 작업에서 일관된 성능 개선을 보였습니다.
• 저정밀도 추론 (Low-precision Inference) 강건성:
  • FP16 및 INT8 양자화: 기존 LN 은 정밀도 감소 시 특징 발산으로 인해 추론이 실패하거나 아티팩트가 발생하지만, i-LN 은 특징이 잘 제한되어 있어 정밀도 손실 없이 안정적으로 작동합니다.
• 공간적 상관관계 모델링:
  • i-LN 을 사용한 모델은 학습된 **상대적 위치 임베딩 (RPE)**이 더 구조화되어 있으며, 이는 네트워크가 픽셀 간의 공간적 관계를 더 잘 이해하고 있음을 시사합니다.

---

5. 의의 및 결론 (Significance)

이 논문은 이미지 복원 분야에서 널리 사용되던 표준 레이어 정규화 방식이 실제로는 작업의 고유한 요구사항 (저수준 특징 보존, 공간적 상관관계) 과 충돌하고 있음을 처음으로 체계적으로 증명했습니다.

• 실용적 가치: i-LN 은 추가적인 계산 비용 없이 기존 모델에 쉽게 적용할 수 있어, IR 트랜스포머의 성능을 즉시 향상시킬 수 있는 실용적인 방법론입니다.
• 이론적 기여: 정규화 기법이 네트워크의 학습 동역학과 특징 분포에 미치는 영향을 깊이 있게 분석하여, 향후 IR 및 저수준 비전 작업의 아키텍처 설계에 중요한 통찰을 제공합니다.
• 확장성: 경량화 환경 (저정밀도 추론) 과 복잡한 실제 세계 (Real-world) 데이터에서도 뛰어난 강건성을 보여, 실제 배포 가능한 모델 개발에 기여합니다.

결론적으로, i-LN은 이미지 복원 트랜스포머의 학습 불안정성을 해결하고 성능을 극대화하는 핵심 요소로, 해당 분야의 새로운 표준이 될 가능성이 높습니다.