확산 모델의 반복 샘플링 오류를 학습형 보정기로 줄인다

확산 모델은 여러 번의 디노이징 단계를 거쳐 무작위 노이즈를 이미지로 바꾼다. 학습 과정에서 디노이저가 주로 받는 입력은 실제 이미지에 노이즈를 섞은 샘플이다. 반면 추론 과정에서는 다음 단계의 입력이 이전 단계에서 모델이 만든 예측에 영향을 받는다. 따라서 모델은 학습 때 직접 경험하지 못한 형태의 입력을 처리하게 된다. 초기 단계에서 발생한 작은 오류가 이후 단계로 전달되면 최종 결과에 흔적이 남거나 오류가 커질 수 있다.

‘Correction Training for Diffusion Models’는 이 문제를 해결하기 위해 전체 확산 모델을 다시 설계하는 대신, 기본 디노이저에 별도의 보정기를 추가한다. 제안 방식의 주요 과정은 다음과 같다.

• 기본 디노이저가 현재 시간 단계의 깨끗한 샘플을 예측한다.
• 보정기는 이 예측을 입력으로 받는다.
• 기본 모델과 동일한 실제 목표를 향하도록 보정기를 학습한다.
• 샘플링 중간에 예측을 수정하고, 추론 시에는 보정기 출력에 같은 과정을 반복 적용할 수도 있다.

핵심은 ‘기본 모델의 예측을 어떻게 고칠 것인가’를 별도의 학습 과제로 만든다는 점이다. 단순히 샘플러 설정을 바꾸거나 외부 가이던스를 강화하는 대신, 기본 디노이저가 실제로 만들어내는 예측을 학습 입력으로 사용한다. 이렇게 하면 모델이 학습 중 접한 입력과 실제 샘플링 경로에서 마주하는 상태 사이의 간극을 줄이는 데 도움이 될 수 있다. 중간 예측을 일찍 수정하면 초기에 생긴 아티팩트가 최종 이미지까지 이어지는 것도 완화할 수 있다.

논문은 DiT, SiT, JiT 아키텍처에서 이 방법을 평가했다. 저자들은 무가이던스 생성에서 기존에 보고된 FID 결과보다 크게 개선됐으며, 고품질 결과를 만들기 어려운 무조건 생성에서도 경쟁력 있는 성능을 보였다고 설명한다. 분류기 없는 가이던스는 이미지 품질을 높이는 데 유용하지만, 생성 다양성을 낮추거나 과도한 채도와 같은 문제를 일으킬 수 있다. 제안 방식은 가이던스를 단순히 강화하는 대신 모델 자체의 예측을 정제해 품질을 높이고, 다양성은 가능한 한 유지하려는 접근이다.

이 방법의 실용적인 특징은 모듈식 구조에 있다. 기존 디노이저가 전체 생성 과정을 계속 담당하고, 새로 추가된 보정기가 반복 샘플링 중 드러나는 예측 오류를 처리한다. 이는 확산 모델 개선이 반드시 더 큰 백본, 강한 가이던스, 복잡한 샘플러에만 의존할 필요는 없다는 가능성을 보여준다. 모델이 자신의 중간 출력을 수정하도록 학습하는 것도 하나의 설계 방향이 될 수 있다.

현재 제공된 자료만으로는 다양한 데이터셋과 해상도, 샘플링 예산에서의 성능이나 추가 계산 비용까지 판단하기 어렵다. 전체 논문과 후속 연구가 필요하지만, 반복 생성에서 발생하는 오류 누적을 직접 겨냥한 간결한 개선책이라는 점에서 의미가 있다.

출처: Hugging Face Daily Papers