Flow Reasoning Models: 반복적 자기 개선(Iterative Self-Refinement)을 통한 추론 스케일링

이산 흐름 모델(Discrete flow models)은 최근 몇 단계의 텍스트 생성(few-step text generation)에서 유망한 성능을 보여주었습니다. 하지만 Sudoku나 Zebra 퍼즐과 같은 구조적 추론(structured reasoning) 작업에 단순히 적용했을 때, 이들은 잘못된 정답으로 확신을 가지고 수렴하는 경향이 있습니다(Sudoku 퍼즐의 약 36%만 해결). 우리는 흐름 모델(flow models)을 이용한 구조적 추론을 위한 학습 및 테스트 시간 스케일링(test-time-scaling) 프레임워크인 Flow Reasoning Models (FRMs)를 소개합니다. 우리는 흐름 모델의 낮은 해결률에도 불구하고, 이들이 스스로의 검증기(verifiers) 역할을 할 수 있다는 관찰을 수행했습니다. 정답은 디노이징 역학(denoising dynamics)의 안정적인 고정점(stable fixed point)이며, 다시 노이즈를 추가하고(re-noised) 다시 해결했을 때(re-solved) 자기 자신으로 돌아옵니다. 이는 테스트 시간 스케일링 패러다임을 가능하게 합니다: 많은 후보 솔루션을 제안하고 동적으로 안정적인 것들만 유지하는 방식이며, 이 방식만으로도 Sudoku-Shah(~100%) 및 Zebra(95.9%)에서 높은 해결률에 도달합니다. 이는 해당 분포로 학습하지 않고도 Sudoku-Extreme(96.1%)과 같은 더 어려운 분포 외(out-of-distribution) 퍼즐로까지 일반화됩니다. 그러나 이러한 순수 탐색(pure search) 방식은 잘못된 후보 솔루션을 생성하는 데 막대한 계산량을 낭비합니다. 따라서 우리는 기본 모델의 효율성을 개선하기 위한 학습 레시피를 설계했습니다. 첫째, 자기 조건화 채널(self-conditioning channel)을 사용하여 흐름 모델을 학습시키고 추론 시에는 이를 닫음으로써, 모델이 자신의 과거 예측을 스스로 개선할 수 있도록 합니다. 둘째, 직접 선호 최적화(direct preference optimization)를 사용하여 모델이 자신의 실패한 생성물을 피하도록 학습시킵니다. 이러한 변화는 기본 모델의 효율성을 실질적으로 향상시켜, 동일한 정확도를 위해 비교 대상인 가장 강력한 매칭 마스크 확산(matched masked-diffusion) 베이스라인보다 8배 이상 적은 단 7번의 순전파(forward passes)만으로 Sudoku에서 99.2%에 도달하게 합니다. 테스트 시간 스케일링과 결합될 때, 이는 흐름 모델이 어려운 분포 외(out-of-distribution) 퍼즐(예: Sudoku-Extreme)을 훨씬 더 효율적으로 해결할 수 있게 합니다.