기술 요약: Diffusion-Proof

문제 정의

형식적 정리 증명(Formal theorem proving)은 복잡한 추론 단계 전반에 걸쳐 장기적인 일관성과 논리적 일관성을 유지하는 것을 요구합니다. 대규모 언어 모델(LLM)이 이 분야에서 크게 발전했음에도 불구하고, 현재의 최첨단 방식들은 자기 회귀(Auto-Regressive, AR) 모델에 크게 의존하고 있습니다. 이러한 모델은 텍p를 이전 토큰만을 기반으로 순차적으로 생성하며 왼쪽에서 오른쪽으로 텍스트를 예측합니다. 이 아키텍처는 다음과 같은 내재적인 한계를 가집니다:

1. 오류 누적: 초기 토큰의 작은 오류가 긴 시퀀스 동안 기하급수적으로 누적되어 논리적 불일치를 초래할 수 있습니다.
2. 장기적 일관성 부족: AR 모델은 증명의 전체적인 글로벌 계획을 유지하는 데 어려움을 겪으며, 종종 중간 단계의 택틱(tactic)을 최종 목표와 일치시키는 데 실패합니다.
3. 인필링(In-filling) 수행 능력 부재: AR 모델은 전체 증명을 다시 작성하지 않고도 중간 섹션을 수정하는 "인필링" 교정(증명의 중간 부분을 수정하면서 주변 문맥을 보존하는 작업)을 자연스럽게 수행할 수 없습니다. 이는 국소적인 오류를 해결하는 데 매우 중요합니다.

디퓨전 LLM(dLLM)은 반복적인 디노이징(denoising)과 양방향 문맥을 활용하여 코딩 및 일반 텍스트 생성에서 가능성을 보여주었으나, 형식적 수학적 추론에 대한 적용은 아직 거의 탐구되지 않았습니다.

방법론

저자들은 형식적 정리 증명을 위해 특별히 설계된 dLLM을 훈련하고 적용하는 첫 번째 프레임워크인 Diffusion-Proof를 제안합니다. 이 프레임워크는 Lean4 환경 내에서 작동하며, dLLM-Prover-7B와 dLLM-Corrector-7B라는 이중 모델 아키텍처를 활용합니다.

1. 데이터 수집 및 전처리

• 데이터셋: 저자들은 이전 연구들로부터 550만 개의 정리 증명 기록을 수집했습니다.
• 정제: 자연어(NL) 문장, 형식 언어(FL) 문장, 그리고 FL 증명을 분리하기 위해 규칙 기반 방법을 사용하여 데이터를 표준화했습니다.
• SFT 데이터셋: 300,000개의 기록을 선택하여 지도 미세 조정(SFT) 데이터셋을 생성하였으며, 순수 Lean 증명과 NL 주석이 포함된 증명의 비율을 1:2로 유지했습니다.
• 교정 데이터: have 키워드에 의해 표시되는 하위 목표 분해(subgoal decomposition)를 특징으로 하는 데이터의 하위 집합을 처리했습니다. 증명 블록을 플레이스홀더 토큰으로 채워 "인필링" 교정을 위한 훈련 데이터를 생성했습니다.

2. 모델 아키텍처 및 훈련

이 프레임워크는 블록 기반 디퓨전 모델인 Fast-dLLM-V2-7B를 기반으로 구축되었습니다. 토큰을 독립적으로 마스킹하는 기존의 dLLM과 달리, 이 모델은 시퀀스를 크기 $B$인 블록으로 나누며, 블록 간의 인과 관계를 보존하면서 개별 블록 내에서만 양방향 어텐션(attention)을 제한합니다.

• dLLM-Prover-7B:

  • 목표: 장기적으로 일관된 택틱 사용을 동반한 전체 증명 생성을 수행합니다.
  • 훈련: 표준 텍-투-텍(text-to-text) 프레임워크를 사용하여 300k SFT 데이터셋으로 미세 조정되었습니다.
  • 전략: 손실 곡선을 안정화하기 위해 증명 복잡도(증명 길이 증가)에 따라 훈련 데이터를 구성하는 **커리큘럼 학습(curriculum learning)**을 채택했습니다. 디퓨전 블록 크기는 32로 설정되었습니다.

• dLLM-Corrector-7B:

  • 목표: 양방향 정보를 사용하여 국소적 증명 교정을 수행합니다.
  • 훈련: 128k 하위 목표 분해 데이터셋으로 미세 조정되었습니다.
  • 혁신: 대규모 블록 교정(Large Block Correction) 방법을 도입했습니다. 디퓨전 블록 크기를 32에서 512로 확장했습니다. 이 모델은 오류가 있는 하위 목표 증명을 일련의 마스크 토큰으로 대체한 후, 확장된 문맥(접두사와 접미사 모두)을 사용하여 해당 블록을 재생성함으로써 인필링 교정을 수행하도록 훈련되었습니다.

3. 추론 파이프라인

1. 생성: dLLM-Prover-7B가 블록 디퓨전 생성을 사용하여 주어진 정리 문장에 대한 완전한 증명을 생성합니다.
2. 검증: Lean에 의해 증명이 검증됩니다.
3. 교정: 증명이 검증에 실패했지만 상위 수준의 스켈레톤(하위 목표 문장들)이 올바른 경우, dLLM-Corrector-7B가 호출됩니다. 이 모델은 오류가 있는 하위 목표 증명을 마스크 토큰으로 교체하고, 512-토큰 디퓨전 블록을 사용하여 이를 재생성함으로써, 글로벌 일관성을 유지하면서 양방향 문맥을 활용해 국소적 오류를 수정합니다.

주요 기여

1. 정리 증명을 위한 최초의 dLLM 프레임워크: Diffusion-Proof는 수학적 추론에서의 장기적 일관성과 오류 교정이라는 특정 과제를 해결하기 위해 디퓨전 모델을 정리 증명에 적응시킨 최초의 프레임워크입니다.
2. 대규모 블록 교정 훈련: 저자들은 효과적인 인필링 교정을 가능하게 하기 위해 큰 디퓨전 블록(512 토큰)과 양방향 정보를 활용하는 새로운 훈련 방법을 도입했습니다. 이는 디퓨전 모델의 고유한 기능이지만 AR 모델에는 없는 기능입니다.
3. 이중 모델 시너지: 이 프레임워크는 프로버(prover)가 글로벌 계획을 담당하고 교정기(corrector)가 국소적 정제를 담당하는 시너지 효과를 입증하며, 단일 모델 베이스라인을 능가하는 성능을 보여줍니다.

실험 결과

프레임워크는 pass@32 지표를 사용하여 MiniF2F-Test(244개 문제)와 ProofNet-Test(186개 문제)에서 평가되었습니다.

• 성능 향상: 동일한 데이터셋으로 훈련된 베이스라인 AR 모델(Qwen-2.5-Lean-SFT-7B)과 비교했을 때:
  • MiniF2F-Test: Diffusion-Proof는 **50.00%**의 정확도를 달성하여 베이스라인(43.85%) 대비 **6.14%**의 절대적 향상을 보였습니다.
  • ProofNet-Test: Diffusion-Proof는 **7.53%**의 정확도를 달성하여 베이스라인(5.91%) 대비 **1.61%**의 절대적 향상을 보였습니다.
• 문제 유형 분석: 모델은 특히 IMO(국제 수학 올림피아드) 문제와 같은 까다로운 카테고리에서 상당한 이득을 보였으며, 여기서 5.00%에서 15.00%로 향상되었습니다.
• 사례 연구: 주목할 만하게도, Diffusion-Proof는 (Long CoT 추론을 갖춘) 더 발전된 DeepSeek-Prover-V2-7B가 pass@32 하에서 해결하지 못한 IMO 문제인 **imo_1962_p2**를 성공적으로 해결했습니다. 분석 결과, AR 모델은 필요한 장거리 의존성과 일관된 택틱 계획을 수행하는 데 어려움을 겪은 반면, 디퓨전 교정기는 양방향 문맥을 성공적으로 활용하여 증명을 해결했습니다.
• 절제 연구(Ablation Studies):
  • 교정기를 제거하면 MiniF2F에서 성능이 1.64% 감소하여 교정 단계의 가치를 확인했습니다.
  • 교정기 없이도 dLLM-Prover 단독으로 AR 베이스라인보다 4.51% 높은 성능을 보여, 디퓨전 아키텍처의 내재적 이점을 입증했습니다.

의의 및 주장

본 논문은 Diffusion-Proof가 형식적 정리 증명에서 디퓨전 모델의 독특한 장점, 특히 장기적 일관성 유지와 양방향 인필링 교정 능력을 강조한다고 주장합니다.

• 아키텍처적 이점: 결과는 디퓨전 생성 패러다임이 긴 시퀀스에 걸쳐 엄격한 논리적 일관성을 요구하는 작업에 대해 AR 모델보다 구조적 우위를 제공함을 시사합니다. 이는 더 낮은 검증 손실 분포와 더 높은 성공률로 입증됩니다.
• 미래 잠재력: 저자들은 이 작업이 dLLM이 강화 학습(RL)이나 Long CoT 없이도 AR 베이스라인을 능가할 수 있음을 보여주는 기초적인 단계라고 위치시킵니다. 현재의 한계로는 계산 규모와 RL 통합의 부재가 있으나, 이 프레임워크는 향상된 추론 능력을 가진 모델에 대한 미래 연구의 유망한 방향을 제시합니다.

저자들은 이 연구의 주된 목적이 거대한 규모 확장이나 RL을 통해 최첨단 성능(SOTA)을 달 achieve하는 것이 아니라, dLLM의 아키텍처적 잠재력을 조사하는 것임을 강조하며, 관찰된 개선 사항이 형식적 추론에서 디퓨전 패러다임의 효용성을 보여주는 강력한 지표임을 밝히고 있습니다.