RestoRect: 흐름 기반 지식 증류로 구현한 초고속 이미지 복원

초록

본 논문은 교사 네트워크의 고품질 특징을 잠재 공간에서 학습 가능한 직선 경로(레티피드 플로우)로 재현하도록 학생 모델을 훈련시키는 새로운 지식 증류 프레임워크인 RestoRect을 제안한다. Retinex 기반 조명‑반사 분해, 학습 가능한 이방성 확산 제약, 삼각형 색공간 편극을 결합하고, 교사·학생 트랜스포머 특징을 교차 정규화·백분위수 기반 이상치 검출을 통해 정렬하는 FLEX 손실을 도입한다. 15개 데이터셋·4개 복원 과제·10개 지표에서 기존 최첨단 모델 대비 연산 효율과 복원 품질을 동시에 향상시킨다.

상세 분석

RestoRect은 기존 정적 특징 매칭 기반 증류가 트랜스포머의 동적 attention 흐름을 포착하지 못한다는 한계를 정확히 짚어낸다. 이를 해결하기 위해 ‘잠재 레티피드 플로우(Latent Rectified Flow)’라는 개념을 도입한다. 레티피드 플로우는 확산 모델에서 제시된 직선 샘플링 기법을 변형해, 노이즈와 목표 특징 사이를 선형 보간하는 ‘속도 필드’를 학습한다. 교사 네트워크가 추출한 다중 스케일 트랜스포머 특징을 목표로 삼고, 학생 네트워크는 동일한 속도 필드를 통해 동일한 잠재 궤적을 재현하도록 훈련된다. 이 과정은 기존 DDIM·DDPM 대비 샘플링 스텝을 크게 줄이면서도 고품질 특징을 복제한다는 점에서 효율성 측면에서 큰 장점을 가진다.

특징 정렬을 담당하는 FLEX 손실은 두 가지 핵심 아이디어를 결합한다. 첫째, 교사와 학생 특징을 각각 학생의 평균·분산으로 정규화한 뒤 교차 정규화(cross‑normalization)를 수행한다. 이는 트랜스포머 레이어마다 서로 다른 스케일과 분포를 갖는 특징들을 직접 비교 가능하게 만든다. 둘째, 백분위수 기반 이상치 검출을 통해 잡음이 심하거나 드물게 나타나는 활성화를 가중치 감소시켜, 손실이 과도하게 특정 영역에 집중되는 현상을 방지한다.

구조적으로는 Retinex 이론을 활용해 입력 이미지를 ‘반사(R)·조명(L)’ 두 성분으로 분해하고, 각각을 별도 ResNet 인코더로 처리한다. 이렇게 얻은 물리적 prior는 트랜스포머 U‑Net의 인코더‑디코더 흐름에 스킵 연결 형태로 주입돼, 조명 변화에 강인한 복원을 가능하게 한다. 또한, 학습 가능한 이방성 확산 연산(A(I)=∇·(c(|∇I|)∇I))을 도입해 에지 보존을 강화하고, 색공간 편극(HVI) 변환을 통해 전통적인 HSV/HSL의 붉은색 경계 불연속성을 삼각함수 기반 연속 매핑으로 해소한다.

정규화 측면에서는 기존 LayerNorm이 채널 차원에만 적용돼 공간적 상관관계를 놓치는 문제를 지적하고, Spatial Channel Layer Normalization(SCLN)을 제안한다. SCLN은 전체 공간‑채널 차원에 대한 평균·분산을 계산해 정규화함으로써, 전역 이미지 통계와 지역 패턴을 동시에 반영한다. 실험적으로 FP16·BF16 환경에서 오버헤드가 0.5% 이하에 그쳐, 정확도 향상 대비 비용이 거의 없음을 입증한다.

학습은 두 단계로 진행된다. 1단계에서는 교사 모델을 픽셀·퍼셉추얼·텍스처·색공간 손실(L_rec, L_tex, L_col)로 사전 학습한다. 2단계는 학생 모델 증류 단계이며, 첫 번째 서브‑페이즈에서는 교사와 학생의 특징을 재현하는 속도 필드만 학습하고, 두 번째 서브‑페이즈에서는 전체 복원 네트워크를 미세 조정한다. 이렇게 단계적 학습을 통해 학생 모델이 초기에는 교사의 특징 분포를 정확히 모방하고, 이후 실제 복원 작업에 적용될 수 있도록 한다.

실험 결과는 15개 데이터셋(저조도, 수중, 역광, 안저 등)에서 PSNR/SSIM·LPIPS·FID 등 10개 지표를 기준으로 기존 최첨단 Diffusion·Transformer 기반 모델 대비 평균 1.2dB 이상의 PSNR 향상과 3배 가량의 추론 속도 개선을 보고한다. 특히, 레티피드 플로우 기반 학생 모델은 DDIM 대비 5~7배 적은 샘플링 스텝으로 동등하거나 더 높은 FID를 달성한다는 점이 눈에 띈다.

비판적으로 보면, 레티피드 플로우와 FLEX 손실을 동시에 구현하려면 복잡한 파이프라인과 다중 모듈(Retinex 분해, anisotropic diffusion, HVI 변환, SCLN 등)이 필요하다. 재현성을 위해서는 각 모듈의 하이퍼파라미터와 초기화 전략이 상세히 공개되어야 하는데, 논문에서는 일부 파라미터(예: anisotropic diffusion 민감도 s, HVI 색상 파라미터 k)만 범위만 제시하고 구체적인 초기값·학습 스케줄을 생략했다. 또한, 교사·학생 네트워크가 서로 다른 아키텍처(ResNet‑Encoder vs. U‑Net‑Transformer)일 경우 교차 정규화가 과연 일반화 가능한지, 다른 도메인(예: 의료 영상 외)에서도 동일한 이득을 보장하는지는 추가 검증이 필요하다. 마지막으로, 레티피드 플로우 자체가 선형 보간에 기반하므로 비선형 복잡한 변환을 완전히 재현하기엔 한계가 있을 수 있다. 향후 연구에서는 비선형 궤적 학습이나 다중 단계 흐름을 결합해 이러한 제약을 보완할 여지가 있다.