스스로 고치는 마코프 모델로 논리 퍼즐 풀기

초록

이 논문은 기존 이산 확산 모델이 마스킹 경로에서 실수를 수정하지 못하고 고정된 단계 수만 사용한다는 한계를 극복하기 위해, 자체 출력으로 학습되는 마코프 전이 커널과 학습 가능한 조기 종료 기준을 도입한다. 토큰을 재마스킹하고 자신감 점수에 따라 언제 예측을 확정할지 결정함으로써 오류를 스스로 교정하고 문제 난이도에 따라 연산량을 조절한다. Sudoku‑Extreme와 Countdown‑4 실험에서 기존 흐름 기반 방법들을 크게 앞서며, 특히 Sudoku에서는 95 % 이상의 유효성을, Countdown‑4에서는 평균 10단계로 96 % 이상의 정확도를 달성한다.

상세 분석

본 연구는 이산 디퓨전 모델이 “마스크 → 토큰” 전이 과정에서 학습 단계는 완전한 마스크 혹은 정답 토큰만을 보게 되므로, 추론 시 모델이 만든 오류 토큰을 되돌릴 메커니즘이 없다는 근본적인 문제를 지적한다. 이를 해결하기 위해 저자들은 두 가지 핵심 설계를 제안한다. 첫째, 마코프 전이 커널 (P_{\theta}) 를 “예측‑확정” 구조로 재구성한다. 입력 상태 (y) 에 대해 토큰별 자신감 점수 (c_{\theta,i}(y)) 를 출력하고, 이 점수가 일정 임계값을 넘으면 해당 위치를 현재 예측 토큰 (p_{\theta,i}^{1}(\cdot|y)) 로 확정하고, 그렇지 않으면 마스크 상태를 유지하거나 재마스크한다. 이렇게 하면 모델이 스스로 “불확실한” 토큰을 남겨두고 이후 단계에서 다시 검토할 수 있다. 둘째, 별도의 진행도 예측 헤드 (\tau_{\theta}(y)) 를 도입해 현재 상태가 목표 해에 얼마나 근접했는지를 추정한다. (\tau_{\theta}) 가 1에 가까워지면 추론을 조기에 종료하도록 학습함으로써, 쉬운 인스턴스는 몇 단계만에 해결하고 어려운 인스턴스는 더 많은 연산을 할당한다.

학습 과정에서는 전통적인 마스킹 경로 (\kappa_t) 를 그대로 사용하되, 모델이 실제 생성한 상태 (\nu_{\theta}^{t}) 를 샘플링해 “온‑폴리시” 데이터를 만든다. 이때 손실 함수는 (i) 자신감이 높은 토큰에 대해 정답 확률을 최대화하고, (ii) 자신감이 낮은 토큰에 대해 마스크 혹은 재마스크 확률을 높이며, (iii) 진행도 예측 (\tau_{\theta}) 와 실제 진행도 (\tau) 사이의 L2 차이를 최소화하는 세 부분으로 구성된다. 결과적으로 모델은 오류가 포함된 상태에서도 학습되므로, 추론 시 오류를 감지하고 재마스킹하여 교정할 수 있다.

실험에서는 동일한 트랜스포머 백본(8 레이어, 8 헤드) 위에 제안된 두 개의 경량 헤드만 추가해 기존 사전학습 모델을 재활용한다. Sudoku‑Extreme 데이터셋에서 기존 DFM 기반 방법이 68 %~89 % 수준에 머물렀던 반면, 제안 모델은 95 % 이상의 유효성을 달성하고, Kaggle Unfiltered 데이터에서도 거의 완전한 정확도를 보였다. 또한, 앙상블 샘플 수를 늘릴수록 정확도가 소폭 상승하지만 평균 단계 수는 62→146 단계로 서브선형적으로 증가해 효율성을 입증한다. Countdown‑4 실험에서는 평균 10단계, 최악의 경우 7 × 10⁶ 스텝 이하로 98.9 % 정확도를 기록했으며, 파라미터 수 34 M으로 기존 Diffusion VDM(85 M)보다 훨씬 가벼우면서도 성능을 앞섰다.

이러한 결과는 “자기 교정(self‑correcting)”과 “난이도 인식(adaptive computation)”이라는 두 축을 동시에 만족시키는 모델 설계가 이산 논리 추론 문제에서 매우 효과적임을 보여준다. 특히, 마코프 전이 커널을 직접 학습하고 진행도 기반 조기 종료를 도입함으로써, 전통적인 고정 스텝 디퓨전 모델이 갖는 비효율성과 오류 전파 문제를 근본적으로 해소한다는 점에서 학술적·실용적 의의가 크다.