텍스트 복원을 위한 통합 디퓨전 트랜스포머 UniT

초록

UniT는 디퓨전 트랜스포머(DiT), 비전‑언어 모델(VLM), 텍스트 스포팅 모듈(TSM)을 순환적으로 결합해 저품질 이미지의 손상된 텍스트를 고품질로 복원한다. VLM이 이미지에서 추출한 텍스트 가이드를 초기화하고, TSM이 디퓨전 과정의 각 단계에서 OCR 예측을 제공해 VLM이 스스로 교정하도록 만든다. 최종적으로 DiT가 이러한 텍스트 힌트를 활용해 세밀한 문자 구조를 복원하고, 텍스트 환각을 크게 감소시킨다. SA‑Text와 Real‑Text 벤치마크에서 최첨단 F1 점수를 달성하였다.

상세 분석

UniT는 텍스트‑인식 이미지 복원(TAIR)이라는 특수한 복원 문제를 해결하기 위해 세 가지 핵심 모듈을 설계하고, 이들을 디퓨전 기반의 반복적 denoising 루프에 통합하였다. 첫 번째 모듈인 Vision‑Language Model(VLM)은 사전 학습된 대규모 멀티모달 모델(Qwen2.5‑VL 7B)을 사용해 저해상도(LQ) 이미지에서 가능한 텍스트 정보를 추출한다. VLM은 풍부한 시각‑언어 사전지식을 바탕으로 복잡한 단어와 문맥을 추론할 수 있지만, 심한 손상이나 문맥과 무관한 텍스트에 대해서는 오류를 범한다. 이를 보완하기 위해 두 번째 모듈인 Text Spotting Module(TSM)은 디퓨전 모델의 중간 특징(denoising timestep마다 추출된 diffusion features)을 입력으로 받아 OCR 예측을 수행한다. TSM은 문자 수준의 로컬 정보를 정확히 포착해 간단한 숫자·문자열에 강점을 보이며, VLM이 제공한 초기 텍스트 가이드가 부정확할 경우 보정 신호를 제공한다. 세 번째이자 핵심인 Diffusion Transformer(DiT) 백본은 UNet 기반 디퓨전 모델보다 전역적인 self‑attention을 활용해 고해상도 텍스트 구조를 정밀하게 재구성한다. DiT는 VLM과 TSM이 제공한 텍스트 힌트를 conditioning vector로 받아, denoising 과정 전반에 걸쳐 텍스트 일관성을 유지한다. 특히, UniT는 “iterative guidance”라는 메커니즘을 도입해 매 timestep마다 TSM의 OCR 결과를 VLM에 피드백하고, VLM은 이를 바탕으로 텍스트 프롬프트를 업데이트한다. 이렇게 순환적인 교정 루프는 텍스트 환각을 억제하고, 점진적으로 텍스트 복원 정확도를 높인다. 실험에서는 SA‑Text lv1lv3, Real‑Text 등 다양한 손상 수준을 포함한 데이터셋에서 기존 최첨단 모델(T eReDiff, DiT4SR) 대비 E2E F1 점수가 47%p 상승했으며, 특히 GT 텍스트 프롬프트를 제공했을 때 DiT 기반 모델이 UNet 기반 모델보다 텍스트 재현 능력이 현저히 우수함을 확인했다. 또한, VLM과 TSM 각각의 성능을 별도 평가한 결과, Qwen2.5‑VL 7B가 LLaVA보다 작은 파라미터에도 불구하고 복잡한 단어 복원에서 더 높은 정확도를 보였으며, TSM은 단순 숫자·문자열에서 거의 완벽에 가까운 OCR 정확도를 달성했다. 이러한 보완적 특성을 활용한 UniT의 설계는 텍스트‑중심 복원 작업에서 기존의 “텍스트 프라미스 + 이미지 손실” 방식보다 더 견고하고 해석 가능한 결과를 제공한다.