💻 computer science

Rethinking Test Time Scaling for Flow-Matching Generative Models

이 논문은 흐름 매칭 생성 모델의 테스트 시간 확장성을 개선하기 위해 다양성 증진 메커니즘인 'Repel'과 노이즈 인식 보상 미세 조정 전략인 'NARF'를 결합한 DOG-Trim 파이프라인을 제안하며, 기존 방법 대비 동일한 계산 비용으로 약 두 배의 성능 향상을 달성함을 보여줍니다.

원저자: Qingtao Yu, Changlin Song, Minghao Sun, Zhengyang Yu, Vinay Kumar Verma, Soumya Roy, Sumit Negi, Hongdong Li, Dylan Campbell

게시일 2026-03-23

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Qingtao Yu, Changlin Song, Minghao Sun, Zhengyang Yu, Vinay Kumar Verma, Soumya Roy, Sumit Negi, Hongdong Li, Dylan Campbell

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

이 논문은 **"생성형 AI(이미지 만드는 AI) 가 더 멋진 그림을 그리게 하는 새로운 방법"**에 대한 연구입니다.

기존의 AI 는 한 번에 그림을 그렸는데, 이 논문은 **"시간과 계산 자원을 조금 더 써서, 여러 번 시도해 보고 가장 좋은 그림을 골라내는 과정"**을 어떻게 더 똑똑하게 만들 수 있는지 설명합니다. 특히 최근 주목받는 'Flow Matching(흐름 매칭)'이라는 새로운 AI 기술에 적용하기 위해 고안된 방법입니다.

이 내용을 일상적인 비유로 쉽게 풀어보겠습니다.

🎨 배경: AI 화가의 딜레마

상상해 보세요. 여러분은 AI 화가를 고용해서 "해변에서 노는 강아지" 그림을 그려달라고 주문했습니다.

기존 방식 (Best-of-N): AI 에게 "그림 6 개를 그려봐. 그중에서 가장 좋은 걸 골라줄게"라고 합니다. AI 는 6 개를 다 그릴 때까지 기다렸다가, 가장 좋은 것을 고릅니다. 하지만 6 개를 다 그리는 데 시간이 너무 오래 걸립니다.
기존의 다른 시도 (국소 탐색): "그림을 그리는 중간중간 멈춰서, '이게 잘 되고 있네?'라고 확인하고 방향을 살짝 바꿔보자"는 시도입니다. 하지만 최신 AI 기술 (Flow Matching) 은 **결정론적 (Deterministic)**이라서, 중간에 방향을 살짝 바꾸는 것이 매우 어렵고 비쌉니다. 마치 "이미 찍힌 사진의 중간 장면을 다시 찍으려다 보니, 카메라를 부수고 다시 찍어야 하는 상황"과 비슷합니다.

💡 이 논문이 제안한 해결책: "DOG-Trim" (개 다듬기)

저자들은 "중간에 방향을 바꾸는 건 비효율적이니까, 처음부터 다양한 그림을 여러 개 동시에 그리다가, 질이 떨어지는 것들을 중간에 잘라내자 (Pruning)"라고 제안합니다. 이를 DOG-Trim이라고 이름 붙였는데, 이는 Diversity(다양성), Order(정렬), Global(전체) 을 의미합니다.

이 방법은 크게 세 가지 핵심 아이디어로 이루어져 있습니다.

1. Repel (반발력): "너무 똑같은 그림은 싫어!"

문제: AI 가 같은 주문 ("해변의 강아지") 을 받으면, 6 개를 그려도 다 거의 똑같은 그림을 그립니다. 마치 동일한 복제인간들이 같은 포즈를 취하는 것과 같습니다. 이렇게 비슷비슷하면 좋은 그림을 찾을 확률이 낮아집니다.
해결책 (Repel): AI 가 그림을 그리는 과정에서, 서로 너무 비슷한 부분 (토큰) 이 생기면 **"이봐, 너랑 너무 닮았잖아! 좀 멀어지자!"**라고 밀어냅니다.
비유: 한 반에서 친구들이 모두 같은 자리에 앉으려 할 때, 선생님이 "너는 저쪽, 너는 저쪽"이라고 자리를 재배치하여 반 전체가 다양한 관점을 갖도록 만드는 것과 같습니다. 이렇게 하면 AI 는 훨씬 더 다양한 스타일의 그림을 동시에 그려낼 수 있습니다.

2. NARF (소음 감지 보상): "초보 그림도 제대로 평가해 줘!"

문제: 그림을 그리는 초기 단계는 아직 흐릿하고 잡음이 많습니다. 이때 AI 가 그린 "초보 스케치"를 평가하는 AI 심사위원 (Reward Model) 은 원래 완성된 그림만 보고 훈련되었습니다. 그래서 "아직 잡음이 많으니 점수를 낮게 매겨라"라고 잘못 판단하거나, 잡음 때문에 생긴 우연한 패턴을 "멋진 디테일"로 오해합니다.
해결책 (NARF): 심사위원에게 **"초보 스케치도 완성된 그림이 될 수 있다는 걸 미리 알려주는 훈련"**을 시킵니다. 즉, "이 흐릿한 그림이 나중에 이렇게 예쁜 그림이 될 거야"라고 미리 예측하도록 가르칩니다.
비유: 요리사가 요리를 시작할 때, 재료를 다듬는 단계 (초기) 에서 "이게 나중에 맛있는 요리가 될 거야"라고 미리 예측할 수 있도록 훈련시킨다면, 좋은 재료를 일찍 골라낼 수 있는 것과 같습니다.

3. Global Pruning (전체 가지치기): "가장 유망한 것만 남기기"

방식: 6 개나 10 개의 그림을 동시에 그리기 시작합니다. 그림이 30% 정도 그려졌을 때, "NARF"로 점수를 매겨서 점수가 낮은 것들은 즉시 폐기합니다. 살아남은 것들만 계속 그립니다.
효과: 모든 그림을 다 그릴 필요 없이, 유망한 것들만 집중적으로 그려서 시간과 에너지를 아끼면서도 최고의 결과물을 얻습니다.

🏆 결과: 왜 이것이 중요한가요?

이 연구는 기존 방법들보다 동일한 계산 비용 (시간/전력) 으로 약 2 배 더 좋은 성능을 보였습니다.

기존 방법: "중간중간 방향을 바꿔보려고 애쓰다가" 에너지를 다 써버림.
이 방법 (DOG-Trim): "다양한 그림을 많이 그려놓고 (Repel), 중간에 점수가 낮은 걸 잘라내고 (Pruning), 초보 그림도 잘 평가하게 훈련된 심사위원 (NARF) 을 써서" 최고의 그림을 찾아냄.

📝 한 줄 요약

**"AI 가 그림을 그릴 때, 너무 똑같은 그림만 그리는 것을 막고 (Repel), 초보 단계에서도 좋은 그림을 잘 골라내게 (NARF) 훈련시켜서, 같은 시간 안에 더 멋진 그림을 찾아내는 방법"**입니다.

이 기술은 AI 가 더 창의적이고 다양한 이미지를 만들어내는 데 큰 도움을 줄 것으로 기대됩니다.

1. 연구 배경 및 문제 제기 (Problem)

배경: 최근 텍스트-이미지 생성 모델 (Text-to-Image) 의 성능 향상은 주로 훈련 시 모델 용량과 데이터 크기를 늘리는 '스케일링'에 의존해 왔으나, 최근에는 추론 시 (Test-Time) 추가적인 계산 자원을 할당하여 성능을 높이는 테스트 타임 스케일링 (Test-Time Scaling, TTS) 이 주목받고 있습니다.
기존 방법의 한계: 기존 TTS 방법들은 확률적 확산 모델 (Stochastic Diffusion Models) 을 기반으로 설계되었습니다. 이들은 노이즈 경로를 국소적으로 탐색 (Local Search) 하거나 경로를 분기시키는 전략을 사용하여 좋은 샘플을 찾습니다.
Flow Matching 모델의 특수성: 최신의 고성능 모델 (예: Flux.1-dev 등) 은 Flow Matching 방식을 사용하며, 결정론적 ODE 솔버 (Deterministic ODE Solvers) 를 통해 샘플링합니다.
- 문제점 1 (국소 탐색의 비효율성): Flow Matching 모델은 결정론적이기 때문에 국소적 경로 분기가 어렵습니다. 이를 위해 SDE(확률 미분 방정식) 를 사용하거나 노이즈를 재주입하는 등의 우회 방법을 쓰면, ODE 에 비해 훨씬 많은 추론 단계가 필요하여 계산 비용이 급증하고 효율성이 떨어집니다.
- 문제점 2 (후보 다양성 부족): Flow Matching 모델은 대규모 데이터셋에서 훈련되어 데이터 매니폴드의 좁고 높은 밀도 영역으로 수렴하는 경향이 있어, 생성된 이미지들이 서로 매우 유사합니다. 이는 탐색 공간 (Search Space) 을 제한합니다.
- 문제점 3 (중간 단계 보상 모델의 편향): 중간 단계 (노이즈가 있는 상태) 의 이미지를 평가할 때, 깨끗한 이미지로만 훈련된 기존 보상 모델 (Reward Model) 은 고주파수 세부 사항 (cartoonish edges 등) 에 편향되어, 실제로는 더 좋은 품질을 가질 수 있는 후보들을 잘못 제거 (Pruning) 하는 문제가 발생합니다.

2. 제안된 방법론 (Methodology)

저자들은 국소 탐색 대신 전역 탐색 (Global Search) 전략을 채택하고, 이를 Flow Matching 모델에 최적화하기 위해 DOG-Trim 파이프라인을 제안합니다.

A. 전역 가지치기 전략 (Global Pruning Strategy)

초기 노이즈에서 $N$ 개의 후보를 병렬로 생성합니다.
중간 시간 단계 (Intermediate Timesteps) 에서 각 후보를 부분적으로 디노이즈 (Denoise) 하고, 이를 이미지 공간으로 디코딩하여 보상 점수를 계산합니다.
점수가 낮은 후보들은 제거하고, 살아남은 후보들만 다음 단계로 진행하여 최종적으로 가장 높은 보상을 가진 이미지를 선택합니다.
이는 제한된 계산 예산 내에서 더 많은 초기 후보를 탐색할 수 있게 합니다.

B. Repel (토큰 반발 메커니즘)

목적: 전역 탐색 시 생성되는 후보들의 다양성을 확보하기 위함.
원리: 병렬로 진행하는 여러 경로 (Trajectories) 간의 이미지 토큰 (Image Tokens) 이 서로 너무 유사해지지 않도록 강제로 밀어냅니다.
구현: Transformer 블록의 토큰 특징 벡터 간 코사인 유사도를 계산하여, 가장 유사한 다른 경로의 토큰을 식별한 후, 해당 토큰을 반대 방향으로 선형적으로 확장 (Linear Extrapolation) 시킵니다.
효과: 초기 디노이즈 단계에서 적용하여 생성 경로의 분산을 극대화하고, 유효한 탐색 공간을 확장합니다.

C. NARF (Noise-Aware Reward Fine-tuning)

목적: 중간 단계 (노이즈가 있는 상태) 의 이미지 평가 정확도를 높이기 위함.
문제: 기존 보상 모델은 깨끗한 이미지로 훈련되어 중간 단계 이미지의 고주파수 노이즈를 과도하게 평가하거나, 실제 품질과 다른 순위를 매깁니다.
해결: 자기 증류 (Self-Distillation) 기법을 사용합니다.
1. 사전 훈련된 Flow 모델로 다양한 노이즈 초기값에서 경로를 생성합니다.
2. 중간 단계의 디노이즈 추정치 (Euler estimates) 를 이미지로 디코딩합니다.
3. 해당 경로의 최종 (완전 디노이즈된) 이미지가 기존 보상 모델로부터 받은 점수를 '의사-정답 (Pseudo-Ground-Truth)'으로 사용합니다.
4. 중간 단계 이미지를 입력으로 받아 최종 점수를 예측하도록 보상 모델을 미세 조정 (Fine-tuning) 합니다.
학습 전략: 쉬운 것에서 어려운 것 (Curriculum Learning) 으로 학습하며, 깨끗한 이미지에서 시작해 점차 노이즈가 많은 단계로 학습 영역을 확장합니다.

3. 주요 기여 (Key Contributions)

Flow Matching 모델에 적합한 TTS 전략 제안: 국소 탐색이 비효율적인 Flow Matching 모델의 특성을 분석하고, 계산 비용을 전역 탐색과 가지치기에 집중하는 DOG-Trim 파이프라인을 제안했습니다.
다양성 증대 메커니즘 (Repel): 토큰 수준의 반발 메커니즘을 도입하여 병렬 생성 시 샘플의 다양성을 높이고, 탐색 공간의 효율성을 극대화했습니다.
노이즈 인식 보상 미세 조정 (NARF): 중간 단계의 노이즈가 있는 이미지를 정확히 평가할 수 있도록 보상 모델을 자기 증류 방식으로 적응시켰으며, 이는 초기 단계의 잘못된 가지치기를 방지합니다.
성능 입증: 기존 SDE 기반의 국소/하이브리드 탐색 방법들보다 동일한 계산 비용 하에 훨씬 우수한 성능을 보임을 실험적으로 증명했습니다.

4. 실험 결과 (Results)

벤치마크: Flux.1-dev 모델을 기반으로 HPSv2.1, PickScore, ImageReward 등 다양한 보상 모델을 사용하여 평가했습니다.
성능 비교:
- 제안된 방법 (DOG-Trim) 은 기존 최상위 방법 (Best-of-N, Greedy Search, SDE 기반 하이브리드 방법 등) 보다 약 2 배 더 큰 성능 향상을 기록했습니다.
- 특히, ImageReward 기준에서 기존 방법 대비 가장 큰 격차 (1.682 vs 1.385 등) 를 보였습니다.
- Ablation Study: 전역 가지치기만으로도 기존 SOTA 를 상회하며, Repel 과 NARF 가 추가될 때 성능이 더욱 향상됨을 확인했습니다.
다양성 및 정렬도:
- Repel 적용 시 생성된 이미지들의 LPIPS 거리 (다양성 지표) 가 유의미하게 증가했습니다.
- NARF 적용 시, 중간 단계의 보상 예측과 최종 단계의 실제 순위를 비교한 Kendall's $\tau$ 계수가 크게 향상되어, 초기 단계에서 우수한 후보를 잘못 제거하는 비율이 줄어듦을 확인했습니다.
계산 효율성: SDE 기반 방법들은 동일한 품질을 내기 위해 ODE 보다 훨씬 많은 단계를 필요로 하여 비효율적이었으며, 제안된 방법은 ODE 기반의 빠른 수렴 속도를 유지하면서 전역 탐색의 이점을 극대화했습니다.

5. 의의 및 결론 (Significance)

이 논문은 현대의 고품질 텍스트-이미지 생성 모델 (Flow Matching 기반) 에 적용 가능한 테스트 타임 스케일링의 새로운 패러다임을 제시합니다.

패러다임 전환: 확산 모델 (Diffusion) 에 적용되던 국소적 경로 탐색 전략이 Flow Matching 모델에서는 비효율적임을 지적하고, 전역 탐색 + 초기 단계 가지치기가 더 효과적임을 증명했습니다.
실용성: 추가적인 모델 학습 없이, 추론 시 계산 자원을 효율적으로 배분하여 생성 품질을 획기적으로 개선할 수 있는 방법을 제시했습니다.
확장성: 제안된 NARF 전략은 다양한 아키텍처 (SDXL, SD3 등) 에도 일반화되어 적용 가능함을 보였습니다.

결론적으로, DOG-Trim 은 제한된 계산 자원 하에서 Flow Matching 모델이 생성할 수 있는 이미지 품질의 한계를 넓히는 강력한 프레임워크입니다.