DRLM은 단일 디노이저로 초안을 병렬 수정하고 에너지로 평가한다
Breaking Diffusion: Parallel Reasoning with Draft Refinement
무엇인가
이 논문은 자기회귀 모델이 완전한 답이 나오기 전에 앞선 토큰을 확정해 버리고, 이산 확산 모델은 병렬 예측을 하더라도 스케줄에 의존하는 역방향 전이를 통해 토큰 갱신이 매개된다는 문제를 다룬다. 비자기회귀적 초안 수정은 전체 초안을 직접 고칠 수 있지만, 개별적으로 그럴듯한 토큰 수정이 함께 모였을 때 더 나은 답이 된다는 보장은 없다. 저자들은 수정을 제안하는 모델과 그 수정을 평가하는 모델을 분리하지 않고, 하나의 디노이저가 전체 초안에 대한 병렬 수정안을 제안하고 한 번의 순전파로 후보 초안을 평가하는 구조를 제안한다.
어떻게 동작하나
DRLM의 학습은 균일 토큰 대체 채널로 손상된 초안에서 깨끗한 토큰을 복원하는 방식이다. 각 위치는 확률 α로 원래 토큰을 유지하거나 1−α 확률로 어휘에서 균일하게 대체되며, 손상된 초안은 모든 위치에 일반 토큰이 들어 있는 완전한 초안이 된다. 틀린 토큰이 따로 표시되지 않기 때문에 모든 위치에 대한 교차엔트로피 학습은 보존과 수리를 동시에 가르친다. 추론 시에는 학습된 예측의 위치별 argmax를 모아 다음 초안을 직접 만들며, 이 과정에는 확산 역방향 전이가 없다.
무엇과 다른가
평가 단계에서는 전체 후보 조건부 확률에 포함된 자기 토큰 증거를 분리한다. 관측된 토큰은 그 자체로 정답일 가능성을 높이므로, 저자들은 알려진 손상 채널의 우도비를 나누고 재정규화해 조건과 나머지 초안이 각 토큰을 얼마나 지지하는지를 추정하는 캐비티 조건부를 만든다. 이 추정값의 음의 로그합이 캐비티 에너지이며, 낮을수록 문맥적 지지가 크다고 해석된다. 이 점수는 추가 채점 모델 없이 후보 전체를 한 번 읽는 순전파로 계산되고, 제안 단계는 전체 후보 확률을 그대로 사용한다.
어떻게 쓰나
생성 루프는 제안, 읽기, 유지의 세 단계로 구성된다. 매 라운드마다 현재 초안을 부분적으로 다시 손상시키고 병렬 예측 맵을 여러 번 적용해 완전한 후보를 만든 뒤, 고정된 읽기 수준에서 후보를 다시 읽어 캐비티 에너지를 계산한다. 후보의 에너지가 현재 초안보다 엄격히 낮을 때만 교체하고, 그렇지 않으면 기존 초안과 캐시된 에너지를 유지한다. 라운드 수 R과 내부 제안 패스 수 J에 대해 RJ번의 모델 평가가 필요하며, 각 라운드가 끝날 때마다 지금까지의 최저 에너지 초안을 바로 반환할 수 있다.
전제와 한계
실험은 Countdown, Sudoku-Extreme, Maze-Hard에서 진행됐다. 6M 파라미터 DRLM은 Countdown-3에서 99.5%, Countdown-4에서 87.6%, Countdown-5에서 44.4%를 기록했다. Countdown-4는 UDLM보다 14.3포인트, LLaMA-13B SFT보다 36.5포인트, LDDM보다 31.3포인트 높다. Countdown-5에서는 LLaMA-13B의 7.4%, LDDM의 10.3%와 비교해 각각 37.0포인트와 34.1포인트 높은 수치다. Sudoku-Extreme에서는 93.9%로 UDLM의 65.5%보다 28.4포인트, LDDM의 24.1%보다 69.8포인트 높았고, Maze-Hard에서는 73.0%로 LDDM의 39.3%보다 33.7포인트 높았다. 저자들은 UDLM이 Maze-Hard에서 약한 이유를 NELBO 가중치가 거의 깨끗한 캔버스에 학습 신호를 집중해 1,800개 토큰을 랜덤 노이즈에서 한 번에 샘플링하는 감독이 부족하기 때문이라고 추측한다.
분석에서는 에너지 기반 유지가 핵심임을 보인다. 모든 후보를 그대로 받아들이면 Countdown-4가 87.6%에서 82.9%, Countdown-5가 44.4%에서 24.0%, Sudoku-Extreme이 93.9%에서 68.6%로 떨어진다. 매 라운드 균일 캔버스에서 다시 시작하고 최저 에너지 초안만 고르는 방식도 Countdown-4 85.3%, Countdown-5 30.6%, Sudoku-Extreme 77.0%, Maze-Hard 66.2%로 일부 이득만 회복한다. 또한 학습 목적과 디코더를 교차 비교했을 때 디코더만 바꾸면 6.9%에서 36.6%, 목적만 바꾸면 22.5%, 둘을 함께 쓰면 44.4%가 되어 두 요소가 각각 기여한다고 주장한다.
개발자에게 이 논문은 전체 답안의 일관성이 중요한 제약 만족, 퍼즐, 계획형 추론에서 고정된 추론 예산 안에 후보를 병렬로 고치고 평가하는 패턴을 제공한다. 구현 시에는 손상 수준을 조건으로 받는 디노이저, 직접 argmax로 초안을 갱신하는 제안 단계, 자기 토큰 우도를 보정한 에너지 평가, 재손상 수준과 읽기 수준 및 라운드 예산 튜닝을 확인해야 한다. 다만 저자들은 캐비티 에너지를 정답 보증서가 아니라 학습된 노이즈 완화 의사우도 점수로 해석하며, 베이즈 최적 예측기와 α_s가 1에 가까워지는 조건에서만 캐비티 조건부가 정확히 복원된다고 전제한다. 유한한 노이즈에서는 나머지 초안 토큰을 노이즈 관측으로 취급하고, 실험도 세 가지 구조적 과제와 6M 모델에 한정되어 일반 언어 추론으로의 확장은 별도 검증이 필요하다.