틀린 예측을 고치면서 그리는 법

Correction Training for Diffusion Models라는 논문을 이번에 살펴보겠습니다. arXiv 번호 2610.14136으로 공개된 이 작업은 diffusion 모델의 반복 샘플링에서 생기는 오차 누적을 매 디노이징 스텝에서 바로잡자고 제안합니다. 원문 링크는 제목 아래에 걸어 두었으니, 이 글에서는 그 아이디어가 왜 나왔고 어떻게 동작하며 어디까지 확인됐는지를 차례로 따라가 보겠습니다.

안녕하세요, 패트릭입니다. 저는 이미지 생성 모델을 쓸 때마다 초반 몇 스텝에서 생긴 어색함이 끝까지 남는 경험을 자주 했습니다. 손가락 모양이 처음부터 조금 이상하면, 나중 스텝에서 아무리 다듬어도 그 이상함이 완전히 지워지지 않는데요. 그렇다면 모델이 매 스텝마다 자기 예측을 스스로 한 번 더 고칠 수 있다면 어떨까요?

처음에는 작은 덧셈처럼 보였습니다

제가 이 논문에서 가장 먼저 멈춘 부분은 제안의 크기보다 방향이었습니다. 보통 생성 품질을 올린다고 하면 더 큰 백본을 쓰거나, 가이던스를 세게 걸거나, 샘플링 스텝을 늘리는 쪽을 떠올리죠. 그런데 이 작업은 모델을 키우는 대신 매 스텝에서 clean sample 예측을 한 번 더 다듬는 작은 모듈을 붙였습니다. 얼핏 보면 사소한 후처리처럼 보이는데, 실제로는 학습과 추론 사이의 어긋남을 정면으로 다루는 선택이었습니다.

왜 이런 선택이 나왔을까요. diffusion 모델은 학습할 때 진짜 이미지에 노이즈를 섞은 입력을 보고, 추론할 때는 자기 예측으로 만든 입력을 봅니다. 입력의 출처가 달라지면 모델이 보는 분포 자체가 달라지고, 초반에 생긴 작은 오차가 뒤 스텝까지 이어질 수 있습니다. 저는 이 대목에서 고개를 끄덕였습니다. 테스트할 때만 다른 문제를 푸는 학생과 비슷해 보였기 때문입니다. 연습 문제에서는 선생님이 준 깨끗한 식을 보다가, 시험에서는 자기가 쓴 중간 풀이를 보고 다음 줄을 이어 쓰는 셈이죠.

이런 어긋남을 exposure bias라고 부르는데, 이름보다 장면이 중요합니다. 새벽에 급하게 뽑은 시안에서 얼굴 윤곽이 첫 번째 스텝부터 살짝 무너지면, 그 뒤 스텝은 무너진 윤곽 위에서 디테일을 얹게 됩니다. 눈을 예쁘게 그려도 턱선이 이미 어긋나 있으면 전체가 어색해지죠. 그래서 저는 이 논문의 질문을 이렇게 읽었습니다. 스텝 수를 늘리기 전에, 매 스텝의 예측 자체를 고칠 수는 없을까.

학습할 때와 만들 때가 다른 그림을 봅니다

diffusion 모델의 학습 장면을 떠올리면, 진짜 이미지 x0에서 출발해 노이즈를 섞은 xt를 만들고, 모델은 여기서 다시 x0나 노이즈를 맞히는 연습을 합니다. 이 과정에서는 정답이 매번 또렷하게 주어지므로, 모델은 안정된 입력 위에서 예측을 배웁니다. 반면 실제 생성에서는 정답이 없습니다. 모델이 전 스텝에서 내놓은 예측으로 다음 입력을 만드는 구조라서, 예측이 조금만 어긋나도 그 어긋남이 입력에 그대로 들어갑니다.

그렇다면 오차는 어떻게 쌓일까요. 초기 스텝에서 생긴 흐릿한 윤곽이나 잘못된 색감이 다음 스텝의 출발점이 되면서, 모델은 틀어진 출발점 위에서 최선을 다하게 됩니다. 뒤늦게 고치려고 해도 이미 전체 구도가 굳어 있어서, 디테일만 덧칠하는 모양이 되기 쉽습니다. 제가 실제로 생성 로그를 볼 때 느끼는 답답함도 여기에서 옵니다. 중간 스텝 이미지를 보면 뭔가 이상한데, 끝까지 가봐야 그 이상함이 어디서 굳었는지 알 수 있으니까요.

여기서 재미있는 질문이 생깁니다. 입력 분포가 달라지는 문제를 모델 크기로만 덮을 수 있을까요. 큰 모델은 틀어진 입력에서도 잘 버티겠지만, 입력 자체가 틀어지는 구조는 그대로 남습니다. 그래서 이 논문은 구조를 그대로 둔 채 입력의 질을 매번 끌어올리는 쪽을 골랐습니다. 매 스텝에서 예측을 고치면, 다음 스텝이 출발하는 지점 자체가 조금씩 좋아지니까요.

고친 것은 노이즈보다 예측이었습니다

이 작업의 구분은 생각보다 또렷합니다. 노이즈를 더 잘 맞히는 쪽으로 가지 않았고, clean sample에 대한 예측을 직접 다듬는 쪽으로 갔습니다. diffusion 샘플링에서는 매 스텝마다 모델이 깨끗한 이미지에 대한 예측을 내놓고, 그 예측을 이용해 다음 xt를 만드는데요. 이 예측이 좋아지면 다음 입력도 좋아지고, 다음 입력이 좋아지면 그다음 예측도 좋아지는 선순환이 생깁니다.

저는 이 숫자보다 그 뒤의 구조가 더 중요하다고 봅니다. 샘플 품질 지표가 좋아졌다는 말만으로는 무엇이 달라졌는지 알기 어렵지만, 예측을 고치는 위치가 매 스텝의 분기점이라는 설명은 곧바로 그림이 그려지기 때문입니다. 분기점에서 방향을 조금씩 틀면 종착지가 크게 달라질 수 있죠. 산길에서 갈림길마다 몇 도씩만 바로잡아도 하산 지점이 완전히 바뀌는 것과 비슷합니다.

물론 이런 아이디어를 들으면 바로 다음 의문이 생깁니다. 예측을 고치는 모듈은 무엇을 보고 배우며, 정답은 어디에서 올까요. 이 논문의 답은 단출합니다. corrector라는 모듈이 베이스 모델의 clean sample 예측을 입력으로 받고, 베이스 디노이저와 같은 ground truth 타깃을 향해 학습합니다. 따로 새로운 정답을 만들지 않고, 이미 있던 정답을 같은 방식으로 한 번 더 맞히는 구조죠. 그래서 학습 목표가 복잡해지지 않고, 기존 파이프라인에 얹기도 쉬워집니다.

디노이저를 하나 더 붙이는 일과 뭐가 다를까요

이 대목에서는 조금 조심해서 읽을 필요가 있습니다. corrector를 둔다는 말을 들으면, 그냥 디노이저를 두 개 쌓은 것과 뭐가 다르냐고 물을 수 있으니까요. 실제로 파라미터만 보면 용량이 늘어난 것처럼 보일 수도 있습니다. 하지만 동작하는 위치와 보는 입력이 다릅니다. 두 번째 디노이저가 노이즈 낀 이미지를 다시 보는 구조라면, corrector는 베이스 모델이 내놓은 예측 자체를 보고 그 예측의 오류를 줄이는 데 집중합니다.

일종의 오답 노트처럼요. 시험지를 두 장 푸는 것이 아니라, 한 장 푼 답안지를 옆에 두고 어디가 틀렸는지 빨간 펜으로 고치는 과정에 가깝습니다. 빨간 펜은 문제를 새로 풀지 않고, 답안의 어색한 줄을 찾아 바로잡죠. corrector도 같은 역할을 맡습니다. 베이스 예측에 남아 있는 흐릿함이나 구조적 어긋남을 입력으로 받아, 정답에 더 가까운 예측으로 고쳐 내놓습니다.

그렇다면 corrector가 베이스 모델의 실수를 그대로 답습하면 어떻게 될까요. 이 질문은 학습 데이터 구성에서 풀립니다. corrector는 깨끗한 정답 대신 베이스 디노이저의 예측을 입력으로 받고 학습하므로, 실제 추론에서 마주칠 틀어진 예측의 분포를 미리 보게 됩니다. 연습할 때부터 틀린 답안지를 갖고 첨삭 연습을 하는 셈이라서, 막상 실전에서 어설픈 예측을 받아도 당황하지 않는 구조죠. 저는 이 대목이 얄팍한 꼼수보다 정직한 설계에 가깝다고 느꼈습니다. 모델이 실제로 볼 입력을 학습 때도 보게 만들었으니까요.

매 스텝마다 예측을 한 번 더 다듬습니다

이제 샘플링 장면을 머릿속으로 따라가 보겠습니다. 어떤 타임스텝 t에서 베이스 모델이 현재 xt를 보고 clean sample 예측을 내놓습니다. 보통은 이 예측을 그대로 이용해 다음 단계로 넘어가는데, 이 방법에서는 corrector가 그 예측을 받아 한 번 더 정제합니다. 정제된 예측으로 다음 xt를 만들면, 다음 스텝은 조금 더 깨끗한 출발선에서 시작합니다. 이런 다듬기가 전 구간에 걸쳐 반복되면서, 초반에 생긴 자잘한 인공물(artifact)이 뒤까지 굳어지는 흐름이 끊깁니다.

문제는 그다음입니다. 매 스텝마다 모듈을 한 번씩 더 돌리면 비용이 들 텐데, 그만한 가치가 있을까요. 논문이 보여주는 그림은 가치가 있다는 쪽입니다. 초반 스텝에서 윤곽이 바로잡히면 중반 이후 스텝은 디테일에 집중할 수 있고, 전체적으로 덜 흔들리는 궤적이 만들어지기 때문입니다. 제가 체감하는 생성 실패의 대부분도 초반 구도가 틀어진 경우라서, 이 설명은 실제 감각과 잘 맞았습니다. 첫 단추를 매 스텝마다 다시 채우는 셈이죠.

여기서 한 가지 실제 시스템 관점의 질문이 바로 생깁니다. 매번 고치면 느려지지 않나요. 이 질문에 대한 답은 뒤에서 증류(distillation) 이야기로 이어지는데, 그 이야기는 마지막에 다시 다루겠습니다. 우선은 원리만 기억하면 됩니다. 예측하고, 고치고, 그 고친 예측으로 다음 입력을 만드는 루프가 매 스텝에서 돌아간다는 것입니다.

자기 출력을 다시 넣어도 왜 무너지지 않을까요

제가 이 논문에서 가장 흥미롭게 본 부분은 반복 적용에 대한 대목입니다. corrector는 학습할 때 베이스 디노이저의 예측만 입력으로 보는데, 추론할 때는 자기 출력에 다시 적용해도 결과가 좋아졌다고 합니다. 보통은 학습 때 못 본 입력을 주면 성능이 흔들리기 마련인데, 여기서는 오히려 다듬을수록 나아졌다는 것이죠. 처음에는 이 주장이 조금 크게 들렸습니다.

그런데 분포를 생각하면 납득이 갑니다. corrector가 배우는 일은 틀어진 예측을 정답 쪽으로 당기는 일이고, 한 번 당겨진 예측은 여전히 같은 종류의 입력 공간에 머뭅니다. 베이스 예측과 자기 출력이 서로 다른 세계의 값이 아니라, 같은 오답 노트 위의 다른 줄처럼 이어져 있기 때문입니다. 그래서 한 번 고친 결과를 다시 넣어도 모듈이 낯설어하지 않고, 남은 오차를 마저 깎아낼 수 있습니다.

생각해보면 사람도 비슷하게 배웁니다. 첨삭을 몇 번 받은 글은 다음 첨삭에서 지적할 부분이 줄지만, 첨삭 자체가 무의미해지지는 않습니다. 남은 어색함을 마저 다듬으면 문장이 더 매끄러워지죠. (참고로 저는 이 반복 적용을 무한정 돌리면 계속 좋아진다는 뜻으로 읽지 않았습니다. 논문이 말하는 것은 학습 때 못 본 자기 출력에서도 쓸모 있게 동작한다는 점이고, 그 이상을 일반화하려면 별도 측정이 필요합니다.)

DiT와 SiT와 JiT에서 모두 통했습니다

구조 이야기가 특정 백본에만 묶여 있으면, 저는 조금 의심합니다. 논문 읽기 모드로 보면 예쁜 아이디어도막상 다른 구조에 옮기면 깨지는 경우가 많기 때문입니다. 그런데 이 작업은 DiT와 SiT, 그리고 JiT라는 서로 다른 아키텍처에서 평가됐고, unguided와 unconditional 생성에서 FID가 크게 좋아졌다고 보고합니다. 한 집에서만 통하는 처방이 아니라 여러 집에서 통하는 습관처럼 보였다는 뜻이죠.

이 대목은 숫자의 크기로 읽기보다 조건으로 읽는 편이 좋습니다. unguided와 unconditional은 모델이 외부 조건이나 세게 건 가이던스 없이 스스로 그림을 끌고 가야 하는 설정입니다. 이런 설정에서 좋아졌다는 말은, corrector가 특정 조건 텍스트에 기대지 않고 생성 궤적 자체를 안정화한다는 해석을 뒷받침합니다. 프롬프트가 자세할 때만 잘 그리는 모델과, 빈 캔버스에서도 구도를 잘 잡는 모델은 성격이 다르니까요.

저희가 새로운 AI 아키텍처를 볼 때도 benchmark 숫자만 보지는 않습니다. 실제로 무엇이 바뀌었고, 그 변화가 시스템에서 어떤 비용 구조를 만드는지를 함께 봅니다. 그런 눈으로 보면 이 결과는 이렇게 읽힙니다. 백본을 바꾸지 않고도 샘플 궤적의 분산이 줄었고, 그 효과가 구조를 가리지 않고 나타났다는 것입니다. 물론 FID 하나만으로 화질과 다양성과 충실도를 다 말할 수는 없습니다. 그래서 다음에 볼 가이던스 비교가 중요합니다.

가이던스 없이 가이던스만큼 좋아졌습니다

이미지 생성에서 classifier-free guidance(CFG)는 품질을 끌어올리는 단골 수단입니다. 조건에 맞는 방향으로 샘플을 밀어주면서 선명도를 챙기죠. 그런데 대가가 있습니다. 세게 밀수록 색이 과해지거나(oversaturation) 표정이 굳고, 다양한 결과가 잘 안 나오는 쪽으로 쏠리기 쉽습니다. 다들 한 번쯤 본 그 번들거리는 질감이 바로 그 흔적입니다.

이 논문에서 눈에 띄는 보고는 correction만으로 CFG에 견줄 만한 성능이 나오면서도, 샘플 다양성이 대체로 유지됐다는 점입니다. 저는 이 문장을 처음에 흘려 읽었다가 다시 돌아왔습니다. 다양성을 지키면서 품질을 올렸다는 말은, 모델이 한 가지 정답 쪽으로 몰아붙인 것이 아니라 궤적마다 어색한 부분을 각자 바로잡았다는 뜻처럼 들렸기 때문입니다. 전원을 올려서 밀어붙인 것이 아니라, 각자 틀린 줄을 고쳐준 셈이죠.

그렇다면 둘을 함께 쓰면 어떻게 될까요. 논문에 따르면 correction과 CFG를 같이 쓰면 다양성이 더 좋아지고 CFG 특유의 과포화가 줄었다고 합니다. 이 결과는 마음에 들지만, 여기까지 일반화하기에는 조건별 측정이 더 필요합니다. 다만 방향은 분명해 보입니다. 가이던스로 방향을 잡고, correction으로 궤적을 다듬으면, 세게 밀지 않아도 되는 여유가 생깁니다. 힘을 빼고도 선명해지는 그림이 좋은 그림이죠.

여기서 흔한 오해를 하나 짚고 넘어가겠습니다. correction이 CFG를 대체하는 만능 열쇠라는 식으로 읽히면 곤란합니다. 논문이 말하는 것은 unguided와 unconditional 같은 특정 조건에서 견줄 만한 수준에 도달했다는 점이고, 모든 조건부 생성에서 CFG를 버려도 된다는 뜻이 아닙니다. 조건이 복잡해지거나 텍스트 충실도가 중요해지면 가이던스의 역할이 다시 커질 수 있습니다. 저는 그래서 이 결과를 대체보다 선택지 확장으로 읽었습니다.

추가 비용은 증류로 돌려받습니다

앞에서 미뤄둔 질문으로 돌아올 차례입니다. 매 스텝마다 고치면 추론이 무거워지지 않나요. 솔직히 이 구조를 처음 보면 연산량이 걱정됩니다. 스텝마다 베이스 예측에 더해 corrector 순전이 추가되니까요. 연구 현장에서는 품질이 좋아도 서빙 비용이 늘면 도입이 망설여지는 경우가 많습니다.

이 논문이 고른 해법은 정제된 예측을 다시 베이스 모델로 증류하는 것입니다. correction으로 다듬어진 결과를 교사로 삼아 베이스 모델을 가르치면, 추론 때는 corrector 없이도 다듬어진 궤적에 가까운 출력을 낼 수 있습니다. 추가 추론 오버헤드 없이 효과를 누리는 구성이죠. 한 번은 천천히 제대로 고치고, 그 요령을 본체에 이식하는 셈입니다. 학원에서는 선생님이 옆에서 첨삭해주다가, 시험에서는 혼자서도 고칠 줄 알게 되는 과정과 비슷합니다.

실제 시스템 관점에서 보면 다음 질문이 바로 생깁니다. 증류하면 품질이 얼마나 유지될까요. 논문의 보고는 추가 비용 없이 효과를 거둘 수 있다는 쪽이지만, 저는 여기서도 조건을 붙여 읽었습니다. 증류는 교사 분포를 얼마나 잘 덮느냐에 따라 결과가 달라지고, 백본 크기와 데이터 구성에 따라서도 편차가 생기기 때문입니다. 그럼에도 방향 자체는 실무적입니다. 추론 그래프를 바꾸지 않고 학습 단계에서 궤적을 개선하는 선택지는, 서빙 라인을 건드리기 어려운 팀에게 특히 반갑습니다. 밤사이 배치로 돌리는 증류 한 번이 낮의 서빙 부담을 덜어주는 구조니까요.

그래서 다음에 볼 것은 고치는 습관입니다

이제 전체를 길게 보면 결국 습관의 문제로 돌아옵니다. 한 번 잘 그리는 모델과 매번 스스로 고치는 모델은 실패하는 모양이 다릅니다. 전자는 틀리면 끝까지 틀리고, 후자는 틀려도 중간에 방향을 틉니다. 저는 그래서 correction training을 점수표보다 운영 감각으로 평가했습니다. 초반 스텝의 작은 어긋남이 뒤로 굳어지는 사고를 줄여주는 장치이기 때문입니다.

제가 앞으로 보고 싶은 것은 이 습관이 어디까지 넓어지는지입니다. 텍스트 조건이 강하게 걸린 생성에서도 같은 안정화가 이어질지, 비디오처럼 스텝과 프레임이 함께 늘어나는 생성에서도 매번 고치는 루프가 버틸지, 그리고 증류 없이 쓸 때와 증류해서 쓸 때의 간극이 구조마다 어떻게 달라지는지가 궁금합니다. 논문이 보여준 DiT와 SiT, JiT 결과는 출발점으로 충분하지만, 여기서 멈추면 아쉬운 것도 사실이죠.

그렇다면 사람의 역할은 무엇일까요. 모델이 스스로 고치는 습관을 갖춰도, 무엇을 잘 그렸다고 할지는 여전히 사람이 정합니다. 저는 이 논문을 읽고 나서 생성 버튼을 누르는 손이 조금 가벼워졌습니다. 완벽한 첫 스텝을 기다리기보다, 고치면서 가는 과정을 믿어봐도 되겠다는 생각이 들었기 때문입니다. 다음에 어색한 시안이 나와도 바로 버리지 말고, 어디에서 굳었는지를 한 번 더 들여다봐야겠습니다.