데이터 균형 스케줄링으로 최적화된 이미지 복원 파운데이션 모델 FoundIR v2

초록

FoundIR‑v2는 확산 기반 이미지 복원 파운데이션 모델에 데이터 균형 스케줄링과 MoE‑드리븐 스케줄러를 도입해, 다양한 복원 작업(디블러, 디헤이즈, 저조도, 초해상도 등)에서 데이터 혼합 비율을 동적으로 최적화한다. 이를 통해 50여 개의 서브태스크에 걸쳐 기존 최첨단 모델보다 일관된 성능 향상을 달성한다.

상세 분석

본 논문은 이미지 복원 파운데이션 모델의 성능이 단순히 데이터 규모와 모델 용량에만 의존하지 않고, 서로 다른 복원 작업 간 데이터 혼합 비율에 크게 좌우된다는 점을 실증한다. 이를 정량화하기 위해 ‘데이터 믹싱 법칙(Data Mixing Law)’을 도입하고, 작업별 성능 변화를 실시간으로 모니터링하면서 샘플링 비율을 조정하는 ‘데이터 균형 스케줄링(Data Equilibrium Scheduling)’ 알고리즘을 설계하였다. 구체적으로, 일정 간격(T)마다 별도 검증셋에서 각 작업의 PSNR/SSIM 변화를 측정하고, 성능이 감소한 작업에 대한 샘플링 가중치를 지수적으로 증가시켜 전체 데이터 분포를 동적으로 재조정한다. 이는 기존의 고정 비율 혹은 무작위 샘플링에 비해 과소 학습된 작업에 더 많은 학습 기회를 제공함으로써 다중 작업 간 성능 편차를 크게 감소시킨다.

모델 측면에서는 Stable Diffusion XL(SDXL) 백본에 ‘Mixture‑of‑Experts(MoE)‑드리븐 스케줄러’를 결합한다. 각 작업의 저품질(LQ) 특징과 확산 단계의 노이즈 라틴트 x_t를 결합한 z_t를 입력으로, n개의 공유 전문가 네트워크가 라우터에 의해 가중합된다. 라우터는 학습 가능한 게이팅 파라미터 g_i를 통해 작업 특성에 맞는 전문가를 선택하고, 이를 통해 작업별 최적화된 확산 사전(prior)을 제공한다. 이렇게 하면 저조도, 디블러, 디헤이즈, 초해상도 등 서로 다른 손상 형태와 레벨을 동시에 고려할 수 있다.

학습 파이프라인은 VAE 인코더로 LQ 이미지를 라틴트 공간에 매핑하고, SDXL 기반 확산 모델이 HQ 라틴트를 복원하도록 한다. 텍스트 설명(LLAVA)도 교차 주의 메커니즘을 통해 라틴트에 통합해, 텍스트‑이미지 정합성을 활용한다. 전체 학습은 96GB 메모리의 NVIDIA H20 GPU 클러스터에서 수행되며, 데이터 균형 스케줄링과 MoE 스케줄러가 공동 최적화된다.

실험에서는 4개의 대표 작업(디블러, 디헤이즈, 저조도, SR)뿐 아니라 50여 개의 실제 시나리오 서브태스크에서 기존 FoundIR, HYPIR, OmniL‑V 등과 비교해 PSNR/SSIM, LPIPS 등 다중 지표에서 일관된 우위를 보였다. 특히 데이터 비율을 동적으로 조정한 경우, 초기 학습 단계부터 전반적인 성능이 급격히 상승하고, 최종 수렴 시점에서도 특정 작업에 대한 과적합 현상이 최소화되었다.

한계점으로는 데이터 균형 스케줄링이 검증셋의 품질에 민감하며, 라우터의 전문가 수가 증가할 경우 메모리·연산 비용이 급증한다는 점을 언급한다. 향후 연구에서는 자동화된 검증셋 생성 및 경량화된 MoE 구조를 탐색할 필요가 있다.