검증자 없이 테스트 시 확장 가능한 확산 모델 VFScale

초록

VFScale는 확산 모델의 내재 에너지 함수를 검증자로 활용해 테스트 시 샘플 수를 늘리는 스케일링을 가능하게 한다. 훈련 단계에서는 MRNCL 손실과 KL 정규화를 도입해 에너지 지형을 정렬하고, 추론 단계에서는 초기에는 Best‑of‑N, 후반에는 MCTS를 결합한 하이브리드 Monte Carlo Tree Search(hMCTS)로 탐색 효율을 높인다. Maze와 Sudoku와 같은 복합 추론 과제에서 기존 확산 모델을 크게 능가한다.

상세 분석

VFScale 논문은 인간의 System 2 사고를 모방하고자 하는 확산 기반 추론 모델의 두 가지 핵심 한계를 명확히 짚는다. 첫 번째는 “검증자 의존성”이다. 기존 연구들은 외부 검증자(예: 사전 학습된 스코어링 네트워크)를 이용해 다수의 샘플 중 최적을 선택했지만, 이는 인간이 스스로 판단하는 내재적 추론과는 거리가 있다. 두 번째는 “탐색 효율성”이다. 단순히 샘플 수를 늘리는 Best‑of‑N(BoN) 방식은 계산량이 급증함에도 불구하고 성능 포화 현상이 나타난다.

이 문제들을 해결하기 위해 저자들은 두 축에서 혁신을 제시한다. 훈련 측면에서는 기존의 MSE 손실과 대비 손실(Contrastive loss)만으로는 에너지 함수와 실제 해결 품질 사이의 일관성을 보장하지 못한다는 점을 지적한다. 따라서 “Monotonic‑Regression Negative Contrastive Learning”(MRNCL) 손실을 설계했다. MRNCL은 동일 입력에 대해 서로 다른 노이즈 레벨에서 생성된 두 개의 부정 샘플과 정답 샘플을 이용해, 샘플 간 L2 거리와 에너지 값 사이에 단조 증가 관계(선형 회귀)를 강제한다. 구체적으로, 세 점(정답, 가까운 부정, 먼 부정)의 에너지‑거리 쌍을 선형 함수에 피팅하고, 기울기가 사전 정의된 임계값 γ보다 작으면 패널티를 부여한다. 이 과정은 “성능‑에너지 일관성(performance‑energy consistency)”을 수학적으로 정의하고, 에너지 지형을 보다 매끄럽고 순차적으로 낮아지도록 만든다.

또한 KL 정규화(L_KL)를 추가해 에너지 함수가 확률 분포를 잘 근사하도록 유도한다. KL 정규화는 역전파를 통해 denoising 과정 전체에 걸쳐 에너지의 변동성을 억제함으로써, 샘플이 최적화 과정에서 급격히 튀는 현상을 방지한다. 결과적으로 학습된 에너지 함수는 외부 검증자 없이도 “내재 검증자” 역할을 수행할 수 있다.

추론 측면에서는 hMCTS라는 하이브리드 탐색 전략을 도입한다. 초기 고노이즈 단계에서는 BoN을 사용해 넓은 후보 풀을 빠르게 생성하고, 이때 에너지 값이 낮은 후보들을 다수 보존한다. 이후 노이즈가 감소하면서 탐색 공간이 수축되면, MCTS를 적용해 남은 후보들에 대해 깊이 있는 시뮬레이션과 백프로파게이션을 수행한다. 특히, MCTS 내부에서 변수 스케줄링을 이용해 DDPM의 α_t 시퀀스를 부분적으로 추출함으로써 샘플링 속도를 가속한다. 이 두 단계의 전환은 “탐색 효율성”을 크게 향상시켜, 샘플 수를 늘려도 계산 비용이 선형적으로 증가하지 않도록 만든다.

실험 결과는 설득력 있다. Maze 문제에서는 6×6 크기로 학습한 모델이 15×15 크기의 미지의 미로를 88 % 정확도로 해결했으며, 기존 확산 모델은 전혀 성공하지 못했다. Sudoku에서는 입력 힌트가 적은 OOD 상황에서도 43 %의 성공률을 기록, 기본 모델의 30 %를 크게 앞섰다. Ablation study는 MRNCL이 없을 경우 에너지‑성능 일관성이 크게 감소하고, hMCTS 없이 BoN만 사용할 경우 샘플 수가 늘어나도 성능 향상이 제한적임을 보여준다.

전체적으로 VFScale는 (1) 에너지 함수를 검증자로 활용하기 위한 훈련 목표 설계, (2) 샘플 수 기반 테스트‑타임 스케일링을 효율적으로 수행할 수 있는 하이브리드 탐색 알고리즘이라는 두 축을 통해, 확산 모델이 인간과 유사한 “내재적” 추론 능력을 갖추도록 만든다. 이는 향후 복합 최적화, 논리 퍼즐, 과학적 시뮬레이션 등 다양한 분야에 확산 기반 추론을 적용할 수 있는 중요한 전환점이 될 것으로 기대된다.