CORA-Diff: Confidence-Oriented Residual Acceptance for Efficient Diffusion Language Model Inference

기술 요약: CORA-Diff

문제 정의
확산 언어 모델(Diffusion Language Models, DLMs)은 순차적인 자기회귀(autoregressive) 모델과 대조되는 병렬 토큰 생성의 잠재력을 제공한다. 그러나 실제 DLM 디코더는 고정된 노이즈 제거 호라이즌(denoising horizon)에 의존하여, 호라이즌에 도달할 때까지 매 단계에서 조밀한(dense) 트랜스포머 패스를 실행하는 경우가 많다. 이러한 방식은 많은 토큰 예측이 궤적(trajectory) 초기에 안정화됨에도 불구하고, 블록이 완전히 해결될 때까지 디코더가 이들을 대상으로 조밀한 패스를 계속 처리하게 함으로써 상당한 계산적 중복을 초อน한다. 기존의 가속화 방법들은 이를 완화하려고 시도하지만, 학습된 수용 필터(acceptance filters), 수정된 결정 점수(logits), 명시적 의존성 모델 또는 특수한 시스템 메커니즘(예: 캐시 재사용)에 의존하는 경우가 빈번하다. 이러한 접근 방식은 보조 구성 요소를 학습시키거나 백본을 변경해야 하므로 범용성과 배포 용이성을 제한한다.

방법론: CORA-Diff
저자들은 원래의 전이 규칙(transfer rule)에 의해 해결되지 않은 "잔여(residual)" 위치들을 오직 네이티브 궤적 신호만을 기반으로 식별하고 수용하도록 설계된, 학습이 필요 없는 방법인 CORA-Diff(Confidence-Oriented Residual Acceptance for Diffusion Language Models)를 제안한다.

• 핵심 메커니즘: CORA-Diff는 원래의 백본과 원래의 전이 규칙을 보존한다. 이 방법은 원래의 규칙이 해결하지 못한 위치에 대해서만 적용되는 경량 게이팅 메커니즘을 도입한다.
• 이중 신호 게이팅: 현재의 디코딩 궤적에서 유도된 두 가지 조건을 만족할 경우, 잔여 위치의 토큰은 조기에 수용된다:
  1. 신뢰도(Confidence): Top-1 예측 확률($c^{(s)}_i$)이 임계값 $\delta_p$를 초과한다.
  2. 지속성(Persistence): 동일한 Top-1 예측이 필요한 연속 실행 단계 수($m^{(s)}_i \ge m$) 동안 지속된다.
• 잔여 수용(Residual Acceptance): 수용된 토k은 현재의 예측값으로 고정되며 이후 단계의 컨텍스트로서 가시성을 유지하지만, 추가적인 노이즈 제거 대상에서는 제외된다. 모든 위치가 (원래의 규칙 또는 CORA 게이트에 의해) 해결되면 블록은 조기에 종료되며, 남은 조밀한 순방향 패스(forward passes)를 건너뛴다.
• 이론적 근거: 이 방법은 고신뢰도 및 지속적인 예측이 통계적으로 고정된 호라이즌의 조밀한 종단점(dense endpoint)과 일치할 가능성이 높다는 가설에 근거한다. 저자들은 "지속성-드리프트 분리(persistence-drift separation)" 가정하에, 조밀한 종단점과의 불일치 확률이 지속성 윈도우 $m$에 따라 지수적으로 감소함을 보여주는 조건부 드리프트-테일 바운드(conditional drift-tail bound)를 도출하였다.
• 캘리브레이션(Calibration): 임계값 $(\delta_p, m)$은 작업 점수 저하나 출력 불일치를 최소화하면서 단계 비율(step ratio)을 최적화하기 위해 별도의 캘리브레이션 서브셋(예: 1,000개의 GSM8K 프롬프트)에서 선택된다. 일단 선택되면, 이 임계값들은 이후의 모든 평가 및 크로스 백본 전이에 대해 고정된다.

주요 기여

1. 학습이 필요 없는 잔여 수용(Training-Free Residual Acceptance): 학습된 수용 모델, 로짓 수정, 또는 백본 변경 없이 DLM 추론을 가속화하는 방법을 도입하였다. 이는 오직 네이티브 신뢰도 및 지속성 신호에만 의존한다.
2. 이론적 정형화: 조밀한 종단점과의 불일치를 통한 조기 수용을 정형화하였으며, 일치를 위한 마진 기반 충분 조건을 제공하고 지속적인 예측의 신뢰성을 설명하는 조건부 드리프트-테일 바운드를 도출하였다.
3. 효율적인 블록 수준 종료: 이 경량 규칙이 디코딩 블록의 조기 종료를 가능하게 하여, 실행되는 조밀한 단계 수를 크게 줄임을 입증하였다.

실험 결과
방법론은 LLaDA-8B-Instruct 모델에 대해 네 가지 벤치마크(GSM8K, MATH, HumanEval, MBPP)에서 평가되었으며, Dream 모델로의 전이 가능성도 테스트되었다.

• 런타임 성능: CORA-Diff는 8가지 작업-길이 설정(256/256 및 1024/1024 디노이징 호라이즌 조합) 모두에서 가장 낮은 측정된 런타임을 달성했다.
  • 고정 호라이즌 메커니즘 격리 프로토콜 하에서, 4.34배~9.49배(256/256) 및 10.68배~13.14배(1024/1024)의 속도 향상을 달了했다.
  • 가장 가까운 런타임 베이스라인인 Learn2PD와 비교했을 때, CORA-Diff는 9.1%~30.9% 더 빨랐다.
  • EOS-aware 배포(종료 토큰이 커밋된 후에만 디코딩이 중단되는 경우) 환경에서, EOS-aware 조밀 디코딩 대비 2.70배(GSM8K) 및 3.32배(HumanEval)의 점진적인 속도 향상을 제공했다.
• 작업 품질: 작업 점수는 8개 설정 중 5개에서 조밀 디코딩과 일치하거나 이를 상회했다. 관찰된 가장 큰 성능 저하는 1.22 포인트였다.
• 전이 가능성: 재튜닝 없이 CORA-Diff는 Dream 백본으로 전이되어, 작업 점수의 변화를 무시할 만한 수준(최대 0.006)으로 유지하면서 3.18배~3.53배의 속도 향상을 달성했다.
• 어블레이션(Ablation): 실험을 통해 신뢰도와 지속성을 결합하는 것이 단일 신호를 사용하는 것보다 더 나은 트레이드오프를 제공함을 확인했다.

의의 및 주장
본 논문은 네이티브 신뢰도 및 지속성 신호가 결정론적인 조밀 종단점과 일치할 잔여 위치를 식별하는 데 실질적으로 충분하다고 주장한다. CORA-Diff는 이러한 신호를 활용함으로써 단순한 학습이 필요 없는 게이트를 통해 상당한 반복적 노이즈 제거 계산을 제거한다. 이 연구는 보조 모델, 시스템 특정 최적화, 또는 사전 학습된 모델의 수정 없이도 상당한 추론 가속을 달wal 수 있음을 보여주며, 작업 품질을 유지한다. 저자들은 이를 기존 가속기들과 구별되는 보완적인 접근 방식으로 위치시키며, 단순한 토큰 수용이나 단계별 시스템 최적화가 아닌 블록 완료(block completion)에 초점을 맞춘다.