DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

기술 요약: DIAMOND – Flow Matching 모델의 아티팩트 완화를 위한 지향적 추론 (Directed Inference for Artifact Mitigation in Flow Matching Models)

문제 정의
최근의 텍스트-이미지 생성 모델, 특히 Rectified Flow 아키텍처(예: FLUX.2) 및 확산 모델(Diffusion Models, 예: SDXL)을 기반으로 한 모델들은 높은 충실도를 보임에도 불구하고, 생성된 이미지에서 구조적 및 해부학적 아티팩트가 빈번하게 발생한다. 여기에는 왜곡된 손, 추가적인 팔다리, 기형적인 얼굴, 잘못된 텍text 등이 포함된다. 기존의 완화 전략들은 다음과 같은 상당한 한계를 가진다:

• 사후 정제 (Post-hoc refinement): 생성 후 이미지를 수정하는 방법들은 핵심 형성 과정 중에 개입하는 데 실패한다.
• 침습적 수정 (Invasive modifications): 모델 미세 조정(fine-tuning)이나 LoRA 가중치를 요구하는 기술은 계산 비용이 많이 들고, 확장하기 어려우며, 현대적 아키텍처에 대한 공개 가중치가 부족한 경우가 많다.
• 휴리스틱의 한계 (Heuristic limitations): 실시간 탐지 방법들은 종종 휴리스틱 임계값이나 무작위 잠재 변수 교체에 의존하여, 부정확한 위치 파악이나 최적화되지 않은 교정으로 이어진다.

방법론
저자들은 DIAMOND(Directed Inference for Artifact Mitigation in Flow Matching Models)를 제안한다. 이는 모델 가중치를 수정하지 않고 추론 궤적(trajectory) 중에 아티팩트를 수정하는 학습이 필요 없는(training-free), 제로샷(zero-shot) 프레임워크이다. 이 방법은 아티팩트가 발생하기 쉬운 잠재 상태로부터 생성을 유도하기 위해, 매 타임스텝마다 깨끗한 샘플의 추정치를 재구성한다는 원리에 따라 작동한다.

1. 깨끗한 잠재 변수 추정 (Clean Latent Estimation):
   노이즈가 섞인 잠재 변수($x_t$)에서 직접 작동하는 일반적인 접근 방식과 달리, DIAMOND는 각 단계에서 깨끗한 데이터 잠재 변수($\hat{x}_{0,t}$)를 추정한다.

   • Rectified Flow 모델(속도 $v_\theta$ 예측): $\hat{x}_{0,t} = x_t - t \cdot v_\theta(x_t, t)$.
   • Diffusion 모델(노이즈 $\epsilon_\theta$ 예측): $\hat{x}_{0,t} = x_t - \sigma_t \cdot \epsilon_\theta(x_t, t)$.
     이 추정된 잠재 변수는 픽셀 공간으로 디코딩($\hat{I}_t = D(\hat{x}_{0,t})$)되어, 아티팩트 탐지기가 중간 노이즈 상태가 아닌 최적의 도메인(깨끗한 데이터)에서 작동하도록 보장한다.

2. 아티팩트 탐지 및 손실 (Artifact Detection and Loss):
   미분 가능한 사전 학습된 아티팩트 탐지기(Artifact Detector, AD)(예: DiffDoctor의 탐지기)가 $\hat{I}_t$를 분석하여 아티팩트 확률 마스크 $M_t$를 생성한다. 아티팩트 손실($L_a$)은 이 마스크와 제로 마스크(아티팩트가 없는 이미지를 나타냄) 사이의 평균 제곱 오차(MSE)로 계산된다:
   $$L_a = \frac{1}{HW} \sum_{i,j} (M_{i,j}^t)^2$$

3. 그래디언트 기반 궤적 교정 (Gradient-Based Trajectory Correction):
   이 방법은 현재 잠재 변수 $x_t$에 대한 손실의 그래디언트($\nabla_{x_t} L_a$)를 계산한다. 교정의 크기가 아티팩트 영역의 크기에 의해 지배되어 안정성을 해치는 것을 방지하기 위해, 그래디언트는 **정규화(normalized)**된다:
   $$\delta_t = \lambda_t \frac{\nabla_{x_t} L_a}{\|\nabla_{x_t} L_a\|^2 + \epsilon}$$
   이 변위 벡터 $\delta_t$는 표준 솔버 업데이트 단계에서 차감되어, 궤적을 아티팩트로부터 멀어지도록 유도한다.

4. 동적 스케줄링 ($\lambda_t$):
   거듭제곱 법칙(power-law) 스케줄이 교정의 강도($\lambda_t$)를 제어한다. 구조적 아티팩트는 노이즈가 지배적인 초기 단계에서 나타나고 고주파 세부 사항은 나중에 정교해진다는 관찰에 근거하여, 이 스케줄은 생성 과정 초기에 강력한 교정을 적용하고 끝으로 갈수록 이를 감소시킨다.

5. 정체성 보존 (Identity Preservation, 선택 사항):
   교정이 전역적 이미지 구조나 프롬프트 정렬을 변경하는 것을 방지하기 위해, 선택적으로 정규화 항($L_{rec}$)을 적용할 수 있다. 이 항은 아티팩트가 없는 영역에서 교정된 이미지와 베이스라인 이미지 간의 평균 절대 오차(MAE)를 최소화한다.

주요 기여

• 학습이 필요 없는 프레임워크 (Training-Free Framework): DIAMOND는 추가적인 학습이나 가중치 수정 없이 Rectified Flow 및 Diffusion 모델의 추론 과정에 기하학적 궤적 가이드를 통합하는 최초의 제로샷 방법이다.
• 온더플라이 추정 (On-the-Fly Estimation): 이 방법은 추론 중에 깨끗한 잠재 변수($\hat{x}_{0,t}$)를 추정하는 새로운 메커니즘을 도입하여, 표준 탐지기가 일반적으로 실패하는 중간 단계에서도 신뢰할 수 있는 아티팩트 탐지를 가능하게 한다.
• 일반화 (Generalization): 이 접근 방식은 FLUX.1, FLUX.2와 같은 Rectified Flow 모델과 SDXL과 같은 표준 Diffusion 모델 모두에서 효과적으로 작동함을 입증하였다.
• 그래디언트 정규화 (Gradient Normalization): 논문은 업데이트 강도를 아티팩트의 공간적 크기와 분리하여, 아티팩트의 범위에 관계없이 일관된 교정을 보장하는 그래디언트 정규화를 식별하고 구현하였다.

실험 결과
FLUX.1 [dev], FLUX.1 [schnell], FLUX.2 [dev], SDXL과 같은 모델들을 사용하여 세 가지 데이터셋(동물, 단어, 사람)에 대해 정량적 평가를 수행하였다.

• 아티팩트 감소: FLUX.1 [dev]를 사용한 "단어(words)" 데이터셋에서, DIAMOND는 평균 아티팩트 빈도를 100%인 베이스라인에서 **9.85%**로 감소시켰으며, 이는 DiffDoctor(46.97%) 및 HPSv2(56.82%)와 같은 베이스라인을 크게 상회하는 성능이다.
• 픽셀 비율 (Pixel Ratio): 아티팩트 픽셀 비율(APR)은 베이스라인에 비해 거의 제로 수준(예: 단어 데이터셋에서 0.009%)으로 감소하였다.
• 품질 보존: 방법론은 높은 CLIP-T 점수(텍스트-이미지 정렬)와 ImageReward 점수(인간 선호도)를 유지하였으며, 이는 아티팩트 제거가 의미론적 일관성이나 시각적 품질을 저하시키지 않았음을 나타낸다.
• 정체성 보존: 절제 연구(Ablation studies)에 따르면, 정규화되지 않은 DIAMD는 아티팩트를 효과적으로 제거하는 반면, $L_{rec}$ 항($\alpha=0.1$)을 추가하면 아티팩트 억제와 베이스 모델의 정체성 보존 사이의 최적의 균형을 맞출 수 있음을 보여주었다.

의의 및 주장
저자들은 DIAMOND가 고충실도의 아티팩트 없는 이미지 합성을 위한 견고한 제로샷 경로를 제공한다고 주장한다. 샘플링 과정 중에 궤적 교정을 통해 직접 개입함으로써, 이 방법은 미세 조정의 계산 비용과 사후 휴리스틱의 부정확성을 극복한다. 본 논문은 이 접근 방식이 전역적 구조와 프롬프트 준수를 유지하면서도, 생성 과정 중에 미세한 공간적 감독과 국소적 아티팩트(예: 손 또는 텍스트)에 대한 정밀한 교정을 가능하게 한다는 점을 강조한다. 이 방법은 이커머스 및 마케팅 비주얼과 같이 해부학적 정확성이 중요한 전문적인 응용 분야를 위한 실용적인 솔루션으로 제시된다.